Spaces:
Running on Zero
Running on Zero
| import math | |
| import os | |
| import json | |
| import re | |
| import cv2 | |
| from dataclasses import dataclass, field | |
| import random | |
| import imageio | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader, Dataset | |
| # from PIL import Image | |
| import PIL | |
| from diffusionGS.utils.typing import * | |
| from torch.utils.data._utils.collate import default_collate_fn_map | |
| class BaseDataModuleConfig: | |
| local_dir: str = '' | |
| local_eval_dir: str = '' | |
| view_idx_file_path: str = 'extra_files/evaluation_index_re10k.json' | |
| batch_size: int = 32 | |
| eval_batch_size: int = 1 | |
| eval_subset: int = -1 | |
| num_workers: int = 0 | |
| num_workers_val: int = 0 | |
| training_res: Optional[List[int]] = field(default_factory=lambda:[256,256]) | |
| patch_size: int = 8 | |
| sel_views_train: int = 4 | |
| sel_views: int = 4 # first view is the input image | |
| scene_scale_factor: float = 1.35 | |
| ################################# Image part ################################# | |
| square_crop: bool = True | |
| load_image: bool = True | |
| class BaseDataset(Dataset): | |
| def __init__(self, cfg: Any, split: str) -> None: | |
| super().__init__() | |
| self.cfg: BaseDataModuleConfig = cfg | |
| self.split = split | |
| if self.split == 'train': | |
| with open(cfg.local_dir, 'r') as f: | |
| all_scene_paths = f.read().splitlines() | |
| all_scene_paths = [path for path in all_scene_paths if path.strip()] | |
| self.uids = all_scene_paths#[10:] | |
| else: | |
| with open(cfg.local_eval_dir, 'r') as f: | |
| all_scene_paths = f.read().splitlines() | |
| all_scene_paths = [path for path in all_scene_paths if path.strip()] | |
| self.view_idx_list = dict() | |
| if self.cfg.view_idx_file_path != '': | |
| if os.path.exists(self.cfg.view_idx_file_path): | |
| with open(self.cfg.view_idx_file_path, 'r') as f: | |
| self.view_idx_list = json.load(f) | |
| # filter out None values, i.e. scenes that don't have specified input and targetviews | |
| self.view_idx_list_filtered = [k for k, v in self.view_idx_list.items() if v is not None] | |
| filtered_scene_paths = [] | |
| for scene in all_scene_paths: | |
| file_name = scene.split("/")[-1] | |
| scene_name = file_name.split(".")[0] | |
| if scene_name in self.view_idx_list_filtered: | |
| filtered_scene_paths.append(scene) | |
| all_scene_paths = filtered_scene_paths | |
| if self.cfg.eval_subset > 0: | |
| all_scene_paths = all_scene_paths[:self.cfg.eval_subset] | |
| self.uids = all_scene_paths | |
| # load anything | |
| if not self.uids: | |
| raise ValueError(f"No {split} scenes found; check the scene list and evaluation index") | |
| print(f"Loaded {len(self.uids)} {split} uids") | |
| # breakpoint() | |
| def __len__(self): | |
| return len(self.uids) | |
| def preprocess_frames(self, frames_chosen, image_paths_chosen): | |
| resize_h = self.cfg.training_res[0] | |
| patch_size = self.cfg.patch_size | |
| square_crop = self.cfg.square_crop | |
| images = [] | |
| intrinsics = [] | |
| for cur_frame, cur_image_path in zip(frames_chosen, image_paths_chosen): | |
| image = PIL.Image.open(cur_image_path) | |
| original_image_w, original_image_h = image.size | |
| resize_w = int(resize_h / original_image_h * original_image_w) | |
| resize_w = int(round(resize_w / patch_size) * patch_size) | |
| # if torch.distributed.get_rank() == 0: | |
| # import ipdb; ipdb.set_trace() | |
| image = image.resize((resize_w, resize_h), resample=PIL.Image.LANCZOS) | |
| if square_crop: | |
| min_size = min(resize_h, resize_w) | |
| start_h = (resize_h - min_size) // 2 | |
| start_w = (resize_w - min_size) // 2 | |
| image = image.crop((start_w, start_h, start_w + min_size, start_h + min_size)) | |
| image = np.array(image) / 255.0 | |
| image = torch.from_numpy(image).permute(2, 0, 1).float() | |
| fxfycxcy = np.array(cur_frame["fxfycxcy"]) | |
| resize_ratio_x = resize_w / original_image_w | |
| resize_ratio_y = resize_h / original_image_h | |
| fxfycxcy *= (resize_ratio_x, resize_ratio_y, resize_ratio_x, resize_ratio_y) | |
| if square_crop: | |
| fxfycxcy[2] -= start_w | |
| fxfycxcy[3] -= start_h | |
| fxfycxcy = torch.from_numpy(fxfycxcy).float() | |
| images.append(image) | |
| intrinsics.append(fxfycxcy) | |
| images = torch.stack(images, dim=0) | |
| intrinsics = torch.stack(intrinsics, dim=0) | |
| w2cs = np.stack([np.array(frame["w2c"]) for frame in frames_chosen]) | |
| c2ws = np.linalg.inv(w2cs) # (num_frames, 4, 4) | |
| c2ws = torch.from_numpy(c2ws).float() | |
| return images, intrinsics, c2ws | |
| def preprocess_poses( | |
| self, | |
| in_c2ws: torch.Tensor, | |
| scene_scale_factor=1.35, | |
| ): | |
| """ | |
| Preprocess the poses to: | |
| 1. translate and rotate the scene to align the average camera direction and position | |
| 2. rescale the whole scene to a fixed scale | |
| """ | |
| # Translation and Rotation | |
| # align coordinate system (OpenCV coordinate) to the mean camera | |
| # center is the average of all camera centers | |
| # average direction vectors are computed from all camera direction vectors (average down and forward) | |
| center = in_c2ws[:, :3, 3].mean(0) | |
| avg_forward = F.normalize(in_c2ws[:, :3, 2].mean(0), dim=-1) # average forward direction (z of opencv camera) | |
| avg_down = in_c2ws[:, :3, 1].mean(0) # average down direction (y of opencv camera) | |
| avg_right = F.normalize(torch.cross(avg_down, avg_forward, dim=-1), dim=-1) # (x of opencv camera) | |
| avg_down = F.normalize(torch.cross(avg_forward, avg_right, dim=-1), dim=-1) # (y of opencv camera) | |
| avg_pose = torch.eye(4, device=in_c2ws.device) # average c2w matrix | |
| avg_pose[:3, :3] = torch.stack([avg_right, avg_down, avg_forward], dim=-1) | |
| avg_pose[:3, 3] = center | |
| avg_pose = torch.linalg.inv(avg_pose) # average w2c matrix | |
| in_c2ws = avg_pose @ in_c2ws | |
| # Rescale the whole scene to a fixed scale | |
| scene_scale = torch.max(torch.abs(in_c2ws[:, :3, 3])) | |
| scene_scale = (scene_scale_factor * scene_scale).clamp_min(1e-8) | |
| in_c2ws[:, :3, 3] /= scene_scale | |
| return in_c2ws | |
| def _get_data(self, index): | |
| uid = self.uids[index] | |
| scene_path = uid.strip() | |
| data_json = json.load(open(scene_path, 'r')) | |
| frames = data_json["frames"] | |
| scene_name = data_json["scene_name"] | |
| #### conduct view selecting #### | |
| if self.split != 'train' and scene_name in self.view_idx_list: | |
| current_view_idx = self.view_idx_list[scene_name] | |
| image_indices= current_view_idx["context"][:1] + current_view_idx["target"] # only take first frame as input | |
| # breakpoint() | |
| else: | |
| # breakpoint() | |
| image_indices = random.sample(range(len(frames)), self.cfg.sel_views + self.cfg.sel_views_train) | |
| ######### load data ######### | |
| image_paths_chosen = [frames[ic]["image_path"] for ic in image_indices] | |
| frames_chosen = [frames[ic] for ic in image_indices] | |
| input_images, input_intrinsics, input_c2ws = self.preprocess_frames(frames_chosen, image_paths_chosen) | |
| # centerize and scale the poses (for unbounded scenes) | |
| scene_scale_factor = self.cfg.scene_scale_factor | |
| input_c2ws = self.preprocess_poses(input_c2ws, scene_scale_factor) | |
| image_indices = torch.tensor(image_indices).long().unsqueeze(-1) # [v, 1] | |
| scene_indices = torch.full_like(image_indices, index) # [v, 1] | |
| indices = torch.cat([image_indices, scene_indices], dim=-1) # [v, 2] | |
| # breakpoint() | |
| ret ={ | |
| 'rgbs': input_images, | |
| 'c2ws': input_c2ws, | |
| 'fxfycxcys': input_intrinsics, | |
| 'image_indices': image_indices, | |
| 'scene_indices': scene_indices, | |
| 'masks': torch.ones_like(input_images[:, :1]), | |
| 'indices': indices, | |
| 'uid': scene_name | |
| } | |
| # breakpoint() | |
| ret['rgbs_input'] = ret['rgbs'][:self.cfg.sel_views+1] | |
| ret['c2ws_input'] = ret['c2ws'][:self.cfg.sel_views+1] | |
| ret['masks_input'] = ret['masks'][:self.cfg.sel_views+1] | |
| ret['fxfycxcys_input'] = ret['fxfycxcys'][:self.cfg.sel_views+1] | |
| # breakpoint() | |
| return ret | |
| def __getitem__(self, index): | |
| try: | |
| return self._get_data(index) | |
| except Exception as e: | |
| raise RuntimeError(f"Could not load scene {self.uids[index]}: {e}") from e | |
| def collate(self, batch): | |
| from torch.utils.data._utils.collate import default_collate_fn_map | |
| ret = {} | |
| for key, value in batch[0].items(): | |
| if isinstance(value, str): | |
| ret[key] = [b[key] for b in batch] | |
| elif isinstance(value, torch.Tensor): | |
| ret[key] = torch.stack([b[key] for b in batch]) | |
| elif isinstance(value, np.ndarray): | |
| ret[key] = torch.stack([torch.from_numpy(b[key]) for b in batch]) | |
| else: | |
| print(key) | |
| ret[key] = default_collate_fn_map[type(batch[0][key])](batch) | |
| return ret | |