-
Notifications
You must be signed in to change notification settings - Fork 12
/
visualization.py
96 lines (73 loc) · 2.88 KB
/
visualization.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
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
# open3d == 0.10.0.1
import open3d
import argparse
import numpy as np
from helper_ply import read_ply
s3dis_color = [[255,248,220], [220,220,220], [139,71,38], [238,197,145], [70,130,180], [179,238,58], [110,139,61], [105,105,105], [0,0,128], [205,92,92], [244,164,96], [147,112,219], [255,228,225]]
semantic3d_color = [[220,220,220], [154,205,50], [0,100,0], [238,220,130], [139,115,85], [70,130,180], [255,231,186], [255,99,71]]
def draw_pc(pc_xyzrgb):
"""
Draw Point Cloud
"""
pc = open3d.geometry.PointCloud()
pc.points = open3d.utility.Vector3dVector(pc_xyzrgb[:, 0:3])
if pc_xyzrgb.shape[1] == 3:
open3d.draw_geometries([pc])
return 0
if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255
pc.colors = open3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.)
else:
pc.colors = open3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6])
open3d.visualization.draw_geometries([pc])
return 0
def s3dis(origin_file_name, label_file_name):
"""
Visualization for S3DIS
"""
origin_data = read_ply(origin_file_name)
label_data = read_ply(label_file_name)
pred = label_data["pred"]
label = label_data["label"]
x = origin_data["x"]
y = origin_data["y"]
z = origin_data["z"]
r = origin_data["red"]
g = origin_data["green"]
b = origin_data["blue"]
label_rgb = np.array([s3dis_color[i] for i in label])
pred_rgb = np.array([s3dis_color[i] for i in pred])
xyzrgb = np.vstack((x,y,z,r,g,b)).T
draw_pc(xyzrgb)
pred_xyzrgb = np.vstack(([x,y,z],pred_rgb.T)).T
draw_pc(pred_xyzrgb)
label_xyzrgb = np.vstack(([x,y,z],label_rgb.T)).T
draw_pc(label_xyzrgb)
def semantic3d(origin_file_name, label_file_name):
"""
Visualization for Semantic3D
"""
origin_data = read_ply(origin_file_name)
with open(label_file_name, 'r') as f: pred = f.readlines()
x = origin_data["x"]
y = origin_data["y"]
z = origin_data["z"]
r = origin_data["red"]
g = origin_data["green"]
b = origin_data["blue"]
pred_rgb = np.array([semantic3d_color[int(i)-1] for i in pred])
xyzrgb = np.vstack((x,y,z,r,g,b)).T
draw_pc(xyzrgb)
pred_xyzrgb = np.vstack(([x,y,z],pred_rgb.T)).T
draw_pc(pred_xyzrgb)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='s3dis', help='s3dis or semantic3d [default: s3dis]')
parser.add_argument('--ply_path', type=str, help='origin ply path')
parser.add_argument('--label_path', type=str, help='label path')
args = parser.parse_args()
if args.dataset == 's3dis':
s3dis(args.ply_path, args.label_path)
elif args.dataset == 'semantic3d':
semantic3d(args.ply_path, args.label_path)
else:
print("Wrong dataset. S3DIS or Semantic3D")