-
Notifications
You must be signed in to change notification settings - Fork 26
/
vis.py
74 lines (68 loc) · 2.72 KB
/
vis.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
# parse arguments
import argparse
parser = argparse.ArgumentParser(description='Tangent convolutions')
parser.add_argument('config', type=str, metavar='config', help='config json file')
parser.add_argument('scan', type=str, metavar='scan', help='scan name')
parser.add_argument('disp', type=str, metavar='disp', help='display type in [g (GT), p (prediction), c (color)]')
parser.add_argument('--raw', action="store_true", help='do training network')
args = parser.parse_args()
from util.path_config import *
from util.dataset_params import *
from util.cloud import *
class param:
def __init__(self, config):
self.dataset_dir = get_tc_path() + config['pre_dataset_dir']
self.parametrized_dir = get_tc_path() + config['pre_output_dir']
self.output_dir = os.path.join(get_tc_path(), config['co_experiment_dir'], config['co_output_dir'])
dataset_type = config['pre_dataset_param']
if dataset_type == "stanford":
self.d_par = stanford_params()
elif dataset_type == "scannet":
self.d_par = scannet_params()
elif dataset_type == "semantic3d":
self.d_par = semantic3d_params()
# reading configuration files
import os, sys
sys.path.append('util')
from config_reader import *
from cloud import *
config = config_reader(args.config)
par = param(config)
scan_name = args.scan
disp = args.disp
if args.raw:
cloud_file_name = os.path.join(par.dataset_dir, scan_name, "scan.pcd")
pcd = read_point_cloud(cloud_file_name)
estimate_normals(pcd, search_param=KDTreeSearchParamHybrid(radius=1.0, max_nn=100))
if disp != 'c':
if disp == 'g':
labels_file_name = os.path.join(par.dataset_dir, scan_name, "scan.labels")
elif disp == 'p':
labels_file_name = os.path.join(par.output_dir, scan_name, "extrapolated.labels")
lbl = read_txt_labels(labels_file_name)
remapped = remap_colors(np.asarray(lbl, dtype='int32'), color_map=par.d_par.color_map) / 255.0
idx = np.in1d(lbl, np.asarray([0]), invert=True)
pcd.points = Vector3dVector(np.asarray(pcd.points)[idx])#[nonzero])
pcd.colors = Vector3dVector(remapped[idx])
draw_geometries([pcd])
else:
if disp == 'p':
cloud_file_path = os.path.join(par.output_dir, scan_name)
else:
cloud_file_path = os.path.join(par.parametrized_dir, scan_name, '0')
scan = ScanData()
scan.load(cloud_file_path, 3)
if disp == 'c':
draw_geometries([scan.clouds[0]])
else:
if disp == 'g':
lbl = scan.labels_gt[0]
elif disp == 'p':
lbl = scan.labels_pr[0]
remapped = remap_colors(np.copy(np.asarray(lbl, dtype='int32')), color_map=par.d_par.color_map) / 255.0
idx = np.in1d(lbl, np.asarray([0]), invert=True)
pcd = PointCloud()
pcd.points = Vector3dVector(np.copy(np.asarray(scan.clouds[0].points)[idx]))
pcd.colors = Vector3dVector(remapped[idx])
estimate_normals(pcd)
draw_geometries([pcd])