(self, cfg: Any)
| 64 | |
| 65 | class CustomImageOrbitDataset(Dataset): |
| 66 | def __init__(self, cfg: Any) -> None: |
| 67 | super().__init__() |
| 68 | self.cfg: CustomImageDataModuleConfig = parse_structured(CustomImageDataModuleConfig, cfg) |
| 69 | |
| 70 | self.n_views = self.cfg.n_test_views |
| 71 | assert self.n_views % self.cfg.num_views_output == 0 |
| 72 | |
| 73 | self.all_scenes = _parse_scene_list(self.cfg.image_list) |
| 74 | |
| 75 | azimuth_deg: Float[Tensor, "B"] = torch.linspace(0, 360.0, self.n_views + 1)[ |
| 76 | : self.n_views |
| 77 | ] |
| 78 | elevation_deg: Float[Tensor, "B"] = torch.full_like( |
| 79 | azimuth_deg, self.cfg.eval_elevation_deg |
| 80 | ) |
| 81 | camera_distances: Float[Tensor, "B"] = torch.full_like( |
| 82 | elevation_deg, self.cfg.eval_camera_distance |
| 83 | ) |
| 84 | |
| 85 | elevation = elevation_deg * math.pi / 180 |
| 86 | azimuth = azimuth_deg * math.pi / 180 |
| 87 | |
| 88 | # convert spherical coordinates to cartesian coordinates |
| 89 | # right hand coordinate system, x back, y right, z up |
| 90 | # elevation in (-90, 90), azimuth from +x to +y in (-180, 180) |
| 91 | camera_positions: Float[Tensor, "B 3"] = torch.stack( |
| 92 | [ |
| 93 | camera_distances * torch.cos(elevation) * torch.cos(azimuth), |
| 94 | camera_distances * torch.cos(elevation) * torch.sin(azimuth), |
| 95 | camera_distances * torch.sin(elevation), |
| 96 | ], |
| 97 | dim=-1, |
| 98 | ) |
| 99 | |
| 100 | # default scene center at origin |
| 101 | center: Float[Tensor, "B 3"] = torch.zeros_like(camera_positions) |
| 102 | # default camera up direction as +z |
| 103 | up: Float[Tensor, "B 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[ |
| 104 | None, : |
| 105 | ].repeat(self.n_views, 1) |
| 106 | |
| 107 | fovy_deg: Float[Tensor, "B"] = torch.full_like( |
| 108 | elevation_deg, self.cfg.eval_fovy_deg |
| 109 | ) |
| 110 | fovy = fovy_deg * math.pi / 180 |
| 111 | |
| 112 | lookat: Float[Tensor, "B 3"] = F.normalize(center - camera_positions, dim=-1) |
| 113 | right: Float[Tensor, "B 3"] = F.normalize(torch.cross(lookat, up), dim=-1) |
| 114 | up = F.normalize(torch.cross(right, lookat), dim=-1) |
| 115 | c2w3x4: Float[Tensor, "B 3 4"] = torch.cat( |
| 116 | [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]], |
| 117 | dim=-1, |
| 118 | ) |
| 119 | c2w: Float[Tensor, "B 4 4"] = torch.cat( |
| 120 | [c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1 |
| 121 | ) |
| 122 | c2w[:, 3, 3] = 1.0 |
| 123 |
nothing calls this directly
no test coverage detected