Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,922 Bytes
1ea89dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import json
import os
import h5py
import numpy as np
import torch
from unik3d.datasets.image_dataset import ImageDataset
from unik3d.datasets.utils import DatasetFromList
class HyperSim(ImageDataset):
min_depth = 0.01
max_depth = 50.0
depth_scale = 1000.0
test_split = "val.txt"
train_split = "train.txt"
intrisics_file = "intrinsics.json"
hdf5_paths = [f"hypersim/hypersim_{i}.hdf5" for i in range(8)]
def __init__(
self,
image_shape,
split_file,
test_mode,
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.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").strip("\n")
intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode("ascii")
intrinsics = json.loads(intrinsics)
# with open(os.path.join(os.environ["TMPDIR"], self.split_file), "w") as f:
# f.write(txt_string)
# with open(os.path.join(os.environ["TMPDIR"], self.intrisics_file), "w") as f:
# json.dump(intrinsics, f)
dataset = []
for line in txt_string.split("\n"):
image_filename, depth_filename, chunk_idx = line.strip().split(" ")
intrinsics_val = torch.tensor(
intrinsics[os.path.join(*image_filename.split("/")[:2])]
).squeeze()[:, :3]
sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]
dataset.append(sample)
h5file.close()
if not self.test_mode:
dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)
if self.test_mode:
dataset = self.chunk(dataset, chunk_dim=1, pct=0.1)
self.dataset = DatasetFromList(dataset)
self.log_load_dataset()
def get_mapper(self):
return {
"image_filename": 0,
"depth_filename": 1,
"K": 2,
"chunk_idx": 3,
}
def pre_pipeline(self, results):
results = super().pre_pipeline(results)
results["dense"] = [True] * self.num_copies
results["synthetic"] = [True] * self.num_copies
results["quality"] = [0] * self.num_copies
return results
|