-
Notifications
You must be signed in to change notification settings - Fork 0
/
visualize.py
126 lines (102 loc) · 3.54 KB
/
visualize.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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import h5py
import argparse
import numpy as np
from time import time
from PIL import Image
from alive_progress import alive_bar
import os
import math
import colorsys
import random
def get_n_hls_colors(num):
hls_colors = []
i = 0
step = 360.0 / num
while i < 360:
h = i
s = 90 + random.random() * 10
l = 50 + random.random() * 10
_hlsc = [h / 360.0, l / 100.0, s / 100.0]
hls_colors.append(_hlsc)
i += step
return hls_colors
def ncolors(num):
rgb_colors = []
if num < 1:
return rgb_colors
hls_colors = get_n_hls_colors(num)
for hlsc in hls_colors:
_r, _g, _b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])
r, g, b = [int(x * 255.0) for x in (_r, _g, _b)]
rgb_colors.append([r, g, b])
return rgb_colors
OUTPUTPATH = "/localhome/hja40/Desktop/Research/proj-motionnet/PC_motion_prediction/visualization"
rgb = np.array([[0, 0, 0], [255, 255, 255]] + ncolors(5))
def getFocalLength(FOV, height, width=None):
# Used to calculate the fixed intrinsic parameters for motionnet
# FOV is in radius, should be vertical angle
if width == None:
f = height / (2 * math.tan(FOV / 2))
return f
else:
fx = height / (2 * math.tan(FOV / 2))
fy = fx / height * width
return (fx, fy)
FOV = 50
img_width = 256
img_height = 256
fx, fy = getFocalLength(FOV / 180 * math.pi, img_height, img_width)
cy = img_height / 2
cx = img_width / 2
def existDir(dir):
if not os.path.exists(dir):
os.makedirs(dir)
def get_parser():
parser = argparse.ArgumentParser(description="Train motion net")
parser.add_argument(
"--result_path",
default=None,
metavar="FILE",
help="hdf5 file which contains the test results",
)
return parser
def renderResults(instance, result, prefix):
existDir(f"{OUTPUTPATH}/{instance}")
existDir(f"{OUTPUTPATH}/{instance}/{prefix}")
output_dir = f"{OUTPUTPATH}/{instance}/{prefix}"
camcs_per_point = result["camcs_per_point"][:]
category_per_point = result[f"{prefix}_category_per_point"][:]
instance_per_point = result[f"{prefix}_instance_per_point"][:]
mtype_per_point = result[f"{prefix}_mtype_per_point"][:]
instance_img = np.zeros((img_width, img_height, 3))
x = camcs_per_point[:, 0]
y = camcs_per_point[:, 1]
z = camcs_per_point[:, 2]
new_x = (x * fx / (-z) + cx).astype(int)
new_y = (-(y * fy / (-z)) + cy).astype(int)
part_index = np.where(category_per_point != 3)
base_index = np.where(category_per_point == 3)
x_min = np.min(new_x[part_index])
x_max = np.max(new_x[part_index])
y_min = np.min(new_y[part_index])
y_max = np.max(new_y[part_index])
instance_img[x_min, y_min] = rgb[2]
instance_img[x_max, y_max] = rgb[3]
# import pdb
# pdb.set_trace()
# instance_img[new_x[base_index], new_y[base_index]] = rgb[1]
instance_img[new_y[part_index], new_x[part_index]] = rgb[instance_per_point[part_index].astype(int) + 2]
image = Image.fromarray(np.uint8(instance_img))
image = image.convert('RGB')
image.save(f"{output_dir}/instance.png")
if __name__ == "__main__":
start = time()
args = get_parser().parse_args()
results = h5py.File(args.result_path)
instances = results.keys()
# with alive_bar(len(instances)) as bar:
for instance in instances:
# renderResults(instance, results[instance], "gt")
renderResults(instance, results[instance], "gt")
stop = time()
print(f'Total time duration: {stop - start}')