WavefrontHelper

 1# wavefront.py
 2import numpy as np
 3
 4
 5class WavefrontOBJ:
 6    def __init__( self, default_mtl='default_mtl' ):
 7        self.path      = None               # path of loaded object
 8        self.mtllibs   = []                 # .mtl files references via mtllib
 9        self.mtls      = [ default_mtl ]    # materials referenced
10        self.mtlid     = []                 # indices into self.mtls for each polygon
11        self.vertices  = []                 # vertices as an Nx3 or Nx6 array (per vtx colors)
12        self.normals   = []                 # normals
13        self.texcoords = []                 # texture coordinates
14        self.polygons  = []                 # M*Nv*3 array, Nv=# of vertices, stored as vid,tid,nid (-1 for N/A)
15
16
17def load_obj( filename: str, default_mtl='default_mtl', triangulate=False ) -> WavefrontOBJ:
18    """Reads a .obj file from disk and returns a WavefrontOBJ instance
19
20    Handles only very rudimentary reading and contains no error handling!
21
22    Does not handle:
23        - relative indexing
24        - subobjects or groups
25        - lines, splines, beziers, etc.
26    """
27    # parses a vertex record as either vid, vid/tid, vid//nid or vid/tid/nid
28    # and returns a 3-tuple where unparsed values are replaced with -1
29    def parse_vertex( vstr ):
30        vals = vstr.split('/')
31        vid = int(vals[0])-1
32        tid = int(vals[1])-1 if len(vals) > 1 and vals[1] else -1
33        nid = int(vals[2])-1 if len(vals) > 2 else -1
34        return (vid,tid,nid)
35
36    with open( filename, 'r' ) as objf:
37        obj = WavefrontOBJ(default_mtl=default_mtl)
38        obj.path = filename
39        cur_mat = obj.mtls.index(default_mtl)
40        for line in objf:
41            toks = line.split()
42            if not toks:
43                continue
44            if toks[0] == 'v':
45                obj.vertices.append( [ float(v) for v in toks[1:]] )
46            elif toks[0] == 'vn':
47                obj.normals.append( [ float(v) for v in toks[1:]] )
48            elif toks[0] == 'vt':
49                obj.texcoords.append( [ float(v) for v in toks[1:]] )
50            elif toks[0] == 'f':
51                poly = [ parse_vertex(vstr) for vstr in toks[1:] ]
52                if triangulate:
53                    for i in range(2,len(poly)):
54                        obj.mtlid.append( cur_mat )
55                        obj.polygons.append( (poly[0], poly[i-1], poly[i] ) )
56                else:
57                    obj.mtlid.append(cur_mat)
58                    obj.polygons.append( poly )
59            elif toks[0] == 'mtllib':
60                obj.mtllibs.append( toks[1] )
61            elif toks[0] == 'usemtl':
62                if toks[1] not in obj.mtls:
63                    obj.mtls.append(toks[1])
64                cur_mat = obj.mtls.index( toks[1] )
65        return obj
66
67
68def save_obj( obj: WavefrontOBJ, filename: str ):
69    """Saves a WavefrontOBJ object to a file
70
71    Warning: Contains no error checking!
72
73    """
74    with open( filename, 'w' ) as ofile:
75        for mlib in obj.mtllibs:
76            ofile.write('mtllib {}\n'.format(mlib))
77        for vtx in obj.vertices:
78            ofile.write('v '+' '.join(['{}'.format(v) for v in vtx])+'\n')
79        for tex in obj.texcoords:
80            ofile.write('vt '+' '.join(['{}'.format(vt) for vt in tex])+'\n')
81        for nrm in obj.normals:
82            ofile.write('vn '+' '.join(['{}'.format(vn) for vn in nrm])+'\n')
83        if not obj.mtlid:
84            obj.mtlid = [-1] * len(obj.polygons)
85        poly_idx = np.argsort( np.array( obj.mtlid ) )
86        cur_mat = -1
87        for pid in poly_idx:
88            if obj.mtlid[pid] != cur_mat:
89                cur_mat = obj.mtlid[pid]
90                ofile.write('usemtl {}\n'.format(obj.mtls[cur_mat]))
91            pstr = 'f '
92            for v in obj.polygons[pid]:
93                # UGLY!
94                vstr = '{}/{}/{} '.format(v[0]+1,v[1]+1 if v[1] >= 0 else 'X', v[2]+1 if v[2] >= 0 else 'X' )
95                vstr = vstr.replace('/X/','//').replace('/X ', ' ')
96                pstr += vstr
97            ofile.write( pstr+'\n')