-
Notifications
You must be signed in to change notification settings - Fork 90
/
project_pixels_to_world_example.py
67 lines (55 loc) · 2.28 KB
/
project_pixels_to_world_example.py
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
import os
import numpy as np
from PIL import Image
import imageio
import torch
import trimesh
from im2mesh.common import (get_tensor_values,
transform_to_world, sample_patch_points
)
'''
This script is an example how pixels with corresponding depth values can be projected to the object-centric world coordinate.
Given the variable PATH, it assumes folder with images, masks, and depth maps, respectively (see below).
Then, n_sample_points are sampled from each image and, if the depth value is not inf, is projected to the object-centric world space.
It loops over the first n_images. The final point cloud is then saved to a ply file in out/pcd.ply.
'''
PATH = 'data/DTU/scan65/scan65'
img_folder = 'image'
depth_folder = 'depth'
mask_folder = 'mask'
n_sample_points = 1000
n_images = 49
img_path = os.path.join(PATH, img_folder)
depth_path = os.path.join(PATH, depth_folder)
img_list = [os.path.join(img_path, i) for i in os.listdir(img_path)]
img_list.sort()
depth_list = [os.path.join(depth_path, i) for i in os.listdir(depth_path)]
depth_list.sort()
cam = np.load(os.path.join(PATH, 'cameras.npz'))
v_out = []
rgb_out = []
for i, imgp in enumerate(img_list):
# load files
img = np.array(Image.open(imgp).convert("RGB")).astype(np.float32) / 255
h, w, c = img.shape
depth = np.array(imageio.imread(depth_list[i]))
depth = depth.reshape(depth.shape[0], depth.shape[1], -1)[..., 0]
hd, wd = depth.shape
assert(h == hd and w == wd)
p = sample_patch_points(1, n_sample_points, patch_size=1, image_resolution=(h, w), continuous=False)
img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
depth = torch.from_numpy(depth).unsqueeze(0).unsqueeze(1)
rgb = get_tensor_values(img, p)
d = get_tensor_values(depth, p)
mask = (d != np.inf).squeeze(-1)
d = d[mask].unsqueeze(0)
rgb = rgb[mask]
p = p[mask].unsqueeze(0)
# transform to world
cm = cam.get('camera_mat_%d' % i).astype(np.float32).reshape(1, 4, 4)
wm = cam.get('world_mat_%d' % i).astype(np.float32).reshape(1, 4, 4)
sm = cam.get('scale_mat_%d' % i, np.eye(4)).astype(np.float32).reshape(1, 4, 4)
p_world = transform_to_world(p, d, cm, wm, sm)[0]
v_out.append(p_world)
v = np.concatenate(v_out, axis=0)
mesh = trimesh.Trimesh(vertices=v).export('out/pcd.ply')