-
Notifications
You must be signed in to change notification settings - Fork 3
/
room_viewer.py
55 lines (48 loc) · 2.36 KB
/
room_viewer.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
import os
import cv2
import torch
import argparse
import numpy as np
import open3d as op
from tqdm import tqdm
def get_3Dpoints(dep,H,W):
x,y = np.meshgrid(np.arange(W),np.arange(H))
theta = (1.0-2*x/float(W))*np.pi
phi = (0.5-y/float(H))*np.pi
ct,st,cp,sp = np.cos(theta),np.sin(theta),np.cos(phi),np.sin(phi)
vec = np.array([(cp*ct),(cp*st),(sp)])
pts = vec * dep
return pts.transpose([1,2,0])
if __name__ == '__main__':
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--img_dir', default='Example/',
help='NOTE: Remeber to quote your glob path. '
'All the given images are assumed to be horizontal.')
parser.add_argument('--data', default='Results/',
help='Directory to depth data')
args = parser.parse_args()
os.makedirs(os.path.join(args.data,'point_clouds'),exist_ok=True)
# Prepare image to processed
img_names = os.listdir(args.img_dir)
img_names.sort()
#Output resolution
H,W = 512, 1024
for name in tqdm(img_names):
img = cv2.imread(os.path.join(args.img_dir,name))
img = cv2.cvtColor(cv2.resize(img,(W,H)),cv2.COLOR_BGR2RGB)
seg = cv2.imread(os.path.join(args.data,'semantic',name[:-4]+'_seg.png'))
seg = cv2.cvtColor(cv2.resize(seg,(W,H),interpolation=cv2.INTER_NEAREST),cv2.COLOR_BGR2RGB)
dep = torch.FloatTensor(np.load(os.path.join(args.data,'depth',name[:-3]+'npy'))).unsqueeze(0).unsqueeze(0)
dep = torch.nn.functional.interpolate(dep,(H,W),mode='nearest').squeeze(0)
dep = dep.numpy()
pts = get_3Dpoints(dep,H,W)
pcd = op.geometry.PointCloud()
pcd.points = op.utility.Vector3dVector(pts.reshape(-1,3))
pcd.colors = op.utility.Vector3dVector(img.reshape(-1,3)/255.)
op.visualization.draw_geometries([pcd],window_name=name)
pcd2 = op.geometry.PointCloud()
pcd2.points = op.utility.Vector3dVector(pts.reshape(-1,3))
pcd2.colors = op.utility.Vector3dVector(seg.reshape(-1,3)/255.)
op.visualization.draw_geometries([pcd2],window_name=name)
op.io.write_point_cloud(os.path.join(args.data,'point_clouds',name[:-3]+'pcd'), pcd)
op.io.write_point_cloud(os.path.join(args.data,'point_clouds',name[:-4]+'_seg.pcd'), pcd2)