Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import h5py | |
| import numpy as np | |
| import torch | |
| from unik3d.datasets.image_dataset import ImageDataset | |
| from unik3d.datasets.sequence_dataset import SequenceDataset | |
| from unik3d.datasets.utils import DatasetFromList | |
| class DiodeIndoor(ImageDataset): | |
| CAM_INTRINSIC = { | |
| "ALL": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]]) | |
| } | |
| min_depth = 0.01 | |
| max_depth = 25.0 | |
| depth_scale = 256.0 | |
| test_split = "val.txt" | |
| train_split = "train.txt" | |
| hdf5_paths = ["DiodeIndoor.hdf5"] | |
| def __init__( | |
| self, | |
| image_shape, | |
| split_file, | |
| test_mode, | |
| crop=None, | |
| benchmark=False, | |
| augmentations_db={}, | |
| normalize=True, | |
| mini=1.0, | |
| **kwargs, | |
| ): | |
| super().__init__( | |
| image_shape=image_shape, | |
| split_file=split_file, | |
| test_mode=test_mode, | |
| benchmark=benchmark, | |
| normalize=normalize, | |
| augmentations_db=augmentations_db, | |
| mini=mini, | |
| **kwargs, | |
| ) | |
| self.test_mode = test_mode | |
| # load annotations | |
| self.load_dataset() | |
| def load_dataset(self): | |
| h5file = h5py.File( | |
| os.path.join(self.data_root, self.hdf5_paths[0]), | |
| "r", | |
| libver="latest", | |
| swmr=True, | |
| ) | |
| txt_file = np.array(h5file[self.split_file]) | |
| txt_string = txt_file.tostring().decode("ascii")[:-1] # correct the -1 | |
| h5file.close() | |
| dataset = [] | |
| for line in txt_string.split("\n"): | |
| image_filename, depth_filename = line.strip().split(" ") | |
| sample = [ | |
| image_filename, | |
| depth_filename, | |
| ] | |
| dataset.append(sample) | |
| if not self.test_mode: | |
| dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini) | |
| self.dataset = DatasetFromList(dataset) | |
| self.log_load_dataset() | |
| def get_intrinsics(self, *args, **kwargs): | |
| return self.CAM_INTRINSIC["ALL"].clone() | |
| def get_mapper(self): | |
| return { | |
| "image_filename": 0, | |
| "depth_filename": 1, | |
| } | |
| def pre_pipeline(self, results): | |
| results = super().pre_pipeline(results) | |
| results["dense"] = [True] * self.num_copies | |
| results["quality"] = [1] * self.num_copies | |
| return results | |
| class DiodeIndoor_F(SequenceDataset): | |
| min_depth = 0.01 | |
| max_depth = 25.0 | |
| depth_scale = 1000.0 | |
| test_split = "train.txt" | |
| train_split = "train.txt" | |
| sequences_file = "sequences.json" | |
| hdf5_paths = ["DiodeIndoor-F.hdf5"] | |
| def __init__( | |
| self, | |
| image_shape: tuple[int, int], | |
| split_file: str, | |
| test_mode: bool, | |
| normalize: bool, | |
| augmentations_db: dict[str, float], | |
| resize_method: str, | |
| mini: float = 1.0, | |
| num_frames: int = 1, | |
| benchmark: bool = False, | |
| decode_fields: list[str] = ["image", "depth"], | |
| inplace_fields: list[str] = ["camera_params", "cam2w"], | |
| **kwargs, | |
| ) -> None: | |
| super().__init__( | |
| image_shape=image_shape, | |
| split_file=split_file, | |
| test_mode=test_mode, | |
| benchmark=benchmark, | |
| normalize=normalize, | |
| augmentations_db=augmentations_db, | |
| resize_method=resize_method, | |
| mini=mini, | |
| num_frames=num_frames, | |
| decode_fields=( | |
| decode_fields if not test_mode else [*decode_fields, "points"] | |
| ), | |
| inplace_fields=inplace_fields, | |
| **kwargs, | |
| ) | |
| def pre_pipeline(self, results): | |
| results = super().pre_pipeline(results) | |
| results["dense"] = [True] * self.num_frames * self.num_copies | |
| results["quality"] = [1] * self.num_frames * self.num_copies | |
| return results | |
| class DiodeOutdoor(ImageDataset): | |
| CAM_INTRINSIC = { | |
| "ALL": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]]) | |
| } | |
| min_depth = 0.1 | |
| max_depth = 80.0 | |
| log_mean = 0 | |
| log_std = 1 | |
| test_split = "diode_outdoor_val.txt" | |
| train_split = "diode_outdoor_train.txt" | |
| hdf5_paths = ["diode.hdf5"] | |
| def __init__( | |
| self, | |
| image_shape, | |
| split_file, | |
| test_mode, | |
| depth_scale=256, | |
| crop=None, | |
| benchmark=False, | |
| augmentations_db={}, | |
| normalize=True, | |
| resize_method="hard", | |
| mini=1.0, | |
| **kwargs, | |
| ): | |
| super().__init__( | |
| image_shape=image_shape, | |
| split_file=split_file, | |
| test_mode=test_mode, | |
| benchmark=benchmark, | |
| normalize=normalize, | |
| augmentations_db=augmentations_db, | |
| resize_method=resize_method, | |
| mini=mini, | |
| **kwargs, | |
| ) | |
| self.test_mode = test_mode | |
| self.depth_scale = depth_scale | |
| self.masker = AnnotationMask( | |
| min_value=self.min_depth, | |
| max_value=self.max_depth if test_mode else None, | |
| custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x, | |
| ) | |
| # load annotations | |
| self.load_dataset() | |
| def load_dataset(self): | |
| self.h5file = h5py.File( | |
| os.path.join(self.data_root, self.hdf5_path), | |
| "r", | |
| libver="latest", | |
| swmr=True, | |
| ) | |
| txt_file = np.array(self.h5file[self.split_file]) | |
| txt_string = txt_file.tostring().decode("ascii")[:-1] | |
| dataset = {"depth_filename": [], "image_filename": []} | |
| for line in txt_string.split("\n"): | |
| depth_filename = line.strip().split(" ")[1] | |
| img_name = line.strip().split(" ")[0] | |
| image_filename = img_name | |
| dataset["depth_filename"].append(depth_filename) | |
| dataset["image_filename"].append(image_filename) | |
| self.dataset = pl.from_dict(dataset) | |
| if not self.test_mode and self.mini: | |
| self.dataset = self.dataset[::2] | |
| class Diode(ImageDataset): | |
| CAM_INTRINSIC = { | |
| "ALL": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]]) | |
| } | |
| log_mean = 0 | |
| log_std = 1 | |
| min_depth = 0.6 | |
| max_depth = 80.0 | |
| test_split = "diode_val.txt" | |
| train_split = "diode_train.txt" | |
| hdf5_paths = ["diode.hdf5"] | |
| def __init__( | |
| self, | |
| image_shape, | |
| split_file, | |
| test_mode, | |
| depth_scale=256, | |
| crop=None, | |
| benchmark=False, | |
| augmentations_db={}, | |
| normalize=True, | |
| mini=1.0, | |
| **kwargs, | |
| ): | |
| super().__init__( | |
| image_shape=image_shape, | |
| split_file=split_file, | |
| test_mode=test_mode, | |
| benchmark=benchmark, | |
| normalize=normalize, | |
| augmentations_db=augmentations_db, | |
| mini=mini, | |
| **kwargs, | |
| ) | |
| self.test_mode = test_mode | |
| self.depth_scale = depth_scale | |
| self.masker = AnnotationMask( | |
| min_value=self.min_depth, | |
| max_value=self.max_depth if test_mode else None, | |
| custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x, | |
| ) | |
| # load annotations | |
| self.load_dataset() | |
| def load_dataset(self): | |
| self.h5file = h5py.File( | |
| os.path.join(self.data_root, self.hdf5_path), | |
| "r", | |
| libver="latest", | |
| swmr=True, | |
| ) | |
| txt_file = np.array(self.h5file[self.split_file]) | |
| txt_string = txt_file.tostring().decode("ascii")[:-1] | |
| dataset = {"depth_filename": [], "image_filename": []} | |
| for line in txt_string.split("\n"): | |
| depth_filename = line.strip().split(" ")[1] | |
| image_filename = line.strip().split(" ")[0] | |
| dataset["depth_filename"].append(depth_filename) | |
| dataset["image_filename"].append(image_filename) | |
| self.dataset = pl.from_dict(dataset) | |
| if not self.test_mode and self.mini: | |
| self.dataset = self.dataset[::2] | |
| def get_intrinsics(self, *args, **kwargs): | |
| return self.CAM_INTRINSIC["ALL"].clone() | |