Spaces:
Sleeping
Sleeping
| """ | |
| Helpers used by both InferenceDataset and DynamicInputDataset. | |
| """ | |
| import os | |
| import numpy as np | |
| import PIL.Image as Image | |
| import torch | |
| import torch.nn.functional as F | |
| from einops import repeat | |
| # --------------------------------------------------------------------------- | |
| # Image + intrinsics helpers | |
| # --------------------------------------------------------------------------- | |
| def _resize_to_cover( | |
| image: Image.Image, | |
| target_h: int, | |
| target_w: int, | |
| ) -> tuple[Image.Image, int, int, int, int]: | |
| """Scale image so both target dimensions are covered, preserving aspect ratio.""" | |
| orig_w, orig_h = image.size | |
| scale = max(target_h / orig_h, target_w / orig_w) | |
| scaled_h = round(orig_h * scale) | |
| scaled_w = round(orig_w * scale) | |
| image = image.resize((scaled_w, scaled_h), Image.LANCZOS) | |
| return image, orig_w, orig_h, scaled_w, scaled_h | |
| def _center_square_crop( | |
| image: Image.Image, | |
| intrinsics: np.ndarray, | |
| crop_size: int, | |
| ) -> tuple[Image.Image, np.ndarray]: | |
| """Crop a square from the center and shift cx/cy accordingly.""" | |
| w, h = image.size | |
| start_h = (h - crop_size) // 2 | |
| start_w = (w - crop_size) // 2 | |
| image = image.crop((start_w, start_h, start_w + crop_size, start_h + crop_size)) | |
| adjusted = intrinsics.copy() | |
| adjusted[2] -= start_w | |
| adjusted[3] -= start_h | |
| return image, adjusted | |
| def _load_frame_image_and_intrinsics( | |
| frame: dict, | |
| image_base_dir: str, | |
| target_h: int, | |
| target_w: int, | |
| square_crop: bool, | |
| ) -> tuple[torch.Tensor, np.ndarray]: | |
| """Load one frame's image and return it with intrinsics adjusted for resize/crop.""" | |
| image_path = os.path.join(image_base_dir, frame["file_path"]) | |
| image, orig_w, orig_h, scaled_w, scaled_h = _resize_to_cover( | |
| Image.open(image_path), target_h, target_w | |
| ) | |
| raw_intrinsics = np.array([frame["fx"], frame["fy"], frame["cx"], frame["cy"]]) | |
| intrinsics = raw_intrinsics * np.array( | |
| [scaled_w / orig_w, scaled_h / orig_h, scaled_w / orig_w, scaled_h / orig_h] | |
| ) | |
| if square_crop: | |
| crop_size = min(target_h, target_w) | |
| image, intrinsics = _center_square_crop(image, intrinsics, crop_size) | |
| image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0 | |
| return image_tensor, intrinsics | |
| # --------------------------------------------------------------------------- | |
| # Pose helpers | |
| # --------------------------------------------------------------------------- | |
| def _w2c_stack_to_c2w(frames: list[dict]) -> torch.Tensor: | |
| """Invert per-frame w2c matrices and pin the last row to [0,0,0,1] for numerical safety.""" | |
| w2c_stack = np.stack([np.array(frame["w2c"]) for frame in frames]) | |
| c2ws = np.linalg.inv(w2c_stack) | |
| # Pinning the last row avoids floating-point drift from inversion | |
| c2w_out = repeat(torch.eye(4, dtype=torch.float32), 'h w -> b h w', b=len(frames)).clone() | |
| c2w_out[:, :3] = torch.from_numpy(c2ws).float()[:, :3] | |
| return c2w_out | |
| def _validate_poses(c2ws: torch.Tensor, label: str) -> None: | |
| """Raise if poses have extreme translations, NaN determinants, or non-unit rotation scale.""" | |
| if (c2ws[:, :3, 3] > 1e3).any(): | |
| raise ValueError(f"Large translation in {label} poses: {c2ws[:, :3, 3].max()}") | |
| dets = torch.det(c2ws[:, :3, :3]) | |
| if torch.isnan(dets).any(): | |
| raise ValueError(f"NaN in {label} pose determinants") | |
| if not torch.allclose(dets, dets.new_ones(dets.shape)): | |
| raise ValueError(f"Det of {label} poses not equal to 1") | |
| def _build_avg_camera_frame(input_c2ws: torch.Tensor) -> torch.Tensor: | |
| """Construct an SE(3) frame from the mean input camera orientation and position. | |
| The axes are built via Gram-Schmidt so they stay orthonormal even when | |
| individual forward vectors don't average to a unit vector. | |
| """ | |
| position_avg = input_c2ws[:, :3, 3].mean(0) | |
| forward_avg = F.normalize(input_c2ws[:, :3, 2].mean(0), dim=0) | |
| down_raw = input_c2ws[:, :3, 1].mean(0) | |
| down_avg = F.normalize(down_raw - down_raw.dot(forward_avg) * forward_avg, dim=0) | |
| right_avg = torch.linalg.cross(down_avg, forward_avg) | |
| return torch.cat([ | |
| torch.stack([right_avg, down_avg, forward_avg, position_avg], dim=1), | |
| torch.tensor([[0, 0, 0, 1]], dtype=torch.float32), | |
| ], dim=0) | |
| def _normalize_poses_to_input_frame( | |
| input_c2ws: torch.Tensor, | |
| target_c2ws: torch.Tensor, | |
| scene_scale_factor: float, | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Re-express all poses in the average input camera frame and unit-scale translations. | |
| Centering on the mean input view makes the model's coordinate system | |
| consistent across scenes regardless of global world-frame placement. | |
| Args: | |
| input_c2ws (torch.Tensor): Input camera-to-world matrices. | |
| target_c2ws (torch.Tensor): Target camera-to-world matrices. | |
| scene_scale_factor (float): Scene-level scale divisor applied before | |
| normalizing by the maximum translation magnitude. | |
| Returns: | |
| tuple[torch.Tensor, torch.Tensor]: Normalized (input_c2ws, target_c2ws). | |
| """ | |
| avg_frame = _build_avg_camera_frame(input_c2ws) | |
| world_to_avg = torch.inverse(avg_frame) | |
| input_c2ws = torch.matmul(world_to_avg.unsqueeze(0), input_c2ws) | |
| target_c2ws = torch.matmul(world_to_avg.unsqueeze(0), target_c2ws) | |
| translation_scale = 1.0 / (scene_scale_factor * input_c2ws[:, :3, 3].abs().max()) | |
| input_c2ws[:, :3, 3] *= translation_scale | |
| target_c2ws[:, :3, 3] *= translation_scale | |
| return input_c2ws, target_c2ws | |
| # --------------------------------------------------------------------------- | |
| # Frame window + ordering helpers | |
| # --------------------------------------------------------------------------- | |
| def _pick_frame_window( | |
| num_frames_total: int, | |
| min_dist: int, | |
| max_dist: int, | |
| ) -> list[int]: | |
| """Randomly sample a contiguous frame window within the given distance bounds.""" | |
| frame_dist = np.random.randint(min_dist, max_dist + 1) | |
| start = np.random.randint(0, num_frames_total - frame_dist) | |
| return list(range(start, start + frame_dist + 1)) | |
| def _select_input_frames( | |
| candidate_indices: list[int], | |
| selection_type: str, | |
| num_frames: int, | |
| scene_eval_info: dict | None = None, | |
| ) -> list[int]: | |
| """Return input frame indices from candidates using the specified selection strategy. | |
| Args: | |
| candidate_indices (list[int]): Pool of frame indices to select from. | |
| selection_type (str): One of 'random', 'uniform', 'kmeans', 'json_context'. | |
| num_frames (int): Number of input frames to select. | |
| scene_eval_info (dict | None): Per-scene metadata required by some strategies. | |
| Returns: | |
| list[int]: Selected input frame indices. | |
| """ | |
| if selection_type == 'random': | |
| return list(np.random.choice(candidate_indices, num_frames, replace=False)) | |
| elif selection_type == 'uniform': | |
| linspace_idx = np.linspace(0, len(candidate_indices) - 1, num_frames, dtype=int) | |
| return [candidate_indices[i] for i in linspace_idx] | |
| elif selection_type == 'kmeans': | |
| # k-means cluster centers precomputed offline for reproducible diverse coverage | |
| kmeans_key = f"fold_8_kmeans_{num_frames}_input" | |
| if not isinstance(scene_eval_info, dict) or kmeans_key not in scene_eval_info: | |
| raise KeyError(f"input_frame_select_type=kmeans requires key '{kmeans_key}' in scene index entry") | |
| return list(scene_eval_info[kmeans_key]) | |
| elif selection_type == 'json_context': | |
| if not isinstance(scene_eval_info, dict) or 'context' not in scene_eval_info: | |
| raise KeyError("input_frame_select_type=json_context requires scene entry to have 'context' list") | |
| return list(scene_eval_info['context']) | |
| else: | |
| raise NotImplementedError(f"Unknown input_frame_select_type: {selection_type}") | |
| def _select_target_frames( | |
| candidate_indices: list[int], | |
| selection_type: str, | |
| num_frames: int, | |
| uniform_every: int, | |
| scene_eval_info: dict | None = None, | |
| ) -> list[int]: | |
| """Return target frame indices from candidates using the specified selection strategy. | |
| Args: | |
| candidate_indices (list[int]): Pool of frame indices to select from. | |
| selection_type (str): One of 'random', 'uniform', 'uniform_every', 'json_target'. | |
| num_frames (int): Number of target frames to select (unused for 'uniform_every'). | |
| uniform_every (int): Step size used by the 'uniform_every' strategy. | |
| scene_eval_info (dict | None): Per-scene metadata required by some strategies. | |
| Returns: | |
| list[int]: Selected target frame indices. | |
| """ | |
| start, end = candidate_indices[0], candidate_indices[-1] | |
| if selection_type == 'random': | |
| return list(np.random.choice(candidate_indices, num_frames, replace=False)) | |
| elif selection_type == 'uniform': | |
| return list(np.linspace(start, end, num_frames, dtype=int)) | |
| elif selection_type == 'uniform_every': | |
| return list(range(start, end + 1, uniform_every)) | |
| elif selection_type == 'json_target': | |
| if not isinstance(scene_eval_info, dict) or 'target' not in scene_eval_info: | |
| raise KeyError("target_frame_select_type=json_target requires scene entry to have 'target' list") | |
| return list(scene_eval_info['target']) | |
| else: | |
| raise NotImplementedError(f"Unknown target_frame_select_type: {selection_type}") | |
| def _apply_input_ordering( | |
| indices: list[int], | |
| reverse: bool, | |
| shuffle: bool, | |
| ) -> list[int]: | |
| """Optionally reverse and/or shuffle input frame order for training-time augmentation.""" | |
| if reverse: | |
| indices = indices[::-1] | |
| if shuffle: | |
| np.random.shuffle(indices) | |
| return indices | |