-
Notifications
You must be signed in to change notification settings - Fork 0
/
abstraction_plot.py
92 lines (70 loc) · 3.08 KB
/
abstraction_plot.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
import torch
import numpy as np
from trimesh import Trimesh
from PIL import Image
import pyvista as pv
from tulsiani.inference import inference as tulsiani_inference
from yang.inference import inference as yang_inference
from loader.load_urocell import load_urocell_preprocessed
from common.transform import predictions_to_mesh_vertices
from graphics.write_mesh import cuboid_faces
from graphics.bruteforce_view import rotate_scene, bruteforce_view
from graphics.colors import colors
from tulsiani.parameters import params
prob_threshold = .5
remove_redundant = False # TODO
n_angles = 8 # Število vrednosti elevation in azimuth kota kamere
shape_image_size = 512
plot_image_height = 2 * shape_image_size
plot_image_width = 5 * shape_image_size
plot_image_size = (plot_image_width, plot_image_height)
def prediction_vertices_to_mesh(vertices):
p = vertices.shape[0]
v = vertices.reshape(-1, 3)
f = cuboid_faces.reshape(1, 6, 4).repeat(p, axis = 0)
f += 8 * np.arange(p).reshape(p, 1, 1)
f = f.reshape(-1, 4)
f = np.concatenate((4 * np.ones((f.shape[0], 1), dtype = int), f), axis = 1)
c = colors[:p].reshape(p, 1, 3).repeat(6, axis = 1)
c = c.reshape(-1, 3)
mesh = pv.PolyData(v, faces = f.reshape(-1), n_faces = 6 * p)
mesh["colors"] = c
return mesh
_, test = load_urocell_preprocessed(params.urocell_dir)
p = pv.Plotter(off_screen = True, window_size = [shape_image_size] * 2)
volume_meshes = [pv.wrap(Trimesh(test[i].vertices, test[i].faces - 1)) for i in range(10)]
test_set_plot = Image.new('RGB', plot_image_size, (0, 0, 0)) # Brez primitivov
best_angles = []
for i, volume_mesh in enumerate(volume_meshes):
volume_actor = p.add_mesh(volume_mesh)
image, best_angle = bruteforce_view(p, n_angles)
p.remove_actor(volume_actor)
best_angles.append(best_angle)
k = i % 5
j = i // 5
test_set_plot.paste(image, box = (k * shape_image_size, j * shape_image_size))
test_set_plot.save('results/test_set_plot.png')
def plot_predictions(X, method, transparent = False):
plot = Image.new('RGBA' if transparent else 'RGB', plot_image_size, (0, 0, 0))
predictions_vertices = predictions_to_mesh_vertices(X).cpu()
for i, volume_mesh in enumerate(volume_meshes):
if method == 'tulsiani':
v = predictions_vertices[i, X.prob[i].cpu() > prob_threshold].numpy()
else:
v = predictions_vertices[i, X.exist[i] == 1.].numpy()
predictions_mesh = prediction_vertices_to_mesh(v)
volume_actor = p.add_mesh(volume_mesh, opacity = .5)
predictions_actor = p.add_mesh(predictions_mesh, scalars = "colors", rgb = True)
k = i % 5
j = i // 5
image = rotate_scene(p, *best_angles[i], transparent = transparent)
plot.paste(image, box = (k * shape_image_size, j * shape_image_size))
p.remove_actor(volume_actor)
p.remove_actor(predictions_actor)
plot.save(f'results/{method}/abstraction_{method}.png')
X = tulsiani_inference(test)
if X is not None:
plot_predictions(X[0], 'tulsiani')
X = yang_inference(test)
if X is not None:
plot_predictions(X[0], 'yang')