model-converter-python/d3/model/formats/ply.py

import os
import sys
import PIL
import struct
from ..basemodel import ModelParser, TextModelParser, Exporter, Vertex, Face, Color, FaceVertex, TexCoord, Material

class UnkownTypeError(Exception):
    def __init__(self, message):
        self.message  = message

def is_ply(filename):
    return filename[-4:] == '.ply'

# List won't work with this function
def _ply_type_size(type):
    if type == 'char' or type == 'uchar':
        return 1
    elif type == 'short' or type == 'ushort':
        return 2
    elif type == 'int' or type == 'uint':
        return 4
    elif type == 'float':
        return 4
    elif type == 'double':
        return 8
    else:
        raise UnkownTypeError('Type ' + type + ' is unknown')

def ply_type_size(type):
    split = type.split()

    if len(split) == 1:
        return [_ply_type_size(type)]
    else:
        if split[0] != 'list':
            print('You have multiple types but it\'s not a list...', file=sys.stderr)
            sys.exit(-1)
        else:
            return list(map(lambda a: _ply_type_size(a), split[1:]))


def bytes_to_element(type, bytes, byteorder = 'little'):
    if type == 'char':
        return ord(struct.unpack('<b', bytes)[0])
    if type == 'uchar':
        return ord(struct.unpack('<c', bytes)[0])
    elif type == 'short':
        return struct.unpack('<h', bytes)[0]
    elif type == 'ushort':
        return struct.unpack('<H', bytes)[0]
    elif type == 'int':
        return struct.unpack('<i', bytes)[0]
    elif type == 'uint':
        return struct.unpack('<I', bytes)[0]
    elif type == 'float':
        return struct.unpack('<f', bytes)[0]
    elif type == 'double':
        return struct.unpack('<d', bytes)[0]
    else:
        raise UnkownTypeError('Type ' + type + ' is unknown')

class PLYParser(ModelParser):

    def __init__(self, up_conversion = None):
        super().__init__(up_conversion)
        self.counter = 0
        self.elements = []
        self.materials = []
        self.inner_parser = PLYHeaderParser(self)
        self.beginning_of_line = ''
        self.header_finished = False

    def parse_bytes(self, bytes, byte_counter):
        if self.header_finished:
            self.inner_parser.parse_bytes(self.beginning_of_line + bytes, byte_counter - len(self.beginning_of_line))
            self.beginning_of_line = b''
            return

        # Build lines for header and use PLYHeaderParser
        current_line = self.beginning_of_line
        for (i, c) in enumerate(bytes):
            char = chr(c)
            if char == '\n':
                self.inner_parser.parse_line(current_line)
                if  current_line == 'end_header':
                    self.header_finished = True
                    self.beginning_of_line = bytes[i+1:]
                    return
                current_line = ''
            else:
                current_line += chr(c)
        self.beginning_of_line = current_line

class PLYHeaderParser:
    def __init__(self, parent):
        self.current_element = None
        self.parent = parent
        self.content_parser = None

    def parse_line(self, string):
        split = string.split()
        if string == 'ply':
            return

        elif split[0] == 'format':
            if split[2] != '1.0':
                print('Only format 1.0 is supported', file=sys.stderr)
                sys.exit(-1)

            if split[1] == 'ascii':
                self.content_parser = PLY_ASCII_ContentParser(self.parent)
            elif split[1] == 'binary_little_endian':
                self.content_parser = PLYLittleEndianContentParser(self.parent)
            elif split[1] == 'binary_big_endian':
                self.content_parser = PLYBigEndianContentParser(self.parent)
            else:
                print('Only ascii, binary_little_endian and binary_big_endian are supported', \
                      file=sys.stderr)
                sys.exit(-1)

        elif split[0] == 'element':
            self.current_element = PLYElement(split[1], int(split[2]))
            self.parent.elements.append(self.current_element)

        elif split[0] == 'property':
            self.current_element.add_property(split[-1], ' '.join(split[1:-1]))

        elif split[0] == 'end_header':
            self.parent.inner_parser = self.content_parser

        elif split[0] == 'comment' and split[1] == 'TextureFile':
            material = Material('mat' + str(len(self.parent.materials)))
            self.parent.materials.append(material)

            try:
                material.map_Kd = PIL.Image.open(os.path.join(os.path.dirname(self.parent.path), split[2]))
            except ImportError:
                pass

