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