|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Test right hand/left hand system compatibility.""" |
|
import sys |
|
import unittest |
|
from os import path |
|
|
|
import numpy as np |
|
import torch |
|
from torch import nn |
|
|
|
|
|
sys.path.insert(0, path.join(path.dirname(__file__), "..")) |
|
devices = [torch.device("cuda"), torch.device("cpu")] |
|
|
|
|
|
n_points = 10 |
|
width = 1_000 |
|
height = 1_000 |
|
|
|
|
|
class SceneModel(nn.Module): |
|
"""A simple model to demonstrate use in Modules.""" |
|
|
|
def __init__(self): |
|
super(SceneModel, self).__init__() |
|
from pytorch3d.renderer.points.pulsar import Renderer |
|
|
|
self.gamma = 1.0 |
|
|
|
torch.manual_seed(1) |
|
vert_pos = torch.rand((1, n_points, 3), dtype=torch.float32) * 10.0 |
|
vert_pos[:, :, 2] += 25.0 |
|
vert_pos[:, :, :2] -= 5.0 |
|
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False)) |
|
self.register_parameter( |
|
"vert_col", |
|
nn.Parameter( |
|
torch.zeros(1, n_points, 3, dtype=torch.float32), requires_grad=True |
|
), |
|
) |
|
self.register_parameter( |
|
"vert_rad", |
|
nn.Parameter( |
|
torch.ones(1, n_points, dtype=torch.float32) * 0.001, |
|
requires_grad=False, |
|
), |
|
) |
|
self.register_parameter( |
|
"vert_opy", |
|
nn.Parameter( |
|
torch.ones(1, n_points, dtype=torch.float32), requires_grad=False |
|
), |
|
) |
|
self.register_buffer( |
|
"cam_params", |
|
torch.tensor( |
|
[ |
|
[ |
|
np.sin(angle) * 35.0, |
|
0.0, |
|
30.0 - np.cos(angle) * 35.0, |
|
0.0, |
|
-angle, |
|
0.0, |
|
5.0, |
|
2.0, |
|
] |
|
for angle in [-1.5, -0.8, -0.4, -0.1, 0.1, 0.4, 0.8, 1.5] |
|
], |
|
dtype=torch.float32, |
|
), |
|
) |
|
self.renderer = Renderer(width, height, n_points) |
|
|
|
def forward(self, cam=None): |
|
if cam is None: |
|
cam = self.cam_params |
|
n_views = 8 |
|
else: |
|
n_views = 1 |
|
return self.renderer.forward( |
|
self.vert_pos.expand(n_views, -1, -1), |
|
self.vert_col.expand(n_views, -1, -1), |
|
self.vert_rad.expand(n_views, -1), |
|
cam, |
|
self.gamma, |
|
45.0, |
|
return_forward_info=True, |
|
) |
|
|
|
|
|
class TestSmallSpheres(unittest.TestCase): |
|
"""Test small sphere rendering and gradients.""" |
|
|
|
def test_basic(self): |
|
for device in devices: |
|
|
|
model = SceneModel().to(device) |
|
angle = 0.0 |
|
for _ in range(50): |
|
cam_control = torch.tensor( |
|
[ |
|
[ |
|
np.sin(angle) * 35.0, |
|
0.0, |
|
30.0 - np.cos(angle) * 35.0, |
|
0.0, |
|
-angle, |
|
0.0, |
|
5.0, |
|
2.0, |
|
] |
|
], |
|
dtype=torch.float32, |
|
).to(device) |
|
result, forw_info = model(cam=cam_control) |
|
sphere_ids = model.renderer.sphere_ids_from_result_info_nograd( |
|
forw_info |
|
) |
|
|
|
for idx in range(n_points): |
|
self.assertTrue( |
|
(sphere_ids == idx).sum() > 0, "Sphere ID %d missing!" % (idx) |
|
) |
|
|
|
|
|
|
|
|
|
|
|
loss = ((result - torch.ones_like(result)).abs()).sum() |
|
loss.backward() |
|
|
|
self.assertTrue(torch.all(model.vert_col.grad[:, :, 0].abs() > 0.0)) |
|
angle += 0.15 |
|
|
|
|
|
if __name__ == "__main__": |
|
unittest.main() |
|
|