class PLYElement:
    def __init__(self, name, number):
        self.name = name
        self.number = number
        self.properties = []

    def add_property(self, name, type):
        self.properties.append((name, type))

class PLY_ASCII_ContentParser:
    def __init__(self, parent):
        self.parent = parent
        self.element_index = 0
        self.counter = 0
        self.current_element = None
        self.beginning_of_line = ''

    def parse_bytes(self, bytes, byte_counter):
        current_line = self.beginning_of_line
        for (i, c) in enumerate(bytes):
            char = chr(c)
            if char == '\n':
                self.parse_line(current_line)
                current_line = ''
            else:
                current_line += chr(c)
        self.beginning_of_line = current_line


    def parse_line(self, string):

        if self.current_element is None:
            self.current_element = self.parent.elements[0]

        split = string.split()
        color = None

        if self.current_element.name == 'vertex':

            vertex = Vertex()
            red = None
            blue = None
            green = None
            alpha = None

            offset = 0
            for property in self.current_element.properties:

                if property[0] == 'x':
                    vertex.x = float(split[offset])
                elif property[0] == 'y':
                    vertex.y = float(split[offset])
                elif property[0] == 'z':
                    vertex.z = float(split[offset])
                elif property[0] == 'red':
                    red = float(split[offset]) / 255
                elif property[0] == 'green':
                    green = float(split[offset]) / 255
                elif property[0] == 'blue':
                    blue = float(split[offset]) / 255
                elif property[0] == 'alpha':
                    alpha = float(split[offset]) / 255

                offset += 1

            self.parent.add_vertex(vertex)

            if red is not None:
                color = Color(red, blue, green)
                self.parent.add_color(color)

        elif self.current_element.name == 'face':

            faceVertexArray = []
            current_material = None

            # Analyse element
            offset = 0
            for property in self.current_element.properties:

                if property[0] == 'vertex_indices':
                    for i in range(int(split[offset])):
                        faceVertexArray.append(FaceVertex(int(split[i+offset+1])))
                    offset += int(split[0]) + 1

                elif property[0] == 'texcoord':
                    offset += 1
                    for i in range(3):
                        # Create corresponding tex_coords
                        tex_coord = TexCoord().from_array(split[offset:offset+2])
                        offset += 2
                        self.parent.add_tex_coord(tex_coord)
                        faceVertexArray[i].tex_coord = len(self.parent.tex_coords) - 1

                elif property[0] == 'texnumber':
                    current_material = self.parent.materials[int(split[offset])]

            face = Face(*faceVertexArray)
            face.material = current_material
            self.parent.add_face(face)

        self.counter += 1

        if self.counter == self.current_element.number:
            self.next_element()

    def next_element(self):
        self.element_index += 1
        if self.element_index < len(self.parent.elements):
            self.current_element = self.parent.elements[self.element_index]

