Spaces:
Running on Zero
Running on Zero
| import json | |
| from dataclasses import dataclass | |
| from functools import cached_property | |
| from io import BytesIO | |
| from pathlib import Path | |
| from typing import Literal, Optional | |
| import torch | |
| import torchvision.transforms as tf | |
| from einops import rearrange, repeat | |
| from jaxtyping import Float, UInt8 | |
| from PIL import Image | |
| from torch import Tensor | |
| from torch.utils.data import IterableDataset | |
| import numpy as np | |
| from ..geometry.projection import get_fov | |
| from .dataset import DatasetCfgCommon | |
| from .shims.augmentation_shim import apply_augmentation_shim | |
| from .shims.crop_shim import apply_crop_shim | |
| from .types import Stage | |
| from .view_sampler import ViewSampler | |
| class DatasetDL3DVCfg(DatasetCfgCommon): | |
| name: Literal["dl3dv"] | |
| roots: list[Path] | |
| max_fov: float | |
| augment: bool | |
| test_len: int | |
| test_chunk_interval: int | |
| train_times_per_scene: int | |
| test_times_per_scene: int | |
| ori_image_shape: list[int] | |
| skip_bad_shape: bool = True | |
| near: float = -1.0 | |
| far: float = -1.0 | |
| shuffle_val: bool = True | |
| min_views: int = 0 | |
| max_views: int = 0 | |
| highres: bool = False | |
| sort_target_index: Optional[bool] = False | |
| overfit_max_views: Optional[int] = None | |
| sort_context_index: Optional[bool] = False | |
| use_index_to_load_chunk: Optional[bool] = False | |
| class DatasetDL3DV(IterableDataset): | |
| cfg: DatasetDL3DVCfg | |
| stage: Stage | |
| view_sampler: ViewSampler | |
| to_tensor: tf.ToTensor | |
| chunks: list[Path] | |
| near: float = 0.1 | |
| far: float = 1000.0 | |
| def __init__( | |
| self, | |
| cfg: DatasetDL3DVCfg, | |
| stage: Stage, | |
| view_sampler: ViewSampler, | |
| ) -> None: | |
| super().__init__() | |
| self.cfg = cfg | |
| self.stage = stage | |
| self.view_sampler = view_sampler | |
| self.to_tensor = tf.ToTensor() | |
| if cfg.near != -1: | |
| self.near = cfg.near | |
| if cfg.far != -1: | |
| self.far = cfg.far | |
| # Collect chunks. | |
| self.chunks = [] | |
| for i, root in enumerate(cfg.roots): | |
| root = root / self.data_stage | |
| if self.cfg.use_index_to_load_chunk: | |
| with open(root / "index.json", "r") as f: | |
| json_dict = json.load(f) | |
| root_chunks = sorted(list(set(json_dict.values()))) | |
| else: | |
| root_chunks = sorted( | |
| [path for path in root.iterdir() if path.suffix == ".torch"] | |
| ) | |
| self.chunks.extend(root_chunks) | |
| if self.cfg.overfit_to_scene is not None: | |
| chunk_path = self.index[self.cfg.overfit_to_scene] | |
| self.chunks = [chunk_path] * len(self.chunks) | |
| if self.stage == "test": | |
| # fast testing | |
| self.chunks = self.chunks[:: cfg.test_chunk_interval] | |
| if self.stage == "val": | |
| self.chunks = self.chunks * int(1e6 // len(self.chunks)) | |
| def shuffle(self, lst: list) -> list: | |
| indices = torch.randperm(len(lst)) | |
| return [lst[x] for x in indices] | |
| def __iter__(self): | |
| # Chunks must be shuffled here (not inside __init__) for validation to show | |
| # random chunks. | |
| if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): | |
| self.chunks = self.shuffle(self.chunks) | |
| # When testing, the data loaders alternate chunks. | |
| worker_info = torch.utils.data.get_worker_info() | |
| if self.stage == "test" and worker_info is not None: | |
| self.chunks = [ | |
| chunk | |
| for chunk_index, chunk in enumerate(self.chunks) | |
| if chunk_index % worker_info.num_workers == worker_info.id | |
| ] | |
| for chunk_path in self.chunks: | |
| # Load the chunk. | |
| chunk = torch.load(chunk_path) | |
| if self.cfg.overfit_to_scene is not None: | |
| item = [x for x in chunk if x["key"] | |
| == self.cfg.overfit_to_scene] | |
| assert len(item) == 1 | |
| if self.stage == "test": | |
| chunk = item | |
| else: | |
| chunk = item * len(chunk) | |
| if self.stage in (("train", "val") if self.cfg.shuffle_val else ("train")): | |
| chunk = self.shuffle(chunk) | |
| times_per_scene = ( | |
| self.cfg.test_times_per_scene | |
| if self.stage == "test" | |
| else self.cfg.train_times_per_scene | |
| ) | |
| for run_idx in range(int(times_per_scene * len(chunk))): | |
| example = chunk[run_idx // times_per_scene] | |
| extrinsics, intrinsics = self.convert_poses(example["cameras"]) | |
| scene = example["key"] | |
| try: | |
| extra_kwargs = {} | |
| if self.cfg.overfit_to_scene is not None and self.stage != "test": | |
| extra_kwargs.update( | |
| { | |
| "max_num_views": ( | |
| 148 | |
| if self.cfg.overfit_max_views is None | |
| else self.cfg.overfit_max_views | |
| ) | |
| } | |
| ) | |
| out_data = self.view_sampler.sample( | |
| scene, | |
| extrinsics, | |
| intrinsics, | |
| min_context_views=self.cfg.min_views, | |
| max_context_views=self.cfg.max_views, | |
| **extra_kwargs, | |
| ) | |
| if isinstance(out_data, tuple): | |
| context_indices, target_indices = out_data[:2] | |
| c_list = [ | |
| ( | |
| context_indices.sort()[0] | |
| if self.cfg.sort_context_index | |
| else context_indices | |
| ) | |
| ] | |
| t_list = [ | |
| ( | |
| target_indices.sort()[0] | |
| if self.cfg.sort_target_index | |
| else target_indices | |
| ) | |
| ] | |
| if isinstance(out_data, list): | |
| c_list = [ | |
| ( | |
| a.context.sort()[0] | |
| if self.cfg.sort_context_index | |
| else a.context | |
| ) | |
| for a in out_data | |
| ] | |
| t_list = [ | |
| ( | |
| a.target.sort()[0] | |
| if self.cfg.sort_target_index | |
| else a.target | |
| ) | |
| for a in out_data | |
| ] | |
| except ValueError: | |
| # Skip because the example doesn't have enough frames. | |
| continue | |
| # Skip the example if the field of view is too wide. | |
| if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): | |
| continue | |
| for context_indices, target_indices in zip(c_list, t_list): | |
| # Load the images. | |
| context_images = [ | |
| example["images"][index.item()] for index in context_indices | |
| ] | |
| try: | |
| context_images = self.convert_images(context_images) | |
| except OSError: | |
| # some data might be corrupted | |
| continue | |
| target_images = [ | |
| example["images"][index.item()] for index in target_indices | |
| ] | |
| try: | |
| target_images = self.convert_images(target_images) | |
| except OSError: | |
| # some data might be corrupted | |
| continue | |
| # Skip the example if the images don't have the right shape. | |
| expected_shape = tuple( | |
| [3, *self.cfg.ori_image_shape] | |
| ) # (3, 270, 480) or (3, 540, 960) | |
| context_image_invalid = context_images.shape[1:] != expected_shape | |
| target_image_invalid = target_images.shape[1:] != expected_shape | |
| if self.cfg.skip_bad_shape and ( | |
| context_image_invalid or target_image_invalid | |
| ): | |
| print( | |
| f"Skipped bad example {example['key']}. Context shape was " | |
| f"{context_images.shape}, target shape was " | |
| f"{target_images.shape}, and expected shape was {expected_shape}" | |
| ) | |
| continue | |
| # check the extrinsics | |
| if any(torch.isnan(torch.det(extrinsics[context_indices][:, :3, :3]))): | |
| # print('invalid extrinsics') | |
| continue | |
| if any(torch.isnan(torch.det(extrinsics[target_indices][:, :3, :3]))): | |
| # print('invalid extrinsics') | |
| continue | |
| # check the extrinsics: translation could be very large | |
| # https://github.com/DL3DV-10K/Dataset/issues/34 | |
| if (extrinsics[context_indices][:, :3, 3] > 1e3).any(): | |
| # print('extremely large camera translation') | |
| continue | |
| if (extrinsics[target_indices][:, :3, 3] > 1e3).any(): | |
| # print('extremely large camera translation') | |
| continue | |
| if not torch.allclose(torch.det(extrinsics[context_indices][:, :3, :3]), torch.det(extrinsics[context_indices][:, :3, :3]).new_tensor(1)): | |
| # print('invalid extrinsics') | |
| continue | |
| if not torch.allclose(torch.det(extrinsics[target_indices][:, :3, :3]), torch.det(extrinsics[target_indices][:, :3, :3]).new_tensor(1)): | |
| # print('invalid extrinsics') | |
| continue | |
| nf_scale = 1.0 | |
| example_out = { | |
| "context": { | |
| "extrinsics": extrinsics[context_indices], | |
| "intrinsics": intrinsics[context_indices], | |
| "image": context_images, | |
| "near": self.get_bound("near", len(context_indices)) | |
| / nf_scale, | |
| "far": self.get_bound("far", len(context_indices)) | |
| / nf_scale, | |
| "index": context_indices, | |
| }, | |
| "target": { | |
| "extrinsics": extrinsics[target_indices], | |
| "intrinsics": intrinsics[target_indices], | |
| "image": target_images, | |
| "near": self.get_bound("near", len(target_indices)) | |
| / nf_scale, | |
| "far": self.get_bound("far", len(target_indices)) | |
| / nf_scale, | |
| "index": target_indices, | |
| }, | |
| "scene": scene, | |
| } | |
| if self.stage == "train" and self.cfg.augment: | |
| example_out = apply_augmentation_shim(example_out) | |
| if self.cfg.image_shape == list(context_images.shape[2:]): | |
| yield example_out | |
| else: | |
| yield apply_crop_shim(example_out, tuple(self.cfg.image_shape)) | |
| def convert_poses( | |
| self, | |
| poses: Float[Tensor, "batch 18"], | |
| ) -> tuple[ | |
| Float[Tensor, "batch 4 4"], # extrinsics | |
| Float[Tensor, "batch 3 3"], # intrinsics | |
| ]: | |
| b, _ = poses.shape | |
| # Convert the intrinsics to a 3x3 normalized K matrix. | |
| intrinsics = torch.eye(3, dtype=torch.float32) | |
| intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone() | |
| fx, fy, cx, cy = poses[:, :4].T | |
| intrinsics[:, 0, 0] = fx | |
| intrinsics[:, 1, 1] = fy | |
| intrinsics[:, 0, 2] = cx | |
| intrinsics[:, 1, 2] = cy | |
| # Convert the extrinsics to a 4x4 OpenCV-style C2W matrix. | |
| w2c = repeat(torch.eye(4, dtype=torch.float32), | |
| "h w -> b h w", b=b).clone() | |
| w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4) | |
| return w2c.inverse(), intrinsics | |
| def convert_images( | |
| self, | |
| images: list[UInt8[Tensor, "..."]], | |
| ) -> Float[Tensor, "batch 3 height width"]: | |
| torch_images = [] | |
| for image in images: | |
| image = Image.open(BytesIO(image.numpy().tobytes())) | |
| torch_images.append(self.to_tensor(image)) | |
| return torch.stack(torch_images) | |
| def get_bound( | |
| self, | |
| bound: Literal["near", "far"], | |
| num_views: int, | |
| ) -> Float[Tensor, " view"]: | |
| value = torch.tensor(getattr(self, bound), dtype=torch.float32) | |
| return repeat(value, "-> v", v=num_views) | |
| def data_stage(self) -> Stage: | |
| if self.cfg.overfit_to_scene is not None: | |
| return "test" | |
| if self.stage == "val": | |
| return "test" | |
| return self.stage | |
| def index(self) -> dict[str, Path]: | |
| merged_index = {} | |
| data_stages = [self.data_stage] | |
| if self.cfg.overfit_to_scene is not None: | |
| data_stages = ("test", "train") | |
| for data_stage in data_stages: | |
| for i, root in enumerate(self.cfg.roots): | |
| if not (root / data_stage).is_dir(): | |
| continue | |
| # Load the root's index. | |
| with (root / data_stage / "index.json").open("r") as f: | |
| index = json.load(f) | |
| index = {k: Path(root / data_stage / v) | |
| for k, v in index.items()} | |
| # The constituent datasets should have unique keys. | |
| assert not (set(merged_index.keys()) & set(index.keys())) | |
| # Merge the root's index into the main index. | |
| merged_index = {**merged_index, **index} | |
| return merged_index | |
| def __len__(self) -> int: | |
| if self.stage in ['train', 'test']: | |
| return ( | |
| min( | |
| len(self.index.keys()) * self.cfg.test_times_per_scene, | |
| self.cfg.test_len, | |
| ) | |
| if self.stage == "test" and self.cfg.test_len > 0 | |
| else len(self.index.keys()) * self.cfg.train_times_per_scene | |
| ) | |
| else: | |
| # set a very large value here to ensure the validation keep going | |
| # and do not exhaust; it will be wrap to length 1 anyway. | |
| return int(1e10) | |