io_mesh_stl/stl_utils.py

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