220 lines
7.4 KiB
Python
220 lines
7.4 KiB
Python
|
import sys
|
||
|
|
from typing import List
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
|
||
|
|
sys_byteorder = ('>', '<')[sys.byteorder == 'little']
|
||
|
|
|
||
|
|
ply_dtypes = {
|
||
|
|
'int8': 'i1', b'char': 'i1', b'uint8': 'u1', b'uchar': 'b1', b'uchar': 'u1',
|
||
|
|
'int16': 'i2', b'short': 'i2', b'uint16': 'u2', b'ushort': 'u2',
|
||
|
|
'int32': 'i4', b'int': 'i4', b'uint32': 'u4', b'uint': 'u4',
|
||
|
|
'float32': 'f4', b'float': 'f4', b'float64': 'f8', b'double': 'f8'
|
||
|
|
}
|
||
|
|
|
||
|
|
valid_formats = {'ascii': '', 'binary_big_endian': '>', 'binary_little_endian': '<'}
|
||
|
|
|
||
|
|
|
||
|
|
def describe_element(name: str, df: pd.DataFrame) -> List[str]:
|
||
|
|
"""
|
||
|
|
Takes the columns of the dataframe and builds a ply-like description
|
||
|
|
:param name:
|
||
|
|
:param df:
|
||
|
|
"""
|
||
|
|
property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'}
|
||
|
|
element = ['element ' + name + ' ' + str(len(df))]
|
||
|
|
|
||
|
|
if name == 'face':
|
||
|
|
element.append("property list uchar int vertex_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
|
||
|
|
|
||
|
|
|
||
|
|
def read_ply(filename):
|
||
|
|
"""
|
||
|
|
Read a .ply (binary or ascii) file and store the elements in pandas DataFrame
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
filename: str
|
||
|
|
Path to the filename
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
data: dict
|
||
|
|
Elements as pandas DataFrames; comments and ob_info as list of string
|
||
|
|
"""
|
||
|
|
if not type(filename) is str:
|
||
|
|
filename = str(filename)
|
||
|
|
with open(filename, 'rb') as ply:
|
||
|
|
if b'ply' not in ply.readline():
|
||
|
|
raise ValueError('The file does not start whith the word ply')
|
||
|
|
# get binary_little/big or ascii
|
||
|
|
fmt = ply.readline().split()[1].decode()
|
||
|
|
# get extension for building the numpy dtypes
|
||
|
|
ext = valid_formats[fmt]
|
||
|
|
|
||
|
|
line = []
|
||
|
|
dtypes = defaultdict(list)
|
||
|
|
count = 2
|
||
|
|
points_size = None
|
||
|
|
mesh_size = None
|
||
|
|
has_texture = False
|
||
|
|
while b'end_header' not in line and line != b'':
|
||
|
|
line = ply.readline()
|
||
|
|
|
||
|
|
if b'element' in line:
|
||
|
|
line = line.split()
|
||
|
|
name = line[1].decode()
|
||
|
|
size = int(line[2])
|
||
|
|
if name == "vertex":
|
||
|
|
points_size = size
|
||
|
|
elif name == "face":
|
||
|
|
mesh_size = size
|
||
|
|
|
||
|
|
elif b'property' in line:
|
||
|
|
line = line.split()
|
||
|
|
# element faces
|
||
|
|
if b'list' in line:
|
||
|
|
|
||
|
|
if b"vertex_indices" in line[-1] or b"vertex_index" in line[-1]:
|
||
|
|
mesh_names = ["n_points", "v1", "v2", "v3"]
|
||
|
|
else:
|
||
|
|
has_texture = True
|
||
|
|
mesh_names = ["n_coords"] + ["v1_u", "v1_v", "v2_u", "v2_v", "v3_u", "v3_v"]
|
||
|
|
|
||
|
|
if fmt == "ascii":
|
||
|
|
# the first number has different dtype than the list
|
||
|
|
dtypes[name].append(
|
||
|
|
(mesh_names[0], ply_dtypes[line[2]]))
|
||
|
|
# rest of the numbers have the same dtype
|
||
|
|
dt = ply_dtypes[line[3]]
|
||
|
|
else:
|
||
|
|
# the first number has different dtype than the list
|
||
|
|
dtypes[name].append(
|
||
|
|
(mesh_names[0], ext + ply_dtypes[line[2]]))
|
||
|
|
# rest of the numbers have the same dtype
|
||
|
|
dt = ext + ply_dtypes[line[3]]
|
||
|
|
|
||
|
|
for j in range(1, len(mesh_names)):
|
||
|
|
dtypes[name].append((mesh_names[j], dt))
|
||
|
|
else:
|
||
|
|
if fmt == "ascii":
|
||
|
|
dtypes[name].append((line[2].decode(), ply_dtypes[line[1]]))
|
||
|
|
else:
|
||
|
|
dtypes[name].append((line[2].decode(), ext + ply_dtypes[line[1]]))
|
||
|
|
count += 1
|
||
|
|
|
||
|
|
# for bin
|
||
|
|
end_header = ply.tell()
|
||
|
|
|
||
|
|
data = {}
|
||
|
|
|
||
|
|
if fmt == 'ascii':
|
||
|
|
top = count
|
||
|
|
bottom = 0 if mesh_size is None else mesh_size
|
||
|
|
|
||
|
|
names = [x[0] for x in dtypes["vertex"]]
|
||
|
|
|
||
|
|
data["points"] = pd.read_csv(filename, sep=" ", header=None, engine="python", skiprows=top,
|
||
|
|
skipfooter=bottom, usecols=names, names=names)
|
||
|
|
|
||
|
|
for n, col in enumerate(data["points"].columns):
|
||
|
|
data["points"][col] = data["points"][col].astype(dtypes["vertex"][n][1])
|
||
|
|
|
||
|
|
if mesh_size:
|
||
|
|
top = count + points_size
|
||
|
|
|
||
|
|
names = np.array([x[0] for x in dtypes["face"]])
|
||
|
|
usecols = [1, 2, 3, 5, 6, 7, 8, 9, 10] if has_texture else [1, 2, 3]
|
||
|
|
names = names[usecols]
|
||
|
|
|
||
|
|
data["faces"] = pd.read_csv(filename, sep=" ", header=None, engine="python",
|
||
|
|
skiprows=top, usecols=usecols, names=names)
|
||
|
|
|
||
|
|
for n, col in enumerate(data["faces"].columns):
|
||
|
|
data["faces"][col] = data["faces"][col].astype(dtypes["face"][n + 1][1])
|
||
|
|
|
||
|
|
else:
|
||
|
|
with open(filename, 'rb') as ply:
|
||
|
|
ply.seek(end_header)
|
||
|
|
points_np = np.fromfile(ply, dtype=dtypes["vertex"], count=points_size)
|
||
|
|
if ext != sys_byteorder:
|
||
|
|
points_np = points_np.byteswap().newbyteorder()
|
||
|
|
data["points"] = pd.DataFrame(points_np)
|
||
|
|
if mesh_size:
|
||
|
|
mesh_np = np.fromfile(ply, dtype=dtypes["face"], count=mesh_size)
|
||
|
|
if ext != sys_byteorder:
|
||
|
|
mesh_np = mesh_np.byteswap().newbyteorder()
|
||
|
|
data["faces"] = pd.DataFrame(mesh_np)
|
||
|
|
data["faces"].drop('n_points', axis=1, inplace=True)
|
||
|
|
|
||
|
|
return data
|
||
|
|
|
||
|
|
|
||
|
|
def write_ply(filename, points=None, faces=None, as_text=True):
|
||
|
|
"""
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
filename: str
|
||
|
|
The created file will be named with this
|
||
|
|
points: ndarray
|
||
|
|
faces: ndarray
|
||
|
|
as_text: boolean
|
||
|
|
Set the write mode of the file. Default: binary
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
boolean
|
||
|
|
True if no problems
|
||
|
|
|
||
|
|
"""
|
||
|
|
filename = str(filename)
|
||
|
|
if not filename.endswith('ply'):
|
||
|
|
filename += '.ply'
|
||
|
|
|
||
|
|
# open in text mode to write the header
|
||
|
|
with open(filename, 'w') as ply:
|
||
|
|
header = ['ply']
|
||
|
|
|
||
|
|
if as_text:
|
||
|
|
header.append('format ascii 1.0')
|
||
|
|
else:
|
||
|
|
header.append('format binary_' + sys.byteorder + '_endian 1.0')
|
||
|
|
|
||
|
|
if points is not None:
|
||
|
|
points = pd.DataFrame(points, columns=['x', 'y', 'z'])
|
||
|
|
header.extend(describe_element('vertex', points))
|
||
|
|
if faces is not None:
|
||
|
|
faces = pd.DataFrame(faces.copy())
|
||
|
|
faces.insert(loc=0, column="n_points", value=3)
|
||
|
|
faces["n_points"] = faces["n_points"].astype("u1")
|
||
|
|
header.extend(describe_element('face', faces))
|
||
|
|
|
||
|
|
header.append('end_header')
|
||
|
|
|
||
|
|
for line in header:
|
||
|
|
ply.write("%s\n" % line)
|
||
|
|
|
||
|
|
if as_text:
|
||
|
|
if points is not None:
|
||
|
|
points.to_csv(filename, sep=" ", index=False, header=False, mode='a', encoding='ascii')
|
||
|
|
if faces is not None:
|
||
|
|
faces.to_csv(filename, sep=" ", index=False, header=False, mode='a', encoding='ascii')
|
||
|
|
else:
|
||
|
|
with open(filename, 'ab') as ply:
|
||
|
|
if points is not None:
|
||
|
|
points.to_records(index=False).tofile(ply)
|
||
|
|
if faces is not None:
|
||
|
|
faces.to_records(index=False).tofile(ply)
|
||
|
|
|
||
|
|
return True
|