class PLYLittleEndianContentParser:
    def __init__(self, parent):
        self.parent = parent
        self.previous_bytes = b''
        self.element_index = 0
        self.counter = 0
        self.current_element = None
        self.started = False

        # Serves for debugging purposes
        # self.current_byte = 0

    def parse_bytes(self, bytes, byte_counter):

        if not self.started:
            # self.current_byte = byte_counter
            self.started = True

        if self.current_element is None:
            self.current_element = self.parent.elements[0]

        bytes = self.previous_bytes + bytes
        current_byte_index = 0

        while True:
            property_values = []

            beginning_byte_index = current_byte_index

            for property in self.current_element.properties:

                size = ply_type_size(property[1])

                if current_byte_index + size[0] > len(bytes):
                    self.previous_bytes = bytes[beginning_byte_index:]
                    # self.current_byte -= len(self.previous_bytes)
                    return

                if len(size) == 1:

                    size = size[0]

                    current_property_bytes = bytes[current_byte_index:current_byte_index+size]
                    property_values.append(bytes_to_element(property[1], current_property_bytes))
                    current_byte_index += size
                    # self.current_byte += size

                elif len(size) == 2:

                    types = property[1].split()[1:]
                    current_property_bytes = bytes[current_byte_index:current_byte_index+size[0]]
                    number_of_elements = bytes_to_element(types[0], current_property_bytes)
                    current_byte_index += size[0]
                    # self.current_byte += size[0]

                    property_values.append([])

                    # Parse list
                    for i in range(number_of_elements):

                        if current_byte_index + size[1] > len(bytes):

                            self.previous_bytes = bytes[beginning_byte_index:]
                            # self.current_byte -= len(self.previous_bytes)
                            return

                        current_property_bytes = bytes[current_byte_index:current_byte_index+size[1]]
                        property_values[-1].append(bytes_to_element(types[1], current_property_bytes))
                        current_byte_index += size[1]
                        # self.current_byte += size[1]


                else:
                    print('I have not idea what this means', file=sys.stderr)

            # Add element
            if self.current_element.name == 'vertex':

                vertex = Vertex()
                red = None
                green = None
                blue = None
                alpha = None
                offset = 0

                for property in self.current_element.properties:

                    if property[0] == 'x':
                        vertex.x = property_values[offset]
                    elif property[0] == 'y':
                        vertex.y = property_values[offset]
                    elif property[0] == 'z':
                        vertex.z = property_values[offset]
                    elif property[0] == 'red':
                        red = property_values[offset] / 255
                    elif property[0] == 'green':
                        green = property_values[offset] / 255
                    elif property[0] == 'blue':
                        blue = property_values[offset] / 255
                    elif property[0] == 'alpha':
                        alpha = property_values[offset] / 255

                    offset += 1

                self.parent.add_vertex(vertex)

                if red is not None:
                    self.parent.add_color(Color(red, blue, green))

            elif self.current_element.name == 'face':

                vertex_indices = []
                tex_coords = []
                material = None

                for (i, property) in enumerate(self.current_element.properties):

                    if property[0] == 'vertex_indices':
                        vertex_indices.append(property_values[i][0])
                        vertex_indices.append(property_values[i][1])
                        vertex_indices.append(property_values[i][2])

                    elif property[0] == 'texcoord':
                        # Create texture coords
                        for j in range(0,6,2):
                            tex_coord = TexCoord(*property_values[i][j:j+2])
                            tex_coords.append(tex_coord)

                    elif property[0] == 'texnumber':
                        material = self.parent.materials[property_values[i]]

                for tex_coord in tex_coords:
                    self.parent.add_tex_coord(tex_coord)

                face = Face(*list(map(lambda x: FaceVertex(x), vertex_indices)))

                counter = 3
                if len(tex_coords) > 0:
                    for face_vertex in [face.a, face.b, face.c]:
                        face_vertex.tex_coord = len(self.parent.tex_coords) - counter
                        counter -= 1

                if material is None and len(self.parent.materials) == 1:
                    material = self.parent.materials[0]

                face.material = material

                self.parent.add_face(face)

            self.counter += 1

            if self.counter == self.current_element.number:
                self.next_element()

    def next_element(self):
        self.counter = 0
        self.element_index += 1
        if self.element_index < len(self.parent.elements):
            self.current_element = self.parent.elements[self.element_index]



class PLYBigEndianContentParser(PLYLittleEndianContentParser):
    def __init__(self, parent):
        super().__init__(self, parent)

    def parse_bytes(self, bytes):
        # Reverse bytes, and then
        super().parse_bytes(self, bytes)

class PLYExporter(Exporter):
    def __init__(self, model):
        super().__init__(model)

    def __str__(self):

        faces = sum([part.faces for part in self.model.parts], [])

        # Header
        string = "ply\nformat ascii 1.0\ncomment Automatically gnerated by model-converter\n"

        # Types : vertices
        string += "element vertex " + str(len(self.model.vertices)) +"\n"
        string += "property float x\nproperty float y\nproperty float z\n"

        # Types : faces
        string += "element face " + str(len(faces)) + "\n"
        string += "property list uint8 int32 vertex_indices\n"

        # End header
        string += "end_header\n"

        # Content of the model
        for vertex in self.model.vertices:
            string += str(vertex.x) + " " + str(vertex.y) + " " + str(vertex.z) + "\n"

        for face in faces:
            string += "3 " + str(face.a.vertex) + " " + str(face.b.vertex) + " " + str(face.c.vertex) + "\n"

        return string