Spaces:
Paused
Paused
File size: 3,958 Bytes
8a32844 470a712 8a32844 470a712 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import numpy as np
import torch
from plyfile import PlyData, PlyElement
import cv2
import trimesh
def get_pcd_base(H, W, u0, v0, fx, fy):
x_row = np.arange(0, W)
x = np.tile(x_row, (H, 1))
x = x.astype(np.float32)
u_m_u0 = x - u0
y_col = np.arange(0, H) # y_col = np.arange(0, height)
y = np.tile(y_col, (W, 1)).T
y = y.astype(np.float32)
v_m_v0 = y - v0
x = u_m_u0 / fx
y = v_m_v0 / fy
z = np.ones_like(x)
pw = np.stack([x, y, z], axis=2) # [h, w, c]
return pw
def reconstruct_pcd(depth, fx, fy, u0, v0, pcd_base=None, mask=None):
if type(depth) == torch.__name__:
depth = depth.cpu().numpy().squeeze()
depth = cv2.medianBlur(depth, 5)
if pcd_base is None:
H, W = depth.shape
pcd_base = get_pcd_base(H, W, u0, v0, fx, fy)
pcd = depth[:, :, None] * pcd_base
if mask:
pcd[mask] = 0
return pcd
def save_point_cloud(pcd, rgb, filename, binary=True):
"""Save an RGB point cloud as a PLY file.
:paras
@pcd: Nx3 matrix, the XYZ coordinates
@rgb: Nx3 matrix, the rgb colors for each 3D point
"""
assert pcd.shape[0] == rgb.shape[0]
if rgb is None:
gray_concat = np.tile(np.array([128], dtype=np.uint8),
(pcd.shape[0], 3))
points_3d = np.hstack((pcd, gray_concat))
else:
points_3d = np.hstack((pcd, rgb))
python_types = (float, float, float, int, int, int)
npy_types = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'),
('green', 'u1'), ('blue', 'u1')]
if binary is True:
# Format into Numpy structured array
vertices = []
for row_idx in range(points_3d.shape[0]):
cur_point = points_3d[row_idx]
vertices.append(
tuple(
dtype(point)
for dtype, point in zip(python_types, cur_point)))
vertices_array = np.array(vertices, dtype=npy_types)
el = PlyElement.describe(vertices_array, 'vertex')
# write
PlyData([el]).write(filename)
else:
x = np.squeeze(points_3d[:, 0])
y = np.squeeze(points_3d[:, 1])
z = np.squeeze(points_3d[:, 2])
r = np.squeeze(points_3d[:, 3])
g = np.squeeze(points_3d[:, 4])
b = np.squeeze(points_3d[:, 5])
ply_head = 'ply\n' \
'format ascii 1.0\n' \
'element vertex %d\n' \
'property float x\n' \
'property float y\n' \
'property float z\n' \
'property uchar red\n' \
'property uchar green\n' \
'property uchar blue\n' \
'end_header' % r.shape[0]
# ---- Save ply data to disk
np.savetxt(filename, np.column_stack([x, y, z, r, g, b]), fmt='%f %f %f %d %d %d', header=ply_head, comments='')
def ply_to_obj(ply_file, obj_file):
mesh = trimesh.load_mesh(ply_file)
# T2 = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
# mesh.apply_transform(T2)
mesh.export(obj_file)
# import numpy as np
# def save_point_cloud_to_obj(points, colors, file_name):
# """
# Save a numpy array of point cloud data with color to an OBJ file.
# Args:
# points (np.ndarray): A numpy array of shape (H, W, 3) where H is height, W is width.
# colors (np.ndarray): A numpy array of color data, shape (H, W, 3), values should be in [0, 1].
# file_name (str): The path to the output .obj file.
# """
# H, W, _ = points.shape
# assert points.shape == colors.shape, "Points and colors must have the same shape"
# with open(file_name, 'w') as file:
# for i in range(H):
# for j in range(W):
# x, y, z = points[i, j]
# r, g, b = colors[i, j]
# file.write(f"v {x} {y} {z} {r} {g} {b}\n")
|