ImodExport.py

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

        # Determine object type
        if objin.Objects[iObject].objType == 'closed':
            mesh = get_mesh(objin, iObject)
            print "Processing closed type object."
        elif objin.Objects[iObject].objType == 'scattered':
            mesh = objin.Objects[iObject]
            print "Processing scattered type object."

        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)

        # Update scale and trans based on MINX info.
        if objin.minx_set:
            scale = list(objin.minx_cscale)
            trans = objin.minx_ctrans

        # If model header info is provided, override the scale provided by MINX
        # info and use this instead. To check if model header has been given,
        # check that the units are set to anything but pixels, which has a value
        # of zero. 
        if objin.units:
            TOL = 1
            # Convert scales to Angstroms
            scalexy = objin.pixelSizeXY / (10 ** -objin.units) * (10 ** 10)
            scalez = objin.pixelSizeZ / (10 ** -objin.units) * (10 ** 10)
            dxy = scalexy - scale[0]
            dz = scalez - scale[2]
            scale[0] = scalexy
            scale[1] = scalexy
            scale[2] = scalez
            if dxy > TOL or dz > TOL:
                print ("WARNING: Scale information of the model and MRC stack "
                      "it was created on are significantly different. The "
                      "model will not load well on the MRC stack in Amira.")
    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 "VRML scale: {}".format(scale)
    print "VRML translation: {}".format(trans)
    print "VRML Z-scale: {}".format(zscale)
    nameStr = 'obj{0}'.format(iObject)
    if name:
        for x in name.split():
            nameStr = nameStr + '_' + x 
    fid = open(fnameout, 'w+')

    obj_type = type(mesh).__name__    

    # 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')

    if obj_type == 'ImodMesh':
        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)
        print "VRML # of vertices: {}".format(nvertlist)
        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)
        print "VRML # of indices: {}".format(nindexlist)
        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')

    elif obj_type == 'ImodObject':
        fid.write('    DEF {0}_SPHERES Group {{   # GROUP OF POINTS:\n'.format(
            nameStr))
        fid.write('      children [\n')
        for icont in range(mesh.nContours):
            for ipoint in range(mesh.Contours[icont].nPoints):
                # Get the x,y,z coordinates of the point and scale them
                pt_coords = list(mesh.Contours[icont].points[ipoint*2+ipoint:
                    ipoint*2+ipoint+3])
                pt_coords[0] = pt_coords[0] * scale[0] + trans[0]
                pt_coords[1] = pt_coords[1] * scale[1] + trans[1]
                pt_coords[2] = pt_coords[2] * scale[2] + trans[2]

                # Get the point radius
                if mesh.Contours[icont].size_set:
                    radius = mesh.Contours[icont].size_vals[ipoint]
                else:
                    radius = mesh.pdrawsize 

                fid.write('        DEF {0}_cont{1}_pt{2} Transform {{\n'.format(
                    nameStr, icont+1, ipoint+1))
                fid.write('          translation {0} {1} {2}\n'.format(
                    pt_coords[0], pt_coords[1], pt_coords[2]))
                fid.write('          children [ Shape {\n')
                fid.write('              appearance USE MAT_{0}\n'.format(
                    nameStr))
                fid.write(('              geometry DEF cont_{0}_{1} Sphere '
                           '{{ radius {2} }} }} ]\n'.format(icont+1, ipoint+1,
                           radius * scale[0])))
                fid.write('        }\n')
        fid.write('      ]\n')
        fid.write('    }\n')
        fid.write('  ]\n')
        fid.write('}\n')
        fid.write('  ]\n')
        fid.write('}\n')
 

    fid.close()
    print "{0} written.".format(fnameout)