home/diploma/lws_root/utils/objparser.py

''' a minimal wavefront .obj file parser 
tested on blender 2.58 exported .obj file
'''
from collections import defaultdict, namedtuple

Vertices = namedtuple('Vertices', 'xs  ys zs')
class Mesh(object):
    def __init__(self, fname=None):
        self.objects = defaultdict(list) # key: object/material name, value: list of triangles
        self.materials = {}
        self.lnums = {}
        self.objfile = None
        self.subobj2text = None
        self.vertices = Vertices([], [], [])
        
        if fname is not None:
            self.parseObjFile(fname)
        
    def parseObjFile(self, fname):
        self.objfile = fname
        self.subobj2text = None
        
        self.objects.clear()
        xs, ys, zs = [], [], [] # hold coords of vertices
        
        lines = open(fname).read().strip().split('\n')
        
        # initialize
        objname = 'default'
        currobj = self.objects[objname]
        self.materials[objname] = 'default'
        self.lnums[objname] = 0
        
        for lnum, l in enumerate(lines):
            if l.startswith('o '): # a new object
                objname = l[2:].strip()
                currobj = self.objects[objname]
                self.lnums[objname] = lnum
                
            elif l.startswith('usemtl '):
                matname = l[7:].strip()
                self.materials[objname] = matname
                #~ currobj = self.objects[objname]
            elif l.startswith('v '): # a vertex
                x, y, z = tuple(float(e) for e in l[2:].strip().split())
                xs.append(x)
                ys.append(y)
                zs.append(z)
                
            elif l.startswith('f '): # a face
                coords = tuple(int(e.split("/")[0]) for e in l[2:].strip().split())
                if len(coords) == 3:
                    v1, v2, v3 = coords
                    face = v1-1, v2-1, v3-1
                    currobj.append(face)
                elif len(coords) == 4:
                    v1, v2, v3, v4 = coords
                    face1 = v1-1, v2-1, v3-1
                    face2 = v3-1, v4-1, v1-1
                    currobj.extend((face1, face2))
                else:
                    #~ print 'skipping face:\n %s' % l
                    pass
        
        # remove empty objects
        for name, triangles in self.objects.items():
            if len(triangles) == 0:
                self.objects.pop(name)
                self.lnums.pop(name)
                self.materials.pop(name)
                
        self.vertices = Vertices(xs, ys, zs)
    
    def itertriangles(self, objname=None):
        ''' iterate over all triangles of total mesh or of given object
        returns 3 points in 3d space
        '''
        if objname is not None:
            vs = self.objects[objname]
        else:
            vs = []
            for _vs in self.objects.values():
                vs.extend(_vs)
        
        xs = self.vertices.xs
        ys = self.vertices.ys
        zs = self.vertices.zs
        for v in vs:
            p1 = (xs[v[0]], ys[v[0]], zs[v[0]])
            p2 = (xs[v[1]], ys[v[1]], zs[v[1]])
            p3 = (xs[v[2]], ys[v[2]], zs[v[2]])
            # return triangle given by 3 points
            yield p1, p2, p3
    
    def splitToText(self):
        if self.subobj2text is None:
            obj2text = {}
            lines = self.objfile.lines()
            lnums = list(sorted(self.lnums.items(), key=lambda e: e[1]))
            for (objname1, l1), (objname2, l2) in zip(lnums, lnums[1:] + [('END', None)]):
                s = '\n'.join('v %s %s %s' % e for e in zip(*self.vertices))
                s += '\n'
                s += '\n'.join(e.strip() for e in lines[l1:l2] if e.startswith('f '))
                
                obj2text[objname1] = s
            self.subobj2text = obj2text
        else:
            obj2text = self.subobj2text
        return obj2text
        
    #~ obj2fnames = {}
    #~ vertices = zip(*mesh.vertices)
    #~ for objname, faces in mesh.objects.items():
        #~ needed_vertices_idx = set()
        #~ for face in faces:
            #~ needed_vertices_idx.update(face)
            
        #~ vertex_lines = []
        #~ idx_mapping = {}
        #~ for i, v_idx  in enumerate(needed_vertices_idx):
            #~ idx_mapping[v_idx] = i
            #~ vertex_lines.append('v %s %s %s' % vertices[v_idx])
        
        #~ face_lines = []
        #~ for f in faces:
            #~ face_lines.append('f %s %s %s' % (idx_mapping[f[0]], idx_mapping[f[1]], idx_mapping[f[2]]))
        
        #~ objfname = tmpdir.joinpath('_%s.obj' % objname)
        #~ objfname.write_text('\n'.join(vertex_lines + face_lines))

        #~ obj2fnames[objname] = objfname