add .ply

lib/ply.py

@@ -0,0 +1,351 @@
+#
+#
+#      0===============================0
+#      |    PLY files reader/writer    |
+#      0===============================0
+#
+#
+# ----------------------------------------------------------------------------------------------------------------------
+#
+#      function to read/write .ply files
+#
+# ----------------------------------------------------------------------------------------------------------------------
+#
+#      Hugues THOMAS - 10/02/2017
+#
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#
+#          Imports and global variables
+#      \**********************************/
+#
+
+# Basic libs
+import numpy as np
+import sys
+
+# Define PLY types
+ply_dtypes = dict([
+    (b'int8', 'i1'),
+    (b'char', 'i1'),
+    (b'uint8', 'u1'),
+    (b'uchar', 'u1'),
+    (b'int16', 'i2'),
+    (b'short', 'i2'),
+    (b'uint16', 'u2'),
+    (b'ushort', 'u2'),
+    (b'int32', 'i4'),
+    (b'int', 'i4'),
+    (b'uint32', 'u4'),
+    (b'uint', 'u4'),
+    (b'float32', 'f4'),
+    (b'float', 'f4'),
+    (b'float64', 'f8'),
+    (b'double', 'f8')
+])
+
+# Numpy reader format
+valid_formats = {'ascii': '', 'binary_big_endian': '>',
+                 'binary_little_endian': '<'}
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#
+#           Functions
+#       \***************/
+#
+
+
+def parse_header(plyfile, ext):
+    # Variables
+    line = []
+    properties = []
+    num_points = None
+
+    while b'end_header' not in line and line != b'':
+        line = plyfile.readline()
+
+        if b'element' in line:
+            line = line.split()
+            num_points = int(line[2])
+
+        elif b'property' in line:
+            line = line.split()
+            properties.append((line[2].decode(), ext + ply_dtypes[line[1]]))
+
+    return num_points, properties
+
+
+def parse_mesh_header(plyfile, ext):
+    # Variables
+    line = []
+    vertex_properties = []
+    num_points = None
+    num_faces = None
+    current_element = None
+
+    while b'end_header' not in line and line != b'':
+        line = plyfile.readline()
+
+        # Find point element
+        if b'element vertex' in line:
+            current_element = 'vertex'
+            line = line.split()
+            num_points = int(line[2])
+
+        elif b'element face' in line:
+            current_element = 'face'
+            line = line.split()
+            num_faces = int(line[2])
+
+        elif b'property' in line:
+            if current_element == 'vertex':
+                line = line.split()
+                vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]]))
+            elif current_element == 'vertex':
+                if not line.startswith('property list uchar int'):
+                    raise ValueError('Unsupported faces property : ' + line)
+
+    return num_points, num_faces, vertex_properties
+
+
+def read_ply(filename, triangular_mesh=False):
+    """
+    Read ".ply" files
+
+    Parameters
+    ----------
+    filename : string
+        the name of the file to read.
+
+    Returns
+    -------
+    result : array
+        data stored in the file
+
+    Examples
+    --------
+    Store data in file
+
+    >>> points = np.random.rand(5, 3)
+    >>> values = np.random.randint(2, size=10)
+    >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values'])
+
+    Read the file
+
+    >>> data = read_ply('example.ply')
+    >>> values = data['values']
+    array([0, 0, 1, 1, 0])
+
+    >>> points = np.vstack((data['x'], data['y'], data['z'])).T
+    array([[ 0.466  0.595  0.324]
+           [ 0.538  0.407  0.654]
+           [ 0.850  0.018  0.988]
+           [ 0.395  0.394  0.363]
+           [ 0.873  0.996  0.092]])
+
+    """
+
+    with open(filename, 'rb') as plyfile:
+
+        # Check if the file start with ply
+        if b'ply' not in plyfile.readline():
+            raise ValueError('The file does not start whith the word ply')
+
+        # get binary_little/big or ascii
+        fmt = plyfile.readline().split()[1].decode()
+        if fmt == "ascii":
+            raise ValueError('The file is not binary')
+
+        # get extension for building the numpy dtypes
+        ext = valid_formats[fmt]
+
+        # PointCloud reader vs mesh reader
+        if triangular_mesh:
+
+            # Parse header
+            num_points, num_faces, properties = parse_mesh_header(plyfile, ext)
+
+            # Get point data
+            vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points)
+
+            # Get face data
+            face_properties = [('k', ext + 'u1'),
+                               ('v1', ext + 'i4'),
+                               ('v2', ext + 'i4'),
+                               ('v3', ext + 'i4')]
+            faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces)
+
+            # Return vertex data and concatenated faces
+            faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T
+            data = [vertex_data, faces]
+
+        else:
+
+            # Parse header
+            num_points, properties = parse_header(plyfile, ext)
+
+            # Get data
+            data = np.fromfile(plyfile, dtype=properties, count=num_points)
+
+    return data
+
+
+def header_properties(field_list, field_names):
+    # List of lines to write
+    lines = []
+
+    # First line describing element vertex
+    lines.append('element vertex %d' % field_list[0].shape[0])
+
+    # Properties lines
+    i = 0
+    for fields in field_list:
+        for field in fields.T:
+            lines.append('property %s %s' % (field.dtype.name, field_names[i]))
+            i += 1
+
+    return lines
+
+
+def write_ply(filename, field_list, field_names, triangular_faces=None):
+    """
+    Write ".ply" files
+
+    Parameters
+    ----------
+    filename : string
+        the name of the file to which the data is saved. A '.ply' extension will be appended to the
+        file name if it does no already have one.
+
+    field_list : list, tuple, numpy array
+        the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a
+        tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered
+        as one field.
+
+    field_names : list
+        the name of each fields as a list of strings. Has to be the same length as the number of
+        fields.
+
+    Examples
+    --------
+    >>> points = np.random.rand(10, 3)
+    >>> write_ply('example1.ply', points, ['x', 'y', 'z'])
+
+    >>> values = np.random.randint(2, size=10)
+    >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values'])
+
+    >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8)
+    >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values']
+    >>> write_ply('example3.ply', [points, colors, values], field_names)
+
+    """
+
+    # Format list input to the right form
+    field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,))
+    for i, field in enumerate(field_list):
+        if field.ndim < 2:
+            field_list[i] = field.reshape(-1, 1)
+        if field.ndim > 2:
+            print('fields have more than 2 dimensions')
+            return False
+
+            # check all fields have the same number of data
+    n_points = [field.shape[0] for field in field_list]
+    if not np.all(np.equal(n_points, n_points[0])):
+        print('wrong field dimensions')
+        return False
+
+        # Check if field_names and field_list have same nb of column
+    n_fields = np.sum([field.shape[1] for field in field_list])
+    if (n_fields != len(field_names)):
+        print('wrong number of field names')
+        return False
+
+    # Add extension if not there
+    if not filename.endswith('.ply'):
+        filename += '.ply'
+
+    # open in text mode to write the header
+    with open(filename, 'w') as plyfile:
+
+        # First magical word
+        header = ['ply']
+
+        # Encoding format
+        header.append('format binary_' + sys.byteorder + '_endian 1.0')
+
+        # Points properties description
+        header.extend(header_properties(field_list, field_names))
+
+        # Add faces if needded
+        if triangular_faces is not None:
+            header.append('element face {:d}'.format(triangular_faces.shape[0]))
+            header.append('property list uchar int vertex_indices')
+
+        # End of header
+        header.append('end_header')
+
+        # Write all lines
+        for line in header:
+            plyfile.write("%s\n" % line)
+
+    # open in binary/append to use tofile
+    with open(filename, 'ab') as plyfile:
+
+        # Create a structured array
+        i = 0
+        type_list = []
+        for fields in field_list:
+            for field in fields.T:
+                type_list += [(field_names[i], field.dtype.str)]
+                i += 1
+        data = np.empty(field_list[0].shape[0], dtype=type_list)
+        i = 0
+        for fields in field_list:
+            for field in fields.T:
+                data[field_names[i]] = field
+                i += 1
+
+        data.tofile(plyfile)
+
+        if triangular_faces is not None:
+            triangular_faces = triangular_faces.astype(np.int32)
+            type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)]
+            data = np.empty(triangular_faces.shape[0], dtype=type_list)
+            data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8)
+            data['0'] = triangular_faces[:, 0]
+            data['1'] = triangular_faces[:, 1]
+            data['2'] = triangular_faces[:, 2]
+            data.tofile(plyfile)
+
+    return True
+
+
+def describe_element(name, df):
+    """ Takes the columns of the dataframe and builds a ply-like description
+
+    Parameters
+    ----------
+    name: str
+    df: pandas DataFrame
+
+    Returns
+    -------
+    element: list[str]
+    """
+    property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'}
+    element = ['element ' + name + ' ' + str(len(df))]
+
+    if name == 'face':
+        element.append("property list uchar int points_indices")
+
+    else:
+        for i in range(len(df.columns)):
+            # get first letter of dtype to infer format
+            f = property_formats[str(df.dtypes[i])[0]]
+            element.append('property ' + f + ' ' + df.columns.values[i])
+
+    return element
+