| 1237 | |
| 1238 | |
| 1239 | class PointersectRecord: |
| 1240 | def __init__( |
| 1241 | self, |
| 1242 | intersection_xyz_w: torch.Tensor, # (b, *m_shape, 3) |
| 1243 | intersection_surface_normal_w: torch.Tensor, # (b, *m_shape, 3) |
| 1244 | intersection_rgb: torch.Tensor, # (b, *m_shape, 3) |
| 1245 | blending_weights: torch.Tensor, # (b, *m_shape, k) k: # neighbor points |
| 1246 | neighbor_point_idxs: torch.Tensor, # long (b, *m_shape, k) |
| 1247 | neighbor_point_valid_len: torch.Tensor, # long (b, *m_shape) |
| 1248 | ray_t: torch.Tensor, # (b, *m_shape) |
| 1249 | ray_hit: torch.Tensor, # (b, *m_shape) bool |
| 1250 | ray_hit_logit: torch.Tensor, # (b, *m_shape) |
| 1251 | model_attn_weights: torch.Tensor, # (b, *m_shape, k+1, n_layers) |
| 1252 | refined_ray_hit: T.Optional[torch.Tensor] = None, # (b, *m_shape) bool |
| 1253 | model_info: T.Optional[T.Dict[str, T.Any]] = None, |
| 1254 | intersection_plane_normals_w: torch.Tensor = None, # (b, *m_shape, 3) |
| 1255 | geometry_weights: torch.Tensor = None, # (b, *m_shape, k) |
| 1256 | valid_neighbor_idx_mask: torch.Tensor = None, # (b, *m_shape, k) whether the neighbor_point_idxs is valid |
| 1257 | valid_plane_normal_mask: torch.Tensor = None, # (b, *m_shape) |
| 1258 | total_time: float = None, |
| 1259 | ): |
| 1260 | self.intersection_xyz_w = intersection_xyz_w |
| 1261 | self.intersection_surface_normal_w = intersection_surface_normal_w |
| 1262 | self.intersection_rgb = intersection_rgb |
| 1263 | self.blending_weights = blending_weights |
| 1264 | self.neighbor_point_idxs = neighbor_point_idxs |
| 1265 | self.neighbor_point_valid_len = neighbor_point_valid_len |
| 1266 | self.ray_t = ray_t |
| 1267 | self.ray_hit = ray_hit |
| 1268 | self.ray_hit_logit = ray_hit_logit |
| 1269 | self.model_attn_weights = model_attn_weights |
| 1270 | self.refined_ray_hit = refined_ray_hit |
| 1271 | self.intersection_plane_normals_w = intersection_plane_normals_w |
| 1272 | self.geometry_weights = geometry_weights |
| 1273 | self.valid_neighbor_idx_mask = valid_neighbor_idx_mask |
| 1274 | self.valid_plane_normal_mask = valid_plane_normal_mask |
| 1275 | self.total_time = total_time |
| 1276 | self.model_info = model_info |
| 1277 | |
| 1278 | # self.cached_info = cached_info # not saved nor concat nor reshaped |
| 1279 | |
| 1280 | self.attr_names = [ |
| 1281 | 'intersection_xyz_w', 'intersection_surface_normal_w', 'intersection_rgb', |
| 1282 | 'blending_weights', 'neighbor_point_idxs', 'neighbor_point_valid_len', |
| 1283 | 'ray_t', 'ray_hit', 'ray_hit_logit', 'model_attn_weights', |
| 1284 | 'refined_ray_hit', 'model_info', |
| 1285 | 'intersection_plane_normals_w', 'geometry_weights', |
| 1286 | 'valid_neighbor_idx_mask', 'valid_plane_normal_mask', |
| 1287 | ] |
| 1288 | |
| 1289 | def to(self, device: torch.device) -> 'PointersectRecord': |
| 1290 | for attr_name in self.attr_names: |
| 1291 | if attr_name == 'model_info': |
| 1292 | continue |
| 1293 | arr = getattr(self, attr_name, None) |
| 1294 | if arr is not None: |
| 1295 | setattr(self, attr_name, arr.to(device=device)) |
| 1296 | return self |
no outgoing calls
no test coverage detected