io_mesh_stl/stl_utils.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8 compliant>
  20
  21 """
  22 Import and export STL files
  23
  24 Used as a blender script, it load all the stl files in the scene:
  25
  26 blender --python stl_utils.py -- file1.stl file2.stl file3.stl ...
  27 """
  28
  29 # TODO: endien
  30
  31
  32 class ListDict(dict):
  33     """
  34     Set struct with order.
  35
  36     You can:
  37        - insert data into without doubles
  38        - get the list of data in insertion order with self.list
  39
  40     Like collections.OrderedDict, but quicker, can be replaced if
  41     ODict is optimised.
  42     """
  43
  44     def __init__(self):
  45         dict.__init__(self)
  46         self.list = []
  47         self._len = 0
  48
  49     def add(self, item):
  50         """
  51         Add a value to the Set, return its position in it.
  52         """
  53         value = self.setdefault(item, self._len)
  54         if value == self._len:
  55             self.list.append(item)
  56             self._len += 1
  57
  58         return value
  59
  60
  61 # an stl binary file is
  62 # - 80 bytes of description
  63 # - 4 bytes of size (unsigned int)
  64 # - size triangles :
  65 #
  66 #   - 12 bytes of normal
  67 #   - 9 * 4 bytes of coordinate (3*3 floats)
  68 #   - 2 bytes of garbage (usually 0)
  69 BINARY_HEADER = 80
  70 BINARY_STRIDE = 12 * 4 + 2
  71
  72
  73 def _header_version():
  74     import bpy
  75     return "Exported from Blender-" + bpy.app.version_string
  76
  77
  78 def _is_ascii_file(data):
  79     """
  80     This function returns True if the data represents an ASCII file.
  81
  82     Please note that a False value does not necessary means that the data
  83     represents a binary file. It can be a (very *RARE* in real life, but
  84     can easily be forged) ascii file.
  85     """
  86
  87     import os
  88     import struct
  89
  90     # Skip header...
  91     data.seek(BINARY_HEADER)
  92     size = struct.unpack('<I', data.read(4))[0]
  93     # Use seek() method to get size of the file.
  94     data.seek(0, os.SEEK_END)
  95     file_size = data.tell()
  96     # Reset to the start of the file.
  97     data.seek(0)
  98
  99     if size == 0:  # Odds to get that result from an ASCII file are null...
 100         print("WARNING! Reported size (facet number) is 0, assuming invalid binary STL file.")
 101         return False  # Assume binary in this case.
 102
 103     return (file_size != BINARY_HEADER + 4 + BINARY_STRIDE * size)
 104
 105
 106 def _binary_read(data):
 107     # Skip header...
 108
 109     import os
 110     import struct
 111
 112     data.seek(BINARY_HEADER)
 113     size = struct.unpack('<I', data.read(4))[0]
 114
 115     if size == 0:
 116         # Workaround invalid crap.
 117         data.seek(0, os.SEEK_END)
 118         file_size = data.tell()
 119         # Reset to after-the-size in the file.
 120         data.seek(BINARY_HEADER + 4)
 121
 122         file_size -= BINARY_HEADER + 4
 123         size = file_size // BINARY_STRIDE
 124         print("WARNING! Reported size (facet number) is 0, inferring %d facets from file size." % size)
 125
 126     # We read 4096 elements at once, avoids too much calls to read()!
 127     CHUNK_LEN = 4096
 128     chunks = [CHUNK_LEN] * (size // CHUNK_LEN)
 129     chunks.append(size % CHUNK_LEN)
 130
 131     unpack = struct.Struct('<12f').unpack_from
 132     for chunk_len in chunks:
 133         if chunk_len == 0:
 134             continue
 135         buf = data.read(BINARY_STRIDE * chunk_len)
 136         for i in range(chunk_len):
 137             # read the normal and points coordinates of each triangle
 138             pt = unpack(buf, BINARY_STRIDE * i)
 139             yield pt[:3], (pt[3:6], pt[6:9], pt[9:])
 140
 141
 142 def _ascii_read(data):
 143     # an stl ascii file is like
 144     # HEADER: solid some name
 145     # for each face:
 146     #
 147     #     facet normal x y z
 148     #     outerloop
 149     #          vertex x y z
 150     #          vertex x y z
 151     #          vertex x y z
 152     #     endloop
 153     #     endfacet
 154
 155     # strip header
 156     data.readline()
 157
 158     curr_nor = None
 159
 160     for l in data:
 161         l = l.lstrip()
 162         if l.startswith(b'facet'):
 163             curr_nor = tuple(map(float, l.split()[2:]))
 164         # if we encounter a vertex, read next 2
 165         if l.startswith(b'vertex'):
 166             yield curr_nor, [tuple(map(float, l_item.split()[1:])) for l_item in (l, data.readline(), data.readline())]
 167
 168
 169 def _binary_write(filepath, faces):
 170     import struct
 171     import itertools
 172     from mathutils.geometry import normal
 173
 174     with open(filepath, 'wb') as data:
 175         fw = data.write
 176         # header
 177         # we write padding at header beginning to avoid to
 178         # call len(list(faces)) which may be expensive
 179         fw(struct.calcsize('<80sI') * b'\0')
 180
 181         # 3 vertex == 9f
 182         pack = struct.Struct('<9f').pack
 183
 184         # number of vertices written
 185         nb = 0
 186
 187         for face in faces:
 188             # calculate face normal
 189             # write normal + vertexes + pad as attributes
 190             fw(struct.pack('<3f', *normal(*face)) + pack(*itertools.chain.from_iterable(face)))
 191             # attribute byte count (unused)
 192             fw(b'\0\0')
 193             nb += 1
 194
 195         # header, with correct value now
 196         data.seek(0)
 197         fw(struct.pack('<80sI', _header_version().encode('ascii'), nb))
 198
 199
 200 def _ascii_write(filepath, faces):
 201     from mathutils.geometry import normal
 202
 203     with open(filepath, 'w') as data:
 204         fw = data.write
 205         header = _header_version()
 206         fw('solid %s\n' % header)
 207
 208         for face in faces:
 209             # calculate face normal
 210             fw('facet normal %f %f %f\nouter loop\n' % normal(*face)[:])
 211             for vert in face:
 212                 fw('vertex %f %f %f\n' % vert[:])
 213             fw('endloop\nendfacet\n')
 214
 215         fw('endsolid %s\n' % header)
 216
 217
 218 def write_stl(filepath="", faces=(), ascii=False):
 219     """
 220     Write a stl file from faces,
 221
 222     filepath
 223        output filepath
 224
 225     faces
 226        iterable of tuple of 3 vertex, vertex is tuple of 3 coordinates as float
 227
 228     ascii
 229        save the file in ascii format (very huge)
 230     """
 231     (_ascii_write if ascii else _binary_write)(filepath, faces)
 232
 233
 234 def read_stl(filepath):
 235     """
 236     Return the triangles and points of an stl binary file.
 237
 238     Please note that this process can take lot of time if the file is
 239     huge (~1m30 for a 1 Go stl file on an quad core i7).
 240
 241     - returns a tuple(triangles, triangles' normals, points).
 242
 243       triangles
 244           A list of triangles, each triangle as a tuple of 3 index of
 245           point in *points*.
 246
 247       triangles' normals
 248           A list of vectors3 (tuples, xyz).
 249
 250       points
 251           An indexed list of points, each point is a tuple of 3 float
 252           (xyz).
 253
 254     Example of use:
 255
 256        >>> tris, tri_nors, pts = read_stl(filepath)
 257        >>> pts = list(pts)
 258        >>>
 259        >>> # print the coordinate of the triangle n
 260        >>> print(pts[i] for i in tris[n])
 261     """
 262     import time
 263     start_time = time.process_time()
 264
 265     tris, tri_nors, pts = [], [], ListDict()
 266
 267     with open(filepath, 'rb') as data:
 268         # check for ascii or binary
 269         gen = _ascii_read if _is_ascii_file(data) else _binary_read
 270
 271         for nor, pt in gen(data):
 272             # Add the triangle and the point.
 273             # If the point is already in the list of points, the
 274             # index returned by pts.add() will be the one from the
 275             # first equal point inserted.
 276             tris.append([pts.add(p) for p in pt])
 277             tri_nors.append(nor)
 278
 279     print('Import finished in %.4f sec.' % (time.process_time() - start_time))
 280
 281     return tris, tri_nors, pts.list
 282
 283
 284 if __name__ == '__main__':
 285     import sys
 286     import bpy
 287     from io_mesh_stl import blender_utils
 288
 289     filepaths = sys.argv[sys.argv.index('--') + 1:]
 290
 291     for filepath in filepaths:
 292         objName = bpy.path.display_name(filepath)
 293         tris, pts = read_stl(filepath)
 294
 295         blender_utils.create_and_link_mesh(objName, tris, pts)