cgkit stl exporter

stlexport.py

# ***** BEGIN LICENSE BLOCK *****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
#
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
#
# The Original Code is the Python Computer Graphics Kit.
#
# The Initial Developer of the Original Code is Matthias Baas.
# Portions created by the Initial Developer are Copyright (C) 2004
# the Initial Developer. All Rights Reserved.
#
# Contributor(s):
#
#   Jeremy Carson <[email protected]>
#
# Alternatively, the contents of this file may be used under the terms of
# either the GNU General Public License Version 2 or later (the "GPL"), or
# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
#
# ***** END LICENSE BLOCK *****

import os.path, sys, re, math
from cgtypes import *
from globalscene import getScene
from geomobject import *
from trimeshgeom import TriMeshGeom
from polyhedrongeom import PolyhedronGeom
import pluginmanager
import cmds

# STLExporter (not part of original cgkit, adapted from OBJExporter)
# MUST be added to /all/__init__.py near bottom: import cgkit.stlexport
class STLExporter:

    _protocols = ["Export"]

    # extension
    def extension():
        """Return the file extensions for this format."""
        return ["stl"]
    extension = staticmethod(extension)

    # description
    def description(self):
        """Return a short description for the file dialog."""
        return "StereoLithography file"
    description = staticmethod(description)

    # exportFile
    def exportFile(self, filename, root=None,**kwargs):
        """Export an STL file.

        root is the root of the subtree that should be exported.
        additional arguments are quietly eaten.

        Output format taken from MeshLab.

        TODO:
            - Support more than vertices and faces.
        """

        self.fhandle = file(filename, "w")

        self.root = cmds.worldObject(root)
        self.v_offset = 0
        self.vt_offset = 0
        self.vn_offset = 0

        self.group_offset = 0

        # Export objects...
        if root!=None:
            self.group_offset = len(self.getGroups(self.root))-1
            self.exportObject(self.root)
        for obj in getScene().walkWorld(self.root):
            self.exportObject(obj)

        self.fhandle.close()      

    @staticmethod
    def cross(a,b):
        """ Cross product
        """
        return [a[1]*b[2] - a[2] * b[1],a[2]*b[0] - a[0] * b[2],a[0] * b[1] - a[1] * b[0]]

    @staticmethod
    def normal_from(a,b,c):
        """ Normal from 3 points
        """
        e1 = [b[0] - a[0],b[1] - a[1],b[2] - a[2]]
        e2 = [c[0] - a[0],c[1] - a[1],c[2] - a[2]]
        n = STLExporter.cross(e1,e2)
        
        length = math.sqrt(n[0]*n[0] + n[1] * n[1] + n[2] * n[2])
        if length > 0.0:
            n[0] = n[0] / length
            n[1] = n[1] / length
            n[2] = n[2] / length
        return n

    # exportObject
    def exportObject(self, obj):
        
        geom = self.convertObject(obj)
        if geom==None:
            return

        self.vt_mode = 0
        self.vn_mode = 0

        # Get the world transform to transform the vertices...
        WT = obj.worldtransform
        WT3 = WT.getMat3()

        verts = [WT*v for v in geom.verts]
        norms = []
        
        # Export normals... (TODO:  Fix this.  For now normals are calculated on the fly during facet creation)
        N = None
        info = geom.findVariable("N")
        if info!=None and info[2]==NORMAL and info[3]==1:
            N = geom.slot("N")
            for norm in N:
                norm = WT3*norm
                try:
                    norm = norm.normalize()
                except:
                    pass
                #norms.append(norm)
                #print >>self.fhandle, "vn %f %f %f"%tuple(norm)
       
        """ MeshLab STL 
        solid vcg
          facet normal 4.962093e-001 -7.952799e-001 -3.482905e-001
            outer loop
              vertex  -1.591758e+000 7.786300e-002 1.514609e+000
              vertex  -1.540988e+000 2.119300e-001 1.280815e+000
              vertex  -1.405034e+000 2.016420e-001 1.498000e+000
            endloop
          endfacet
        endsolid vcg
        """
       
        # Export facets...
        if isinstance(geom, TriMeshGeom):
            faces = self.extract_tri_faces(geom)
        else:
            faces = self.extract_poly_faces(geom)

        print >>self.fhandle, "solid vcg"
        for face in faces:
            normal = STLExporter.normal_from(verts[face[0]],verts[face[1]],verts[face[2]]) # use first vertex's normal as face normal (TODO:  avg all 3 corners)
            print >>self.fhandle, "  facet normal %f %f %f"%tuple(normal)
            print >>self.fhandle, "    outer loop"
            for index in face:
                vertex = verts[index]
                print >>self.fhandle, "       vertex %f %f %f"%tuple(vertex)
            print >>self.fhandle, "    endloop"
            print >>self.fhandle, "  endfacet"
        print >> self.fhandle, "endsolid vcg"

    # exportTriFaces
    def extract_tri_faces(self, geom):
        """Export the faces of a TriMesh geom.
        """
        faces = []
        for i in range(geom.faces.size()):
            f = geom.faces[i]
            faces.append(f)
                
        return faces
        

    # exportPolyFaces
    def extract_poly_faces(self, geom):
        """Export the faces of a polyhedron geom.
        """
        faces = []
        for i in range(geom.getNumPolys()):
            poly = geom.getPoly(i)
            faces.append(poly[0])
        return faces        
    # convertObject
    def convertObject(self, obj):
        """Converts an object into a polyhedron or trimesh if necessary.

        The return value is a GeomObject (TriMeshGeom or PolyhedronGeom)
        or None.
        """
        geom = obj.geom
        if isinstance(geom, TriMeshGeom):
            return geom
        
        if not isinstance(geom, PolyhedronGeom):
            # Try to convert into a polyhedron...
            pg = PolyhedronGeom()
            try:
                geom.convert(pg)
                geom = pg
            except:
                pass

        # Is it a PolyhedronGeom that has no polys with holes? then return
        # the geom...
        if isinstance(geom, PolyhedronGeom) and not geom.hasPolysWithHoles():
            return geom

        # Try to convert into a triangle mesh...
        tm = TriMeshGeom()
        try:
            geom.convert(tm)
            return tm
        except:
            pass

        return None
######################################################################

# Register the exporter class as a plugin class
pluginmanager.register(STLExporter)