Trains random gaussians to fit an image.
| 14 | |
| 15 | |
| 16 | class SimpleTrainer: |
| 17 | """Trains random gaussians to fit an image.""" |
| 18 | |
| 19 | def __init__( |
| 20 | self, |
| 21 | gt_image: Tensor, |
| 22 | num_points: int = 2000, |
| 23 | ): |
| 24 | self.device = torch.device("cuda:0") |
| 25 | self.gt_image = gt_image.to(device=self.device) |
| 26 | self.num_points = num_points |
| 27 | |
| 28 | fov_x = math.pi / 2.0 |
| 29 | self.H, self.W = gt_image.shape[0], gt_image.shape[1] |
| 30 | self.focal = 0.5 * float(self.W) / math.tan(0.5 * fov_x) |
| 31 | self.img_size = torch.tensor([self.W, self.H, 1], device=self.device) |
| 32 | |
| 33 | self._init_gaussians() |
| 34 | |
| 35 | def _init_gaussians(self): |
| 36 | """Random gaussians""" |
| 37 | bd = 2 |
| 38 | |
| 39 | self.means = bd * (torch.rand(self.num_points, 3, device=self.device) - 0.5) |
| 40 | self.scales = torch.rand(self.num_points, 3, device=self.device) |
| 41 | d = 3 |
| 42 | self.rgbs = torch.rand(self.num_points, d, device=self.device) |
| 43 | |
| 44 | u = torch.rand(self.num_points, 1, device=self.device) |
| 45 | v = torch.rand(self.num_points, 1, device=self.device) |
| 46 | w = torch.rand(self.num_points, 1, device=self.device) |
| 47 | |
| 48 | self.quats = torch.cat( |
| 49 | [ |
| 50 | torch.sqrt(1.0 - u) * torch.sin(2.0 * math.pi * v), |
| 51 | torch.sqrt(1.0 - u) * torch.cos(2.0 * math.pi * v), |
| 52 | torch.sqrt(u) * torch.sin(2.0 * math.pi * w), |
| 53 | torch.sqrt(u) * torch.cos(2.0 * math.pi * w), |
| 54 | ], |
| 55 | -1, |
| 56 | ) |
| 57 | self.opacities = torch.ones((self.num_points), device=self.device) |
| 58 | |
| 59 | self.viewmat = torch.tensor( |
| 60 | [ |
| 61 | [1.0, 0.0, 0.0, 0.0], |
| 62 | [0.0, 1.0, 0.0, 0.0], |
| 63 | [0.0, 0.0, 1.0, 8.0], |
| 64 | [0.0, 0.0, 0.0, 1.0], |
| 65 | ], |
| 66 | device=self.device, |
| 67 | ) |
| 68 | self.background = torch.zeros(d, device=self.device) |
| 69 | |
| 70 | self.means.requires_grad = True |
| 71 | self.scales.requires_grad = True |
| 72 | self.quats.requires_grad = True |
| 73 | self.rgbs.requires_grad = True |