| 2126 | |
| 2127 | |
| 2128 | class RGBDImage: |
| 2129 | def __init__( |
| 2130 | self, |
| 2131 | rgb: torch.Tensor, # (b, q, h, w, 3) b: different scene, q: multiple imgs of same scene |
| 2132 | depth: T.Optional[torch.Tensor], # (b, q, h, w) |
| 2133 | camera: Camera, # (b, q) |
| 2134 | normal_w: T.Optional[torch.Tensor] = None, # (b, q, h, w, 3) surface normal in world coord |
| 2135 | hit_map: T.Optional[torch.Tensor] = None, # (b, q, h, w) bool, 1: valid |
| 2136 | feature: T.Optional[torch.Tensor] = None, # (b, q, h, w, f) feature |
| 2137 | ): |
| 2138 | self.rgb = rgb |
| 2139 | self.depth = depth |
| 2140 | self.camera = camera |
| 2141 | self.normal_w = normal_w |
| 2142 | self.hit_map = hit_map |
| 2143 | self.feature = feature |
| 2144 | |
| 2145 | def compute_ray_normal_dot_product(self) -> T.Union[torch.Tensor, None]: |
| 2146 | """ |
| 2147 | Compute the dot product between the normal_w and the camera ray |
| 2148 | |
| 2149 | Returns: |
| 2150 | (b, q, h, w) the dot product (cos(theta)) between normal_w and camera ray |
| 2151 | """ |
| 2152 | if self.camera is None or self.normal_w is None: |
| 2153 | return None |
| 2154 | ray: Ray = self.camera.generate_camera_rays(device=self.normal_w.device) # (b, q, h, w) |
| 2155 | ray_direction_w = ray.directions_w # (b, q, h, w, 3) |
| 2156 | dot_prod = (ray_direction_w * self.normal_w).sum(dim=-1) # (b, q, h, w) |
| 2157 | return dot_prod # can be negative |
| 2158 | |
| 2159 | def index_select(self, dim: int, index: torch.Tensor) -> 'RGBDImage': |
| 2160 | rgbd_image = self.clone() |
| 2161 | for attr_name in ['rgb', 'depth', 'normal_w', 'hit_map', 'camera', 'feature']: |
| 2162 | arr = getattr(rgbd_image, attr_name, None) |
| 2163 | if arr is not None: |
| 2164 | setattr(rgbd_image, attr_name, arr.index_select(dim=dim, index=index)) |
| 2165 | return rgbd_image |
| 2166 | |
| 2167 | def chunk(self, chunks: int, dim: int = 0) -> T.List['RGBDImage']: |
| 2168 | out_dict = dict() |
| 2169 | total = None |
| 2170 | for attr_name in ['rgb', 'depth', 'normal_w', 'hit_map', 'camera', 'feature']: |
| 2171 | arr = getattr(self, attr_name, None) |
| 2172 | if arr is not None: |
| 2173 | chunked_arr = arr.chunk(chunks=chunks, dim=dim) |
| 2174 | out_dict[attr_name] = chunked_arr |
| 2175 | if total is None: |
| 2176 | total = len(chunked_arr) |
| 2177 | else: |
| 2178 | assert len(chunked_arr) == total |
| 2179 | rgbd_images = [] |
| 2180 | for i in range(total): |
| 2181 | d = dict() |
| 2182 | for attr_name in out_dict: |
| 2183 | d[attr_name] = out_dict[attr_name][i] |
| 2184 | rgbd_image = RGBDImage(**d) |
| 2185 | rgbd_images.append(rgbd_image) |
no outgoing calls
no test coverage detected