import os

import numpy as np

from plyfile import PlyData, PlyElement

import pandas as pd

from scannet200_constants import *

def read_plymesh(filepath):

"""Read ply file and return it as numpy array. Returns None if emtpy."""

with open(filepath, 'rb') as f:

plydata = PlyData.read(f)

if plydata.elements:

vertices = pd.DataFrame(plydata['vertex'].data).values

faces = np.array([f[0] for f in plydata["face"].data])

return vertices, faces

def save_plymesh(vertices, faces, filename, verbose=True, with_label=True):

"""Save an RGB point cloud as a PLY file.

Args:

points_3d: Nx6 matrix where points_3d[:, :3] are the XYZ coordinates and points_3d[:, 4:] are

the RGB values. If Nx3 matrix, save all points with [128, 128, 128] (gray) color.

"""

assert vertices.ndim == 2

if with_label:

if vertices.shape[1] == 7:

python_types = (float, float, float, int, int, int, int)

npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),

('blue', 'u1'), ('label', 'u4')]

if vertices.shape[1] == 8:

python_types = (float, float, float, int, int, int, int, int)

npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),

('blue', 'u1'), ('label', 'u4'), ('instance_id', 'u4')]

else:

if vertices.shape[1] == 3:

gray_concat = np.tile(np.array([128], dtype=np.uint8), (vertices.shape[0], 3))

vertices = np.hstack((vertices, gray_concat))

elif vertices.shape[1] == 6:

python_types = (float, float, float, int, int, int)

npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'),

('blue', 'u1')]

else:

pass

vertices_list = []

for row_idx in range(vertices.shape[0]):

cur_point = vertices[row_idx]

vertices_list.append(tuple(dtype(point) for dtype, point in zip(python_types, cur_point)))

vertices_array = np.array(vertices_list, dtype=npy_types)

elements = [PlyElement.describe(vertices_array, 'vertex')]

if faces is not None:

faces_array = np.empty(len(faces), dtype=[('vertex_indices', 'i4', (3,))])

faces_array['vertex_indices'] = faces

elements += [PlyElement.describe(faces_array, 'face')]

# Write

PlyData(elements).write(filename)

if verbose is True:

print('Saved point cloud to: %s' % filename)

# Map the raw category id to the point cloud

def point_indices_from_group(points, seg_indices, group, labels_pd, CLASS_IDs):

group_segments = np.array(group['segments'])

label = group['label']

# Map the category name to id

label_ids = labels_pd[labels_pd['raw_category'] == label]['id']

label_id = int(label_ids.iloc[0]) if len(label_ids) > 0 else 0

# Only store for the valid categories

if not label_id in CLASS_IDs:

label_id = 0

# get points, where segment indices (points labelled with segment ids) are in the group segment list

point_IDs = np.where(np.isin(seg_indices, group_segments))

return points[point_IDs], point_IDs[0], label_id

# Uncomment out if mesh voxelization is required

# import trimesh

# from trimesh.voxel import creation

# from sklearn.neighbors import KDTree

# import MinkowskiEngine as ME

# VOXELIZE the scene from sampling on the mesh directly instead of vertices

def voxelize_pointcloud(points, colors, labels, instances, faces, voxel_size=0.2):

# voxelize mesh first and determine closest labels with KDTree search

trimesh_scene_mesh = trimesh.Trimesh(vertices=points, faces=faces)

voxel_grid = creation.voxelize(trimesh_scene_mesh, voxel_size)

voxel_cloud = np.asarray(voxel_grid.points)

orig_tree = KDTree(points, leaf_size=8)

_, voxel_pc_matches = orig_tree.query(voxel_cloud, k=1)

voxel_pc_matches = voxel_pc_matches.flatten()

# match colors to voxel ids

points = points[voxel_pc_matches] / voxel_size

colors = colors[voxel_pc_matches]

labels = labels[voxel_pc_matches]

instances = instances[voxel_pc_matches]

# Voxelize scene

quantized_scene, scene_inds = ME.utils.sparse_quantize(points, return_index=True)

quantized_scene_colors = colors[scene_inds]

quantized_labels = labels[scene_inds]

quantized_instances = instances[scene_inds]

return quantized_scene, quantized_scene_colors, quantized_labels, quantized_instances