Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,136 Bytes
d82e7f9 |
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 |
import os
import numpy as np
import utils3d
from plyfile import PlyData, PlyElement
from PIL import Image
def sphere_uv2dirs(uv: np.ndarray):
theta, phi = (1 - uv[..., 0]) * (2 * np.pi), uv[..., 1] * np.pi
directions = np.stack([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)], axis=-1)
return directions
def save_3d_points(points: np.array, colors: np.array, mask: np.array, save_path: str):
points = points.reshape(-1, 3)
colors = colors.reshape(-1, 3)
mask = mask.reshape(-1, 1)
vertex_data = np.empty(mask.sum(), dtype=[
('x', 'f4'),
('y', 'f4'),
('z', 'f4'),
('red', 'u1'),
('green', 'u1'),
('blue', 'u1'),
])
vertex_data['x'] = [a for i, a in enumerate(points[:, 0]) if mask[i]]
vertex_data['y'] = [a for i, a in enumerate(points[:, 1]) if mask[i]]
vertex_data['z'] = [a for i, a in enumerate(points[:, 2]) if mask[i]]
vertex_data['red'] = [a for i, a in enumerate(colors[:, 0]) if mask[i]]
vertex_data['green'] = [a for i, a in enumerate(colors[:, 1]) if mask[i]]
vertex_data['blue'] = [a for i, a in enumerate(colors[:, 2]) if mask[i]]
vertex_element = PlyElement.describe(vertex_data, 'vertex', comments=['vertices with color'])
save_dir = os.path.dirname(save_path)
if not os.path.exists(save_dir): os.makedirs(save_dir, exist_ok=True)
PlyData([vertex_element], text=True).write(save_path)
def colorize_normal(normal: np.ndarray, normal_mask: np.ndarray):
normal_rgb = (((normal + 1) * 0.5) * 255).astype(np.uint8)
normal_mask = np.repeat(np.expand_dims(normal_mask, axis=-1), 3, axis=-1)
normal_mask = normal_mask.astype(np.uint8)
normal_rgb = normal_rgb * normal_mask
return normal_rgb
def normal_normalize(normal: np.ndarray):
normal_norm = np.linalg.norm(normal, axis=-1, keepdims=True)
normal_norm[normal_norm < 1e-6] = 1e-6
return normal / normal_norm
def distance2pointcloud(
distance: np.ndarray,
image: np.ndarray,
mask: np.ndarray,
save_path: str = None,
return_normal: bool = False,
save_distance: bool = False
):
if distance.ndim >= 3: distance = distance.squeeze()
if save_distance:
save_path_dis = save_path.replace('3dpc', 'depth').replace('.ply', '.npy')
save_dir_dis = os.path.dirname(save_path_dis)
if not os.path.exists(save_dir_dis): os.makedirs(save_dir_dis, exist_ok=True)
np.save(save_path_dis, distance)
height, width = distance.shape[:2]
points = distance[:, :, None] * sphere_uv2dirs(utils3d.numpy.image_uv(width=width, height=height))
save_3d_points(points, image, mask, save_path)
if return_normal:
normal, normal_mask = utils3d.numpy.points_to_normals(points, mask)
normal = normal * np.array([-1, -1, 1])
normal = normal_normalize(normal)
normal_1 = normal[..., 0]
normal_2 = normal[..., 1]
normal_3 = normal[..., 2]
normal = np.stack([normal_1, normal_3, normal_2], axis=-1)
normal_img = colorize_normal(normal, normal_mask)
return Image.fromarray(normal_img)
|