Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """Create multiview data.""" | |
| import sys | |
| from os import path | |
| # Making sure you can run this, even if pulsar hasn't been installed yet. | |
| sys.path.insert(0, path.join(path.dirname(__file__), "..", "..")) | |
| def create_multiview(): | |
| """Test multiview optimization.""" | |
| import imageio | |
| # import cv2 | |
| # import skvideo.io | |
| import numpy as np | |
| import torch | |
| from pytorch3d.renderer.points.pulsar import Renderer | |
| from torch import nn | |
| from torch.autograd import Variable | |
| # Constructor. | |
| n_points = 10 | |
| width = 1000 | |
| height = 1000 | |
| class Model(nn.Module): | |
| """A dummy model to test the integration into a stacked model.""" | |
| def __init__(self): | |
| super(Model, self).__init__() | |
| self.gamma = 0.1 | |
| self.renderer = Renderer(width, height, n_points) | |
| def forward(self, vp, vc, vr, cam_params): | |
| # self.gamma *= 0.995 | |
| # print("gamma: ", self.gamma) | |
| return self.renderer.forward(vp, vc, vr, cam_params, self.gamma, 45.0) | |
| # Generate sample data. | |
| torch.manual_seed(1) | |
| vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0 | |
| vert_pos[:, 2] += 25.0 | |
| vert_pos[:, :2] -= 5.0 | |
| # print(vert_pos[0]) | |
| vert_col = torch.rand(n_points, 3, dtype=torch.float32) | |
| vert_rad = torch.rand(n_points, dtype=torch.float32) | |
| # Distortion. | |
| # vert_pos[:, 1] += 0.5 | |
| vert_col *= 0.5 | |
| # vert_rad *= 0.7 | |
| for device in [torch.device("cuda")]: | |
| model = Model().to(device) | |
| vert_pos = vert_pos.to(device) | |
| vert_col = vert_col.to(device) | |
| vert_rad = vert_rad.to(device) | |
| for angle_idx, angle in enumerate([-1.5, -0.8, -0.4, -0.1, 0.1, 0.4, 0.8, 1.5]): | |
| vert_pos_v = Variable(vert_pos, requires_grad=False) | |
| vert_col_v = Variable(vert_col, requires_grad=False) | |
| vert_rad_v = Variable(vert_rad, requires_grad=False) | |
| 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, | |
| ], | |
| dtype=torch.float32, | |
| ).to(device) | |
| cam_params_v = Variable(cam_params, requires_grad=False) | |
| result = model.forward(vert_pos_v, vert_col_v, vert_rad_v, cam_params_v) | |
| result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) | |
| imageio.imsave( | |
| "reference/examples_TestRenderer_test_multiview_%d.png" % (angle_idx), | |
| result_im, | |
| ) | |
| if __name__ == "__main__": | |
| create_multiview() | |