io_mesh_stl/stl_utils.py

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