Surf2StlWriter.py

import numpy as np
import datetime
import math
from scipy.spatial import Delaunay 
import struct

"""
Example:

import numpy as np
import matplotlib.tri as mtri
from scipy.spatial import Delaunay
from Surf2StlWriter import *


# Create 1st x/y/z datas
x = np.linspace(-6, 6, 30)
y = np.linspace(-6, 6, 30)
X1, Y1 = np.meshgrid(x, y)
Z1 = np.sin(np.sqrt(X1 ** 2 + Y1 ** 2))

# Create 2nd x/y/z datas and Delaunay objects
u = np.linspace(0, 2.0 * np.pi, endpoint=True, num=50)
v = np.linspace(-0.5, 0.5, endpoint=True, num=10)
u, v = np.meshgrid(u, v)
u, v = u.flatten(), v.flatten()

X2 = (1 + 0.5 * v * np.cos(u / 2.0)) * np.cos(u)
Y2 = (1 + 0.5 * v * np.cos(u / 2.0)) * np.sin(u)
Z2 = 0.5 * v * np.sin(u / 2.0)

tri = mtri.Triangulation(u, v)
delaunay_tri = Delaunay(np.array([u, v]).T)

# Write ascii STL file
s2sWriter_ascii = Surf2StlWriter()
with s2sWriter_ascii("writer_test_ascii.stl", mode='ascii'):
    s2sWriter_ascii.write_surf(X1, Y1, Z1)
    s2sWriter_ascii.write_surf_tri(X2, Y2, Z2, delaunay_tri)

# Write binary STL file
s2sWriter_binary = Surf2StlWriter()
with s2sWriter_binary("writer_test_binary.stl", mode='binary'):
    s2sWriter_binary.write_surf(X1, Y1, Z1)
    s2sWriter_binary.write_surf_tri(X2, Y2, Z2, delaunay_tri)

"""

class Surf2StlWriter:
    
    file_name = ""
    file = None
    mode = 'binary'
    title_str = ""
    solid_id = 0
    total_facets = 0
    
    def __init__(self):
        pass

    def __enter__(self):
        self.close()
        self.open(self.file_name, self.mode)
        return self

    def __call__(self, file_name=None, mode='binary'):
        if file_name != None:
            self.file_name = file_name
        self.mode = mode
        return self
        
    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        
    def get_file_name(self):
        return self.file_name
    
    def open(self, file_name, mode='binary'):
        self.file_name = file_name
        self.file = open(self.file_name, 'w' if self.mode == 'ascii' else 'wb')
        self.title_str = 'Created by Surf2StlWriter.py %s' % datetime.datetime.now().strftime('%d-%b-%Y %H:%M:%S')
    
        if self.mode != 'ascii':
            title_str_ljust = self.title_str.ljust(80)
            # f.write(title_str_ljust.encode('utf-8')) # same as 'ascii' for alphabet characters
            self.file.write(title_str_ljust.encode('ascii'))
            self.file.write(struct.pack('i', 0))

    def close(self):
        if self.file != None:
            if self.mode != 'ascii':
                self.file.seek(80, 0)
                self.file.write(struct.pack('i', self.total_facets))

            self.file.close()
            self.file = None

    def write_surf(self, x, y, z):
        if self.file == None:
            raise IOError("File not opened.")
        
        if len(x.shape) == 1 and x.shape[0] == z.shape[1] \
                and len(y.shape) == 1 and y.shape[0] == z.shape[0]:
            x, y = np.meshgrid(x, y)

        if len(x.shape) != len(z.shape) \
                or len(y.shape) != len(z.shape) \
                or x.shape[1] != z.shape[1] \
                or y.shape[0] != z.shape[0]:
            raise Exception('Unable to resolve x and y variables')

        if self.mode == 'ascii':
            self.file.write('solid %s\n' % self.title_str)

        for i in range(z.shape[0]-1):
            for j in range(z.shape[1]-1):
                p1 = np.array([x[i,j], y[i,j], z[i,j]])
                p2 = np.array([x[i,j+1], y[i,j+1], z[i,j+1]])
                p3 = np.array([x[i+1,j+1], y[i+1,j+1], z[i+1,j+1]])
                val = self.local_write_facet(self.file, p1, p2, p3, self.mode, self.solid_id)
                self.total_facets += val

                p1 = np.array([x[i+1,j+1], y[i+1,j+1], z[i+1,j+1]])
                p2 = np.array([x[i+1,j], y[i+1,j], z[i+1,j]])
                p3 = np.array([x[i,j], y[i,j], z[i,j]])
                val = self.local_write_facet(self.file, p1, p2, p3, self.mode, self.solid_id)
                self.total_facets += val

        if self.mode == 'ascii':
            self.file.write('endsolid %s\n' % self.title_str)
        self.solid_id += 1

    def write_surf_tri(self, x, y, z, tri):
        if self.file == None:
            raise IOError("File not opened.")

        if len(x.shape) != 1 \
                or len(y.shape) != 1 \
                or len(z.shape) != 1:
            raise Exception('Each variable x,y,z must be a 1-dimensional array')

        if x.shape[0] != z.shape[0] \
                or y.shape[0] != z.shape[0]:
            raise Exception('Number of x,y,z elements must be equal')

        if self.mode == 'ascii':
            self.file.write('solid %s\n' % self.title_str)

        indices = tri.simplices
        verts = tri.points[indices]
        for i in range(0, indices.shape[0], 1):
            p = indices[i]
            p1 = np.array([x[p[0]], y[p[0]], z[p[0]]])
            p2 = np.array([x[p[1]], y[p[1]], z[p[1]]])
            p3 = np.array([x[p[2]], y[p[2]], z[p[2]]])
            val = self.local_write_facet(self.file, p1, p2, p3, self.mode, self.solid_id)
            self.total_facets += val

        if self.mode == 'ascii':
            self.file.write('endsolid %s\n' % self.title_str)
        self.solid_id += 1
        
    def local_write_facet(self, f, p1, p2, p3, mode, solid_id):
        if np.isnan(p1).any() or np.isnan(p2).any() or np.isnan(p3).any():
                return 0;

        n = self.local_find_normal(p1, p2, p3)
        if self.mode == 'ascii':
            self.file.write('facet normal %.7f %.7f %.7f\n' % (n[0], n[1], n[2]))
            self.file.write('outer loop\n')
            self.file.write('vertex %.7f %.7f %.7f\n' % (p1[0], p1[1], p1[2]))
            self.file.write('vertex %.7f %.7f %.7f\n' % (p2[0], p2[1], p2[2]))
            self.file.write('vertex %.7f %.7f %.7f\n' % (p3[0], p3[1], p3[2]))
            self.file.write('endloop\n')
            self.file.write('endfacet\n')
        else:
            self.file.write(struct.pack('%sf' % len(n), *n))
            self.file.write(struct.pack('%sf' % len(p1), *p1))
            self.file.write(struct.pack('%sf' % len(p2), *p2))
            self.file.write(struct.pack('%sf' % len(p3), *p3))
            self.file.write(struct.pack('h', solid_id))
        return 1

    def local_find_normal(self, p1, p2, p3):
        v1 = p2 - p1
        v2 = p3 - p1
        v3 = np.cross(v1, v2)
        n = v3 / math.sqrt(np.sum(v3*v3))
        return n