-
Notifications
You must be signed in to change notification settings - Fork 1
/
render.py
133 lines (110 loc) · 3.8 KB
/
render.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
127
128
129
130
131
132
133
""" # noqa
___ __________________ ___________
/ _/__ ____ / __/ ___/ _/ __/ |/ / ___/ __/
/ _/ _ \/ __/ _\ \/ /___/ // _// / /__/ _/ # noqa
/_/ \___/_/ /___/\___/___/___/_/|_/\___/___/ # noqa
Author : Benjamin Blundell - [email protected]
render.py - Spit out an image using our renderer, given
an obj file, sigma and optional rotation and translation.
"""
if __name__ == "__main__":
"""
Render and save an image given a model.
Given a model, rotation and a sigma, spit out an image for us.
Parameters
----------
sigma : float
The sigma of the images in question - default 1.8
obj : string
The path to the obj file - default none.
rot : string
The rotation in angle/axis format - X,Y,Z.
Pass in as a string with comma separation - default none.
quat : string
The rotation in quaternion format - W,X,Y,Z.
Pass in as a string with comma separation - default none.
Returns
-------
None
"""
import argparse
import torch
import math
import util.plyobj as plyobj
from util.image import save_image, save_fits
from net.renderer import Splat
from util.math import TransTen, PointsTen, VecRot, angles_to_axis
from util.image import NormaliseBasic, NormaliseNull
from pyquaternion import Quaternion
parser = argparse.ArgumentParser(description="Render an image.")
parser.add_argument(
"--sigma",
type=float,
default=1.8,
help="sigma to learn from (default: 1.8)",
)
parser.add_argument(
"--obj",
help="The object file to render",
)
parser.add_argument(
"--rot",
help="The rotations in euler angles form, separated by commas, in degrees",
)
parser.add_argument(
"--quat",
help="The rotations as a Quaternion, W, X, Y and Z",
)
parser.add_argument(
"--norm",
default=False,
action="store_true",
help="Normalise with the basic normaliser.",
required=False,
)
args = parser.parse_args()
use_cuda = False
device = torch.device("cuda" if use_cuda else "cpu")
base_points = PointsTen(device=device)
base_points.from_points(plyobj.load_obj(args.obj))
# Which normalisation are we using?
normaliser = NormaliseNull()
if args.norm:
normaliser = NormaliseBasic()
mask = []
for _ in range(len(base_points)):
mask.append(1.0)
mask = torch.tensor(mask, device=device)
xt = torch.tensor([0.0], dtype=torch.float32)
yt = torch.tensor([0.0], dtype=torch.float32)
splat = Splat(device=device)
r = VecRot(0, 0, 0).to_ten(device=device)
if args.rot is not None:
tokens = args.rot.replace('"', "").split(",")
assert(len(tokens) == 3)
rx = math.radians(float(tokens[0]))
ry = math.radians(float(tokens[1]))
rz = math.radians(float(tokens[2]))
r = angles_to_axis(rx, ry, rz).to_ten(device=device)
# r = VecRot().to_ten(device=device)
if args.quat is not None:
tokens = args.quat.replace('"', "").split(",")
assert(len(tokens) == 4)
qw = math.radians(float(tokens[0]))
qx = math.radians(float(tokens[1]))
qy = math.radians(float(tokens[2]))
qz = math.radians(float(tokens[3]))
q = Quaternion(qw, qx, qy, qz)
v = q.get_axis()
a = q.radians
v[0] *= a
v[1] *= a
v[2] *= a
r = VecRot(v[0], v[1], v[2]).to_ten(device=device)
t = TransTen(xt, yt)
model = splat.render(base_points, r, t, mask, sigma=args.sigma)
model = model.reshape(1, 1, 128, 128)
model = normaliser.normalise(model)
model = model.squeeze()
save_image(model, name="renderer.jpg")
save_fits(model, name="renderer.fits")