Writer for .ply geometry files (pointclouds, meshes, lines)

plywriter.py

import traceback

"""
A simple writer for .ply geometry files. Very useful for debugging!
Buffers the contents and writes out only on exit (necessary for writing
the header without knowing the vertex/edge/face count in advance).
Closes the fd in the case a with block raises an exception.
We could wait until the exit to even open the fd; however, I decided
against that to match expected behaviour. This could accomodate a
design where the vertex/edge/face count is provided in advance.

Usage:
```
from PLYWriter import PLYWriter
...
with PLYWriter('some_file.ply`, hasEdges=True) as f:
    if len(start_pts) != len(end_pts):
        raise ValueError('must have an end for each start!')
    for v in start_pts:
        f.addPoint(v)
    for v in end_pts:
        f.addPoint(v)
    for i in range(len(start_pts)):
        f.addEdge(i, i + len(start_pts))
```
"""

class PLYWriter:
    def __init__(
        self,
        filename,
        hasNormals=False,
        hasColours=False,
        hasEdges=False,
        hasFaces=False
    ):
        self.filename = filename
        if not filename.endswith('.ply'):
            self.filename += '.ply'
        self.hasNormals = hasNormals
        self.hasColours = hasColours
        self.hasEdges = hasEdges
        self.hasFaces = hasFaces
        self.vertices = []
        self.edges = []
        self.faces = []

    def __enter__(self):
        self.fd = open(self.filename, 'w')
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is not None:
            traceback.print_exception(exc_type, exc_value, tb)
            self.fd.close()
            return False # carry exception on
        # Now we can write out everything!
        self.__writeHeader()
        for vert in self.vertices:
            self.fd.write(vert)
            self.fd.write('\n')
        for edge in self.edges:
            self.fd.write(edge)
            self.fd.write('\n')
        for face in self.faces:
            self.fd.write(face)
            self.fd.write('\n')
        self.fd.close()
        return True

    def addPoint(self, pos, normal=None, colour=None):
        vert = ''
        if len(pos) != 3:
            raise ValueError('invalid position length')
        vert += '{:.6f} {:.6f} {:.6f}'.format(pos[0], pos[1], pos[2])
        if self.hasNormals:
            if len(normal) != 3:
                raise ValueError('invalid normal length')
            vert += ' {:.6f} {:.6f} {:.6f}'.format(normal[0], normal[1], normal[2])
        else:
            if normal != None:
                raise ValueError('unexpected normal provided!')
        if self.hasColours:
            if len(colour) != 3:
                raise ValueError('invalid colour length')
            vert += ' {:.0f} {:.0f} {:.0f}'.format(colour[0], colour[1], colour[2])
        else:
            if colour != None:
                raise ValueError('unexpected colour provided!')
        self.vertices.append(vert)
    
    def addEdge(self, v1, v2):
        if not self.hasEdges:
            print('unexpected edge provided! setting hasEdges to True')
            self.hasEdges = True
        self.edges.append('{} {}'.format(v1, v2))

    def addFace(self, v1, v2, v3, v4=None):
        if not self.hasFaces:
            print('unexpected face provided! setting hasFaces to True')
            self.hasFaces = True
        if v4 is not None:
            self.faces.append('4 {} {} {} {}'.format(v1, v2, v3, v4))
        else:
            self.faces.append('3 {} {} {}'.format(v1, v2, v3))

    def __writeHeader(self):
        self.fd.write('ply\n')
        self.fd.write('format ascii 1.0\n')
        self.fd.write('comment File generated by PLYWriter v0.6.10\n')
        self.fd.write('element vertex {}\n'.format(len(self.vertices)))
        self.fd.write('property float x\n')
        self.fd.write('property float y\n')
        self.fd.write('property float z\n')
        if self.hasNormals:
            self.fd.write('property float nx\n')
            self.fd.write('property float ny\n')
            self.fd.write('property float nz\n')
        if self.hasColours:
            self.fd.write('property uchar red\n')
            self.fd.write('property uchar green\n')
            self.fd.write('property uchar blue\n')
        if self.hasEdges:
            self.fd.write('element edge {}\n'.format(len(self.edges)))
            self.fd.write('property int vertex1\n')
            self.fd.write('property int vertex2\n')
        if self.hasFaces:
            self.fd.write('element face {}\n'.format(len(self.faces)))
            self.fd.write('property list uchar int vertex_indices\n')
        self.fd.write('end_header\n')