-
Notifications
You must be signed in to change notification settings - Fork 2
/
visualization.py
82 lines (63 loc) · 4.01 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
import torch
import numpy as np
def make_pointcloud_obj(locs: torch.Tensor, actual: torch.Tensor, volume: torch.Tensor, side_len: int,
path: str) -> None:
"""
Plots a given pair of point cloud predictions and labels.
Taken from:
:param locs: (torch.Tensor) Predicted locations
:param actual: (torch.Tensor) Predicted label
:param volume: (torch.Tensor) Input volume to estimate shape of obj plot
:param side_len: (torch.Tensor) Scale parameter
:param path: (torch.Tensor) Path to store plot
"""
to_write = locs.cpu().numpy().astype(np.short)
to_write_act = actual.cpu().numpy().astype(np.short)
# Mean (shape) centering
mean = np.array([volume.shape[2] * side_len / 2, volume.shape[3] * side_len / 2, volume.shape[4] * side_len / 2])
to_write_act = to_write_act - mean
to_write = to_write - mean
with open((path + '_prediction.obj'), 'w') as f:
for line in to_write:
f.write("v " + " " + str(line[0]) + " " + str(line[1]) + " " + str(line[2]) +
" " + "0.5" + " " + "0.5" + " " + "1.0" + "\n")
# Corners of volume
f.write("v " + " " + "0" + " " + "0" + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + "0" + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + str(volume.shape[3] * side_len) + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + str(volume.shape[3] * side_len) + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + "0" + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + "0" + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + str(volume.shape[3] * side_len) + " " + str(
volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + str(volume.shape[3] * side_len) + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
with open((path + '_label.obj'), 'w') as f:
for line in to_write_act:
f.write("v " + " " + str(line[0]) + " " + str(line[1]) + " " + str(line[2]) +
" " + "0.19" + " " + "0.8" + " " + "0.19" + "\n")
# Corners of volume
f.write("v " + " " + "0" + " " + "0" + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + "0" + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + str(volume.shape[3] * side_len) + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + str(volume.shape[3] * side_len) + " " + "0" +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + "0" + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + "0" + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + str(volume.shape[2] * side_len) + " " + str(volume.shape[3] * side_len) + " " + str(
volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")
f.write("v " + " " + "0" + " " + str(volume.shape[3] * side_len) + " " + str(volume.shape[4] * side_len) +
" " + "1.0" + " " + "0.5" + " " + "0.5" + "\n")