You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
274 lines
9.6 KiB
274 lines
9.6 KiB
import struct
|
|
import numpy as np
|
|
import collections
|
|
|
|
|
|
def _write_ply_point(fp, x, y, z, color=None, normal=None, binary=False):
|
|
args = [x, y, z]
|
|
if color is not None:
|
|
args += [int(color[0]), int(color[1]), int(color[2])]
|
|
if normal is not None:
|
|
args += [normal[0], normal[1], normal[2]]
|
|
if binary:
|
|
fmt = '<fff'
|
|
if color is not None:
|
|
fmt = fmt + 'BBB'
|
|
if normal is not None:
|
|
fmt = fmt + 'fff'
|
|
fp.write(struct.pack(fmt, *args))
|
|
else:
|
|
fmt = '%f %f %f'
|
|
if color is not None:
|
|
fmt = fmt + ' %d %d %d'
|
|
if normal is not None:
|
|
fmt = fmt + ' %f %f %f'
|
|
fmt += '\n'
|
|
fp.write(fmt % tuple(args))
|
|
|
|
|
|
def _write_ply_triangle(fp, i0, i1, i2, binary):
|
|
if binary:
|
|
fp.write(struct.pack('<Biii', 3, i0, i1, i2))
|
|
else:
|
|
fp.write('3 %d %d %d\n' % (i0, i1, i2))
|
|
|
|
|
|
def _write_ply_header_line(fp, str, binary):
|
|
if binary:
|
|
fp.write(str.encode())
|
|
else:
|
|
fp.write(str)
|
|
|
|
|
|
def write_ply(path, verts, trias=None, color=None, normals=None, binary=False):
|
|
if verts.shape[1] != 3:
|
|
raise Exception('verts has to be of shape Nx3')
|
|
if trias is not None and trias.shape[1] != 3:
|
|
raise Exception('trias has to be of shape Nx3')
|
|
if color is not None and not callable(color) and not isinstance(color, np.ndarray) and color.shape[1] != 3:
|
|
raise Exception('color has to be of shape Nx3 or a callable')
|
|
|
|
mode = 'wb' if binary else 'w'
|
|
with open(path, mode) as fp:
|
|
_write_ply_header_line(fp, "ply\n", binary)
|
|
if binary:
|
|
_write_ply_header_line(fp, "format binary_little_endian 1.0\n", binary)
|
|
else:
|
|
_write_ply_header_line(fp, "format ascii 1.0\n", binary)
|
|
_write_ply_header_line(fp, "element vertex %d\n" % (verts.shape[0]), binary)
|
|
_write_ply_header_line(fp, "property float32 x\n", binary)
|
|
_write_ply_header_line(fp, "property float32 y\n", binary)
|
|
_write_ply_header_line(fp, "property float32 z\n", binary)
|
|
if color is not None:
|
|
_write_ply_header_line(fp, "property uchar red\n", binary)
|
|
_write_ply_header_line(fp, "property uchar green\n", binary)
|
|
_write_ply_header_line(fp, "property uchar blue\n", binary)
|
|
if normals is not None:
|
|
_write_ply_header_line(fp, "property float32 nx\n", binary)
|
|
_write_ply_header_line(fp, "property float32 ny\n", binary)
|
|
_write_ply_header_line(fp, "property float32 nz\n", binary)
|
|
if trias is not None:
|
|
_write_ply_header_line(fp, "element face %d\n" % (trias.shape[0]), binary)
|
|
_write_ply_header_line(fp, "property list uchar int32 vertex_indices\n", binary)
|
|
_write_ply_header_line(fp, "end_header\n", binary)
|
|
|
|
for vidx, v in enumerate(verts):
|
|
if color is not None:
|
|
if callable(color):
|
|
c = color(vidx)
|
|
elif color.shape[0] > 1:
|
|
c = color[vidx]
|
|
else:
|
|
c = color[0]
|
|
else:
|
|
c = None
|
|
if normals is None:
|
|
n = None
|
|
else:
|
|
n = normals[vidx]
|
|
_write_ply_point(fp, v[0], v[1], v[2], c, n, binary)
|
|
|
|
if trias is not None:
|
|
for t in trias:
|
|
_write_ply_triangle(fp, t[0], t[1], t[2], binary)
|
|
|
|
|
|
def faces_to_triangles(faces):
|
|
new_faces = []
|
|
for f in faces:
|
|
if f[0] == 3:
|
|
new_faces.append([f[1], f[2], f[3]])
|
|
elif f[0] == 4:
|
|
new_faces.append([f[1], f[2], f[3]])
|
|
new_faces.append([f[3], f[4], f[1]])
|
|
else:
|
|
raise Exception('unknown face count %d', f[0])
|
|
return new_faces
|
|
|
|
|
|
def read_ply(path):
|
|
with open(path, 'rb') as f:
|
|
# parse header
|
|
line = f.readline().decode().strip()
|
|
if line != 'ply':
|
|
raise Exception('Header error')
|
|
n_verts = 0
|
|
n_faces = 0
|
|
vert_types = {}
|
|
vert_bin_format = []
|
|
vert_bin_len = 0
|
|
vert_bin_cols = 0
|
|
line = f.readline().decode()
|
|
parse_vertex_prop = False
|
|
while line.strip() != 'end_header':
|
|
if 'format' in line:
|
|
if 'ascii' in line:
|
|
binary = False
|
|
elif 'binary_little_endian' in line:
|
|
binary = True
|
|
else:
|
|
raise Exception('invalid ply format')
|
|
if 'element face' in line:
|
|
splits = line.strip().split(' ')
|
|
n_faces = int(splits[-1])
|
|
parse_vertex_prop = False
|
|
if 'element camera' in line:
|
|
parse_vertex_prop = False
|
|
if 'element vertex' in line:
|
|
splits = line.strip().split(' ')
|
|
n_verts = int(splits[-1])
|
|
parse_vertex_prop = True
|
|
if parse_vertex_prop and 'property' in line:
|
|
prop = line.strip().split()
|
|
if prop[1] == 'float':
|
|
vert_bin_format.append('f4')
|
|
vert_bin_len += 4
|
|
vert_bin_cols += 1
|
|
elif prop[1] == 'uchar':
|
|
vert_bin_format.append('B')
|
|
vert_bin_len += 1
|
|
vert_bin_cols += 1
|
|
else:
|
|
raise Exception('invalid property')
|
|
vert_types[prop[2]] = len(vert_types)
|
|
line = f.readline().decode()
|
|
|
|
# parse content
|
|
if binary:
|
|
sz = n_verts * vert_bin_len
|
|
fmt = ','.join(vert_bin_format)
|
|
verts = np.ndarray(shape=(1, n_verts), dtype=np.dtype(fmt), buffer=f.read(sz))
|
|
verts = verts[0].astype(vert_bin_cols * 'f4,').view(dtype='f4').reshape((n_verts, -1))
|
|
faces = []
|
|
for idx in range(n_faces):
|
|
fmt = '<Biii'
|
|
length = struct.calcsize(fmt)
|
|
dat = f.read(length)
|
|
vals = struct.unpack(fmt, dat)
|
|
faces.append(vals)
|
|
faces = faces_to_triangles(faces)
|
|
faces = np.array(faces, dtype=np.int32)
|
|
else:
|
|
verts = []
|
|
for idx in range(n_verts):
|
|
vals = [float(v) for v in f.readline().decode().strip().split(' ')]
|
|
verts.append(vals)
|
|
verts = np.array(verts, dtype=np.float32)
|
|
faces = []
|
|
for idx in range(n_faces):
|
|
splits = f.readline().decode().strip().split(' ')
|
|
n_face_verts = int(splits[0])
|
|
vals = [int(v) for v in splits[0:n_face_verts + 1]]
|
|
faces.append(vals)
|
|
faces = faces_to_triangles(faces)
|
|
faces = np.array(faces, dtype=np.int32)
|
|
|
|
xyz = None
|
|
if 'x' in vert_types and 'y' in vert_types and 'z' in vert_types:
|
|
xyz = verts[:, [vert_types['x'], vert_types['y'], vert_types['z']]]
|
|
colors = None
|
|
if 'red' in vert_types and 'green' in vert_types and 'blue' in vert_types:
|
|
colors = verts[:, [vert_types['red'], vert_types['green'], vert_types['blue']]]
|
|
colors /= 255
|
|
normals = None
|
|
if 'nx' in vert_types and 'ny' in vert_types and 'nz' in vert_types:
|
|
normals = verts[:, [vert_types['nx'], vert_types['ny'], vert_types['nz']]]
|
|
|
|
return xyz, faces, colors, normals
|
|
|
|
|
|
def _read_obj_split_f(s):
|
|
parts = s.split('/')
|
|
vidx = int(parts[0]) - 1
|
|
if len(parts) >= 2 and len(parts[1]) > 0:
|
|
tidx = int(parts[1]) - 1
|
|
else:
|
|
tidx = -1
|
|
if len(parts) >= 3 and len(parts[2]) > 0:
|
|
nidx = int(parts[2]) - 1
|
|
else:
|
|
nidx = -1
|
|
return vidx, tidx, nidx
|
|
|
|
|
|
def read_obj(path):
|
|
with open(path, 'r') as fp:
|
|
lines = fp.readlines()
|
|
|
|
verts = []
|
|
colors = []
|
|
fnorms = []
|
|
fnorm_map = collections.defaultdict(list)
|
|
faces = []
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line.startswith('#') or len(line) == 0:
|
|
continue
|
|
|
|
parts = line.split()
|
|
if line.startswith('v '):
|
|
parts = parts[1:]
|
|
x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
|
|
if len(parts) == 4 or len(parts) == 7:
|
|
w = float(parts[3])
|
|
x, y, z = x / w, y / w, z / w
|
|
verts.append((x, y, z))
|
|
if len(parts) >= 6:
|
|
r, g, b = float(parts[-3]), float(parts[-2]), float(parts[-1])
|
|
rgb.append((r, g, b))
|
|
|
|
elif line.startswith('vn '):
|
|
parts = parts[1:]
|
|
x, y, z = float(parts[0]), float(parts[1]), float(parts[2])
|
|
fnorms.append((x, y, z))
|
|
|
|
elif line.startswith('f '):
|
|
parts = parts[1:]
|
|
if len(parts) != 3:
|
|
raise Exception('only triangle meshes supported atm')
|
|
vidx0, tidx0, nidx0 = _read_obj_split_f(parts[0])
|
|
vidx1, tidx1, nidx1 = _read_obj_split_f(parts[1])
|
|
vidx2, tidx2, nidx2 = _read_obj_split_f(parts[2])
|
|
|
|
faces.append((vidx0, vidx1, vidx2))
|
|
if nidx0 >= 0:
|
|
fnorm_map[vidx0].append(nidx0)
|
|
if nidx1 >= 0:
|
|
fnorm_map[vidx1].append(nidx1)
|
|
if nidx2 >= 0:
|
|
fnorm_map[vidx2].append(nidx2)
|
|
|
|
verts = np.array(verts)
|
|
colors = np.array(colors)
|
|
fnorms = np.array(fnorms)
|
|
faces = np.array(faces)
|
|
|
|
# face normals to vertex normals
|
|
norms = np.zeros_like(verts)
|
|
for vidx in fnorm_map.keys():
|
|
ind = fnorm_map[vidx]
|
|
norms[vidx] = fnorms[ind].sum(axis=0)
|
|
N = np.linalg.norm(norms, axis=1, keepdims=True)
|
|
np.divide(norms, N, out=norms, where=N != 0)
|
|
|
|
return verts, faces, colors, norms
|
|
|