PyIMOD/ImodExport.py at master · coleslaw481/PyIMOD · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from __future__ import division

import os
from .utils import is_string

def ImodExport(objin, fnameout, **kwargs):
    iObject = kwargs.get('object', 0)
    is_string(fnameout, 'Output filename')
    objType = type(objin).__name__
    mats = []
    scale = [1, 1, 1]
    trans = [0, 0, 0]
    if objType == 'ImodModel':
        if iObject >= 1:
            iObject-=1
        mesh = get_mesh(objin, iObject)
        name = objin.Objects[iObject].name
        mats.append(objin.Objects[iObject].ambient / 255)
        mats.append(objin.Objects[iObject].diffuse / 255)
        mats.append(objin.Objects[iObject].specular / 255)
        mats.append(objin.Objects[iObject].shininess / 255)
        mats.append(objin.Objects[iObject].transparency)
        if objin.minx_set:
            scale = objin.minx_cscale
            trans = objin.minx_ctrans
    else:
        raise ValueError('input is not a valid class type.')

    fext = os.path.splitext(fnameout)[1]

    if fext.lower() == '.vrml' or fext.lower() == '.wrl':
        export_vrml2(mesh, iObject, name, mats, scale, trans, fnameout)

def get_mesh(imodModel, iObject):
    nObjects = imodModel.nObjects
    if iObject > nObjects:
        raise ValueError('Value specified by object kwarg exceeds nObjects.')
    nMeshes = imodModel.Objects[iObject].nMeshes
    if nMeshes > 1:
        raise ValueError('Object {0} contains > 1 mesh.'.format(iObject+1))
    mesh = imodModel.Objects[iObject].Meshes[0]
    return mesh

def export_vrml2(mesh, iObject, name, mats, scale, trans, fnameout):
    iObject+=1
    zscale = scale[2] / scale[0]
    print scale, trans, zscale
    nameStr = 'obj{0}'.format(iObject)
    if name:
        for x in name.split():
            nameStr = nameStr + '_' + x
    fid = open(fnameout, 'w+')

    # Write VRML 2.0 header data
    fid.write('#VRML V2.0 utf8\n')
    fid.write('#Generated by pyimod\n\n')
    fid.write('DEF imod_model Transform {\n')
    fid.write('  children [\n\n')
    fid.write('#MATERIAL FOR OBJECT {0}:\n'.format(iObject))
    fid.write('Shape {\n')
    fid.write('  appearance DEF {0} Appearance {{\n'.format('MAT_' + nameStr))
    fid.write('    material Material {\n')
    fid.write('      ambientIntensity {0}\n'.format(mats[0]))
    fid.write('      diffuseColor {0} {0} {0}\n'.format(mats[1]))
    fid.write('      specularColor {0} {0} {0}\n'.format(mats[2]))
    fid.write('      emissiveColor 0 0 0\n')
    fid.write('      shininess {0}\n'.format(mats[3]))
    fid.write('      transparency {0}\n'.format(mats[4]))
    fid.write('    }\n')
    fid.write('  }\n')
    fid.write('}\n\n')
    fid.write('#DATA FOR OBJECT {0}:\n'.format(iObject))
    fid.write('DEF {0} Transform {{\n'.format(nameStr))
    fid.write('  children [\n')
    fid.write('    Shape {   #MESH\n')
    fid.write('      appearance USE {0}\n'.format('MAT_' + nameStr))
    fid.write('      geometry DEF {0} IndexedFaceSet {{\n'.format(nameStr))
    fid.write('        ccw FALSE\n')
    fid.write('        solid FALSE\n')
    fid.write('        creaseAngle 1.56207\n')
    fid.write('        coord Coordinate {\n')
    fid.write('          point [   # list of all points in mesh\n')

    # Write VRML 2.0 mesh vertex data
    C = 0
    nvertlist = len(mesh.vertices)
    while C < nvertlist:
        x, y, z = [mesh.vertices[C+x] for x in range(3)]
        x = x * scale[0] + trans[0]
        y = y * scale[1] + trans[1]
        z = z * scale[2] + trans[2]
        x = '{0:.1f}'.format(x) if x % 1 else int(x)
        y = '{0:.1f}'.format(y) if y % 1 else int(y)
        z = '{0:.1f}'.format(z) if z % 1 else int(z)
        fid.write('            {0} {1} {2},\n'.format(x, y, z))
        C+=6
    fid.write('          ]\n')
    fid.write('        }\n')

    # Write VRML 2.0 mesh index data
    C = 0
    nindexlist = len(mesh.indices)
    fid.write('        coordIndex [   # connect triangles\n')
    while C < nindexlist:
        if mesh.indices[C] < 0:
            C+=1
            continue
        else:
            fid.write('          {0},{1},{2},-1,\n'.format(
                int(mesh.indices[C+2] / 2),
                int(mesh.indices[C+1] / 2),
                int(mesh.indices[C] / 2)))
            C+=3
    fid.write('        ]\n')
    fid.write('      }\n')
    fid.write('    }\n')
    fid.write('  ]\n')
    fid.write('}\n\n')
    fid.write('  ]\n')
    fid.write('}\n')
    fid.close()
    print "{0} written.".format(fnameout)