Added binary stl save option

geomdl/exchange.py

@@ -8,6 +8,7 @@
 """
 
 import warnings
+import struct
 
 from . import Abstract
 from . import Multi
@@ -163,7 +164,7 @@ def save_obj_multi(surface_list=(), file_name=None, vertex_spacing=2):
         print("Cannot open " + str(file_name) + " for writing")
 
 
-def save_stl(surf_in=None, file_name=None, vertex_spacing=2):
+def save_stl(surf_in=None, file_name=None, binary=True, vertex_spacing=2):
     """ Saves surfaces as a .stl file.
 
     This function serves as a router between :py:func:`save_stl_single` and :py:func:`save_stl_multi`.
@@ -172,17 +173,25 @@ def save_stl(surf_in=None, file_name=None, vertex_spacing=2):
     :type surf_in: Abstract.Surface or Multi.MultiSurface
     :param file_name: name of the output file
     :type file_name: str
+    :param binary: True if the saved STL file is going to be in binary format
+    :type binary: bool
     :param vertex_spacing: size of the triangle edge in terms of points sampled on the surface
     :type vertex_spacing: int
     """
     if isinstance(surf_in, Multi.MultiSurface):
-        save_stl_multi(surf_in, file_name, vertex_spacing)
+        if binary:
+            save_stl_binary_multi(surf_in, file_name, vertex_spacing)
+        else:
+            save_stl_ascii_multi(surf_in, file_name, vertex_spacing)
     else:
-        save_stl_single(surf_in, file_name, vertex_spacing)
+        if binary:
+            save_stl_binary_single(surf_in, file_name, vertex_spacing)
+        else:
+            save_stl_ascii_single(surf_in, file_name, vertex_spacing)
 
 
-def save_stl_single(surface=None, file_name=None, vertex_spacing=2):
-    """ Saves a single surface as a .stl file.
+def save_stl_ascii_single(surface=None, file_name=None, vertex_spacing=2):
+    """ Saves a single surface as an ASCII .stl file.
 
     :param surface: surface to be saved
     :type surface: Abstract.Surface
@@ -202,27 +211,28 @@ def save_stl_single(surface=None, file_name=None, vertex_spacing=2):
     # Create the file and start saving triangulated surface points
     try:
         with open(file_name, 'w') as fp:
-            fp.write("solid Surface\n")
             vertices, triangles = exh.make_obj_triangles(surface.surfpts,
                                                          int((1.0 / surface.delta) + 1),
                                                          int((1.0 / surface.delta) + 1),
                                                          vertex_spacing)
 
+            fp.write("solid Surface\n")
             for t in triangles:
-                fp.write("\tfacet normal " + str(t.normal[0]) + " " + str(t.normal[1]) + " " + str(t.normal[2]) + "\n")
+                line = "\tfacet normal " + str(t.normal[0]) + " " + str(t.normal[1]) + " " + str(t.normal[2]) + "\n"
+                fp.write(line)
                 fp.write("\t\touter loop\n")
                 for v in t.vertices:
-                    fp.write("\t\t\tvertex " + str(v.x) + " " + str(v.y) + " " + str(v.z) + "\n")
+                    line = "\t\t\tvertex " + str(v.x) + " " + str(v.y) + " " + str(v.z) + "\n"
+                    fp.write(line)
                 fp.write("\t\tendloop\n")
                 fp.write("\tendfacet\n")
-
             fp.write("endsolid Surface\n")
     except IOError:
         print("Cannot open " + str(file_name) + " for writing")
 
 
-def save_stl_multi(surface_list=(), file_name=None, vertex_spacing=2):
-    """ Saves multiple surfaces as a .stl file.
+def save_stl_ascii_multi(surface_list=(), file_name=None, vertex_spacing=2):
+    """ Saves multiple surfaces as an ASCII .stl file.
 
     :param surface_list: list of surfaces to be saved
     :type surface_list: Multi.MultiAbstract
@@ -258,14 +268,104 @@ def save_stl_multi(surface_list=(), file_name=None, vertex_spacing=2):
                                                              int((1.0 / surface.delta) + 1),
                                                              vertex_spacing)
                 for t in triangles:
-                    fp.write(
-                        "\tfacet normal " + str(t.normal[0]) + " " + str(t.normal[1]) + " " + str(t.normal[2]) + "\n")
+                    line = "\tfacet normal " + str(t.normal[0]) + " " + str(t.normal[1]) + " " + str(t.normal[2]) + "\n"
+                    fp.write(line)
                     fp.write("\t\touter loop\n")
                     for v in t.vertices:
-                        fp.write("\t\t\tvertex " + str(v.x) + " " + str(v.y) + " " + str(v.z) + "\n")
+                        line = "\t\t\tvertex " + str(v.x) + " " + str(v.y) + " " + str(v.z) + "\n"
+                        fp.write(line)
                     fp.write("\t\tendloop\n")
                     fp.write("\tendfacet\n")
 
             fp.write("endsolid Surface\n")
     except IOError:
         print("Cannot open " + str(file_name) + " for writing")
+
+
+def save_stl_binary_single(surface, file_name, vertex_spacing=2):
+    """ Saves a single surface as a binary .stl file.
+
+    Inspired from https://github.com/apparentlymart/python-stl
+
+    :param surface: surface to be saved
+    :type surface: Abstract.Surface
+    :param file_name: name of the output file
+    :type file_name: str
+    :param vertex_spacing: size of the triangle edge in terms of points sampled on the surface
+    :type vertex_spacing: int
+    """
+    # Input validity checking
+    if not isinstance(surface, Abstract.Surface):
+        raise ValueError("Input is not a surface")
+    if not file_name:
+        raise ValueError("File name field is required")
+    if vertex_spacing < 1 or not isinstance(vertex_spacing, int):
+        raise ValueError("Vertex spacing must be an integer value and it must be bigger than zero")
+
+    # Create the file and start saving triangulated surface points
+    try:
+        with open(file_name, 'wb') as fp:
+            vertices, triangles = exh.make_obj_triangles(surface.surfpts,
+                                                         int((1.0 / surface.delta) + 1),
+                                                         int((1.0 / surface.delta) + 1),
+                                                         vertex_spacing)
+
+            # Write triangle list to the binary STL file
+            fp.write(b'\0' * 80)  # header
+            fp.write(struct.pack('<i', len(triangles)))  # number of triangles
+            for t in triangles:
+                fp.write(struct.pack('<3f', *t.normal))  # normal
+                for v in t.vertices:
+                    fp.write(struct.pack('<3f', *v.data))  # vertices
+                fp.write(b'\0\0')  # attribute byte count
+    except IOError:
+        print("Cannot open " + str(file_name) + " for writing")
+
+
+def save_stl_binary_multi(surface_list=(), file_name=None, vertex_spacing=2):
+    """ Saves multiple surfaces as a binary .stl file.
+
+    :param surface_list: list of surfaces to be saved
+    :type surface_list: Multi.MultiAbstract
+    :param file_name: name of the output file
+    :type file_name: str
+    :param vertex_spacing: size of the triangle edge in terms of points sampled on the surface
+    :type vertex_spacing: int
+    """
+    # Input validity checking
+    if not isinstance(surface_list, Multi.MultiAbstract):
+        raise ValueError("Input must be a list of surfaces")
+    if not file_name:
+        raise ValueError("File name field is required")
+    if vertex_spacing < 1 or not isinstance(vertex_spacing, int):
+        raise ValueError("Vertex spacing must be an integer value and it must be bigger than zero")
+
+    # Create the file and start saving triangulated surface points
+    try:
+        with open(file_name, 'wb') as fp:
+            # Loop through MultiSurface object
+            triangles_list = []
+            for surface in surface_list:
+                if not isinstance(surface, Abstract.Surface):
+                    warnings.warn("Encountered a non-surface object")
+                    continue
+
+                # Set surface delta
+                surface.delta = surface_list.delta
+
+                vertices, triangles = exh.make_obj_triangles(surface.surfpts,
+                                                             int((1.0 / surface.delta) + 1),
+                                                             int((1.0 / surface.delta) + 1),
+                                                             vertex_spacing)
+                triangles_list += triangles
+
+            # Write triangle list to the binary STL file
+            fp.write(b'\0' * 80)  # header
+            fp.write(struct.pack('<i', len(triangles_list)))  # number of triangles
+            for t in triangles_list:
+                fp.write(struct.pack('<3f', *t.normal))  # normal
+                for v in t.vertices:
+                    fp.write(struct.pack('<3f', *v.data))  # vertices
+                fp.write(b'\0\0')  # attribute byte count
+    except IOError:
+        print("Cannot open " + str(file_name) + " for writing")