home/diploma/lws_root/runphoton.py

from zipfile import ZipFile
import subprocess
import time
from datetime import datetime

import logging
from collections import namedtuple, defaultdict

try:
    import rpyc
except ImportError:
    rpyc = None

import hashlib
from dateutil import parser
from utils.path import path
from utils.xml import StringToTree, FileToTree, subelement, TreeToFile, prettyPrint
from utils.objparser import Mesh
from utils.materials import materialsAsXml

log = logging.getLogger('runphoton')
Resolution = namedtuple('Resolution', 'x y z')
BBox = namedtuple('BBox', 'x1 x2 y1 y2 z1 z2')

MESH_CACHE = {}
def prepareControlFile(ctrlfile, objfile, materials, light, resolution=None, bbox=None,
                        density=None, numphotons=None, outname=None, disabledObjects=None, 
                        mesh=None, tmpdir=None, sceneName=None, rootAttribs=None, remapMaterials=None):
    ''' read ctrlfile and fill with values (materials, light and simulation params)
    @param materials: the material parameters
    @type materials: dict with key: matname and value namedtuple(reflect alpha ior)
    @param light: the params of the light source
    @type light: namedtuple(power x y z)
    @param objfile: file containing the scene in wavefront .obj format
    @type objfile: path
    @param mesh: a mesh from an allready parsed obj file
    @type mesh: Mesh()
    @param tmpdir: directory where the generated files should be stored, if None use parent dir of ctrlfile
    @type tmpdir: path
    @param disabledObjects: objects by name, that should be removed from ctrlfile
    @type disabledObjects: set(str)
    '''
    
    
    
    if mesh is None:
        cache_key = (objfile, objfile.mtime)
        if cache_key in MESH_CACHE:
            mesh = MESH_CACHE[cache_key]
        else:
            mesh = Mesh()
            mesh.parseObjFile(objfile)
            MESH_CACHE[cache_key] = mesh
            
    if tmpdir is None:
        tmpdir = ctrlfile.parent.joinpath('_tmp')
        
    if not tmpdir.exists():
        tmpdir.mkdir()
        
    tree = FileToTree(ctrlfile)
    
    know_materials = set()
    # add <brdf> material definitions to <scene>
    for e in materialsAsXml(materials).getroot().xpath('brdf'):
        know_materials.add(e.attrib['name'])
        tree.getroot().insert(0, e)
    
    obj2fnames = {}
    for objname, text in mesh.splitToText().items():
        objfname = tmpdir.joinpath('_%s.obj' % objname)
        objfname.write_text(text)
        obj2fnames[objname] = objfname
    
    for objname in mesh.objects:
        if disabledObjects is not None and objname in disabledObjects:
            continue
        
        matname = mesh.materials.get(objname, 'default')
        
        # XXX: hackish, only for thesis eval
        if remapMaterials is not None:
            matname = remapMaterials[matname]
            
        if matname == '(null)':
            log.warn('no material defined for object %s' % objname)
            continue
        elif not matname in know_materials:
            log.warn('unknown material: %s' % matname)
            continue
            
        attribs = dict([('m%s%s' % (i, j), 0.0 if i != j else 1.0) for i in range(4) for j in range(4)])
        tEl = subelement(tree.getroot(), 'transform', attribs=attribs)
        objEl = subelement(tEl, 'object', attribs={'name': objname, 'brdf_name': matname})
        meshEl = subelement(objEl, 'mesh', attribs={'autosmooth': 180})
        hash = hashlib.sha1(obj2fnames[objname].text()).hexdigest()
        subelement(meshEl, 'include', attribs={'file': obj2fnames[objname].name, 'hash': hash})
    
    lightEl = tree.getroot().xpath('light')[0]
    lightEl.attrib['power'] = str(light.power)
    
    fromEl = lightEl.xpath('from')[0]
    fromEl.attrib['x'] = str(light.x)
    fromEl.attrib['y'] = str(light.y)
    fromEl.attrib['z'] = str(light.z)
    
    if outname is not None:
        outputEl = tree.getroot().xpath('rwp/volumerender/output')[0]
        outputEl.attrib['filename'] = outname
    
    if density is not None:
        radiusEl = tree.getroot().xpath('rwp/volumerender/density_estimator/radius')[0]
        radiusEl.attrib['value'] = str(density)
    
    if resolution is not None:
        sizeEl = tree.getroot().xpath('rwp/volumerender/size')[0]
        sizeEl.attrib['x'] = str(resolution.x)
        sizeEl.attrib['y'] = str(resolution.y)
        sizeEl.attrib['z'] = str(resolution.z)
    
    if bbox is not None:
        minEl = tree.getroot().xpath('rwp/simulation/bbox/min')[0]
        minEl.attrib['x'] = str(bbox.x1)
        minEl.attrib['y'] = str(bbox.y1)
        minEl.attrib['z'] = str(bbox.z1)
        
        maxEl = tree.getroot().xpath('rwp/simulation/bbox/max')[0]
        maxEl.attrib['x'] = str(bbox.x2)
        maxEl.attrib['y'] = str(bbox.y2)
        maxEl.attrib['z'] = str(bbox.z2)
    
    if numphotons is not None:
        photonsEl = tree.getroot().xpath('rwp/simulation/photons')[0]
        photonsEl.attrib['value'] = str(numphotons)
    
    if rootAttribs is not None:
        tree.getroot().attrib.update(rootAttribs)
        
    material_ctrlfile = tmpdir.joinpath('_' + ctrlfile.name)
    material_ctrlfile.write_text(prettyPrint(tree))
    
    return material_ctrlfile


