root/trunk/RBRapier/RBRapier/Formats/Lightwave/Loader.py

#!/usr/bin/env python
##~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
##~ License 
##~ 
##- The RuneBlade Foundation library is intended to ease some 
##- aspects of writing intricate Jabber, XML, and User Interface (wxPython, etc.) 
##- applications, while providing the flexibility to modularly change the 
##- architecture. Enjoy.
##~ 
##~ Copyright (C) 2002  TechGame Networks, LLC.
##~ 
##~ This library is free software; you can redistribute it and/or
##~ modify it under the terms of the BSD style License as found in the 
##~ LICENSE file included with this distribution.
##~ 
##~ TechGame Networks, LLC can be reached at:
##~ 3578 E. Hartsel Drive #211
##~ Colorado Springs, Colorado, USA, 80920
##~
##~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#~ Imports 
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

import os
import struct
import chunk

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#~ Definitions 
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class _LWOSubChunk(chunk.Chunk):
    """Adapt python's chunk.Chunk class to be able to read LWO sub-chunks..."""
    def __init__(self, file, align = 1, bigendian = 1, inclheader = 0):
        import struct
        self.closed = 0
        self.align = align      # whether to align to word (2-byte) boundaries
        if bigendian:
            strflag = '>'
        else:
            strflag = '<'
        self.file = file
        self.chunkname = file.read(4)
        if len(self.chunkname) < 4:
            raise EOFError
        try:
            self.chunksize = struct.unpack(strflag+'h', file.read(2))[0]
        except struct.error:
            raise EOFError
        if inclheader:
            self.chunksize = self.chunksize - 6 # subtract header
        self.size_read = 0
        try:
            self.offset = self.file.tell()
        except (AttributeError, IOError):
            self.seekable = 0
        else:
            self.seekable = 1

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

def _ReadChunkString(chunk, readstep=2):
    result = ''
    while result[-1:] != '\0':
        result += chunk.read(readstep)
    result = result[:result.find('\0')]
    return result

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class _LWOSurface(object):
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #~ Constants / Variables / Etc. 
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    #BaseColor = (255, 255, 255, 0)
    #Flag = 0

    #iLuminous = 0
    #iDiffuse = 0
    #iSpecular = 0
    #iReflection = 0
    #iTransparency = 0

    #Luminous = 0.
    #Diffuse = 0.
    #Specular = 0.
    #Reflection = 0.
    #Transparency = 0.

    #Shininess = 1

    #ReflectiveMode = 3
    #ReflectiveImage = ''
    #ReflectiveSeamAngle = 0.
    #ReflectiveIndex = 1.

    #EdgeTransparency = 1.
    #SmoothingAngle = math.pi/2.

    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #~ Public Methods 
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    def __init__(self, Name):
        self.Name = Name
        self.Faces = []

    def ReadFormat(self, SurfaceDefChunk, TaskProgress=None):
        FormatMap = {
            'COLR': ('BaseColor', '>BBBB'),
            'FLAG': ('Flag', '>H'),

            'LUMI': ('iLuminous', '>h'),
            'DIFF': ('iDiffuse', '>h'),
            'SPEC': ('iSpecular', '>h'),
            'REFL': ('iReflection', '>h'),
            'TRAN': ('iTransparency', '>h'),

            'VLUM': ('Luminous', '>f'),
            'VDIF': ('Diffuse', '>f'),
            'VSPC': ('Specular', '>f'),
            'VRFL': ('Reflection', '>f'),
            'VTRN': ('Transparency', '>f'),

            'GLOS': ('Shininess', '>h'),

            'RFLT': ('ReflectiveMode', '>h'),
            'RIMG': ('ReflectiveImage', _ReadChunkString),
            'RSAN': ('ReflectiveSeamAngle', '>f'),
            'RIND': ('ReflectiveIndex', '>f'),

            'EDGE': ('EdgeTransparency', '>f'),
            'SMAN': ('SmoothingAngle', '>f'),
            }
        try:
            while 1:
                SubChunk = _LWOSubChunk(SurfaceDefChunk)
                name, fmt = FormatMap.get(SubChunk.getname(), (None, None))
                if isinstance(fmt, str):
                    value = struct.unpack(fmt, SubChunk.read(SubChunk.getsize()))
                    if len(value) == 1: value = value[0]
                elif callable(fmt):
                    value = fmt(SubChunk)
                else:
                    print "[%s] UNKNOWN: %s %r" % (self.Name, SubChunk.getname(), SubChunk.read())
                    SubChunk.skip()
                    continue
                if name: 
                    setattr(self, name, value)
                    #print name, value
        except EOFError:
            pass # end of chunk

        Flag = getattr(self, 'Flag', 0)
        self.Luminous = Flag & 1; Flag >>= 1
        self.Outline = Flag & 1; Flag >>= 1
        self.Smoothing = Flag & 1; Flag >>= 1
        self.ColorHighlights = Flag & 1; Flag >>= 1
        self.ColorFilter = Flag & 1; Flag >>= 1
        self.OpaqueEdge = Flag & 1; Flag >>= 1
        self.TransparentEdge = Flag & 1; Flag >>= 1
        self.SharpTerminator = Flag & 1; Flag >>= 1
        self.DoubleSided = Flag & 1; Flag >>= 1
        self.Additive = Flag & 1; Flag >>= 1
        self.ShadowAlpha = Flag & 1; Flag >>= 1

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class _LWOFace(object):
    def ReadFormat(self, FacesChunk, TaskProgress=None):
        NumIndices = struct.unpack(">H", FacesChunk.read(2))[0]
        self.IndexList = [struct.unpack(">H", FacesChunk.read(2))[0] for each in xrange(NumIndices)]
        self.SurfaceIndex = struct.unpack(">h", FacesChunk.read(2))[0]
        ByteLen = len(self.IndexList)*2 + 4
        if self.SurfaceIndex < 0:
            self.SurfaceIndex = -self.SurfaceIndex
            # This polygon has "detail polygons", but the reference says to ignore them
            # as they are mostly obsolete
            NumDetailPolys = struct.unpack(">H", FacesChunk.read(2))[0]
            ByteLen += 2
            # we exit at this point so the detail polys are handled as normal polys... call me lazy!
        self.SurfaceIndex -= 1 # Make index zero based
        return ByteLen

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class LightwaveLWOBLoader(object):
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #~ Constants / Variables / Etc. 
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    RaiseMissingElements = 0
    WarnMissingElements = 0

    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #~ Public Methods 
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    def ReadFormat(self, LightwaveLWOFile, TaskProgress=None):
        self.Name = None
        self.Vertices = []
        self.Surfaces = {}

        ChunkTypeTable = {
            'SRFS': self._ReadSurfaces,
            'PNTS': self._ReadPoints,
            'POLS': self._ReadFaces,
            'SURF': self._ReadSurfaceDefinition,
            }

        self._FilePath = os.path.split(LightwaveLWOFile.name)[0]
        FORMChunk = chunk.Chunk(LightwaveLWOFile)
        if FORMChunk.getname() != 'FORM':
            raise ValueError, "File does not contain a FORM chunk to contain a Lightwave3D Object"

        if TaskProgress: 
            TaskProgress = TaskProgress.NewSubtask('Load Lightwave Object', FORMChunk.tell(), FORMChunk.getsize())
        else: TaskProgress = None

        FormType = FORMChunk.read(4)
        if FormType == 'LWOB':
            pass # This is what we're looking for
        elif FormType != 'LWO2':
            raise NotImplemented, "File contains a Lightwave3D 2 Object (LWO2), which is not (yet) supported."
        else: raise ValueError, "File does not contain a Lightwave3D Object"

        try:
            while 1:
                NextChunk = chunk.Chunk(FORMChunk)
                try: 
                    ex = ChunkTypeTable[NextChunk.getname()]
                except KeyError:
                    if self.RaiseMissingElements:
                        raise NotImplemented, "Unknown LightwaveOBJ chunk type '%s'" % NextChunk.getname()
                    elif self.WarnMissingElements:
                        print "Unknown LightwaveOBJ chunk type '%s'" % NextChunk.getname()
                    NextChunk.skip()
                else:
                    ex(NextChunk)
                if TaskProgress: 
                    TaskProgress.Progress = FORMChunk.tell()
        except EOFError:
            pass # we're done

    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    #~ Protected Methods 
    #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    def _ReadSurfaces(self, SurfacesChunk):
        self.SurfaceNames = filter(None, SurfacesChunk.read().split('\0'))
        self.Surfaces = dict([(name, _LWOSurface(name)) for name in self.SurfaceNames])

    def _ReadPoints(self, PointsChunk):
        NumPoints = PointsChunk.getsize() / 12
        self.Vertices = [struct.unpack(">3f", PointsChunk.read(12)) for idx in xrange(NumPoints)]

    def _ReadFaces(self, FacesChunk):
        BytesLeft = FacesChunk.getsize()
        while BytesLeft > 0:
            face = _LWOFace()
            BytesConsumed = face.ReadFormat(FacesChunk)
            FaceSurfaceName = self.SurfaceNames[face.SurfaceIndex-1]
            self.Surfaces[FaceSurfaceName].Faces.append(face)
            BytesLeft -= BytesConsumed
        assert BytesLeft == 0

    def _ReadSurfaceDefinition(self, SurfaceDefChunk):
        SurfaceName = _ReadChunkString(SurfaceDefChunk)
        self.Surfaces[SurfaceName].ReadFormat(SurfaceDefChunk)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#~ Testing 
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

if __name__=='__main__':
    print "Testing..."
    test = LightwaveLWOLoader()
    #test.ReadFormat(open('data/ki162a.lwo', 'rb'))
    test.ReadFormat(open('data/fi110a.lwo', 'rb'))
    print "Test complete."