def runphoton(ctrlfile, verbose=True):
    ''' run photon remotely via rpyc by using a (prepared) ctrlfile
    
    return ctrlfile
    '''
    
    
    ctrlfile = path(ctrlfile)
    workdir = ctrlfile.parent
    conn = rpyc.classic.connect("127.0.0.1", port=18812)
    remote_open = conn.modules.__builtin__.open
    remote_check_output = conn.modules.subprocess.check_output
    remote_getsize = conn.modules.os.path.getsize
    remote_ZipFile = conn.modules.zipfile.ZipFile
    
    tree = FileToTree(ctrlfile)
    
    s = ctrlfile.text()
    log.info('transmitting control file "%s" [%.2f kb] ...' % (ctrlfile, len(s) / 1024.0))
    remote_open('tmp/%s' % ctrlfile.name, 'wb').write(s)
    
    objfiles = tree.xpath('//mesh/include[@file]/@file')
    
    # full remote hash building
    conn.execute("import hashlib")
    remote_hash = lambda f:  conn.eval("hashlib.sha1(open('tmp/%s').read()).hexdigest()" % f)
    try:
        file2hash = dict([(f, remote_hash(f)) for f in objfiles])
    except Exception:
        # probably a file does not exists
        print 'got not remote hashes for objfiles'
        file2hash = defaultdict(str)
        
    
    # send local files to server
    for objfile in objfiles:
        localfile = workdir.joinpath(objfile)
        if hashlib.sha1(open(localfile).read()).hexdigest() == file2hash[objfile]:
            print '%s is uptodate' % objfile
            continue
        
        s = localfile.text()
        log.info('transmitting obj file "%s" [%.2f kb] ...' % (objfile, len(s) / 1024.0))
        remote_open('tmp/%s' % objfile, 'wb').write(s)
        
    # run photon 
    scenefile = 'tmp/%s' % ctrlfile.name
    log.info('running photon remotely on %s' % scenefile)
    
    try:
        cmd = ['photon', '-s', scenefile]
        
        output = remote_check_output(cmd)
        if verbose:
            print output
    except subprocess.CalledProcessError, e:
        log.error('error during executing: %s' % ' '.join(cmd))
        log.error('output:\n%s' % e.output)
        #~ raise
    except Exception:
        print 'XXXXX'
    
    outfname = tree.getroot().xpath('string(rwp/volumerender/output/@filename)')
    assert outfname != ''
    
    # zip .raw file to reduce download size
    rawfile = '%s.raw' % outfname
    log.info('zipping %s' % rawfile)
    zf = remote_ZipFile(rawfile + '.zip', mode='w', compression=conn.modules.zipfile.ZIP_DEFLATED)
    zf.write(rawfile)
    zf.close()

    # download result files
    for ext in '.raw.zip', '.dat':
        fn = outfname + ext
        log.info('transmitting %s file [%.2f kb] ...' % (fn, remote_getsize(fn) / 1024.0))
        t = time.time()
        s = remote_open(fn, 'rb').read()
        log.info('in %.2f sec' % (time.time() - t))
        
        fn = workdir.joinpath('%s%s' % (outfname, ext))
        open(fn, 'wb').write(s)

    fn = workdir.joinpath('%s.raw.zip' % outfname)
    log.info('unzipping %s' % fn)
    zf = ZipFile(fn)
    s = zf.read(rawfile)
    open(workdir.joinpath('%s.raw' % outfname), 'wb').write(s)

    conn.close()
    
    # return .dat file
    return {'datfile': workdir.joinpath('%s.dat' % outfname), 'ctrlfile': ctrlfile}

if __name__ == '__main__':
    name = '../maps/onewall/One-Wall'
    runphoton(name)