| 13 | import random |
| 14 | |
| 15 | class Re10kMapDataset(Dataset): |
| 16 | def __init__(self, |
| 17 | opt, |
| 18 | shuffle=False, |
| 19 | training=True, |
| 20 | nearby_range=15): |
| 21 | self.opt = opt |
| 22 | self.shuffle = shuffle |
| 23 | self.training = training |
| 24 | self.seed = 42 |
| 25 | if nearby_range < opt.output_frames - 1: |
| 26 | print(f"Warning: nearby_range is less than output_frames - 1 on re10k, nearby_range: {nearby_range}, output_frames: {opt.output_frames}, change nearby_range to {opt.output_frames - 1}") |
| 27 | nearby_range = opt.output_frames - 1 |
| 28 | self.nearby_range = nearby_range |
| 29 | |
| 30 | self.root_path = opt.root_path |
| 31 | self.augmentation = self.opt.augmentation |
| 32 | self.set_transform() |
| 33 | |
| 34 | if self.training: |
| 35 | self.root_path = osp.join(self.root_path, "train") |
| 36 | else: |
| 37 | self.root_path = osp.join(self.root_path, "test") |
| 38 | |
| 39 | assert os.path.exists(os.path.join(self.root_path, "index.json")), "index.json not found" |
| 40 | with open(osp.join(self.root_path, "index.json"), "r") as f: |
| 41 | self.index = json.load(f) |
| 42 | |
| 43 | self.src_chunk_paths = sorted(glob.glob(osp.join(self.root_path, "*.torch"))) |
| 44 | assert len(self.src_chunk_paths) > 0, "No chunks found in the data path" |
| 45 | self._get_depth_path = lambda x: osp.join(osp.dirname(x), 'depth_{}'.format(osp.basename(x).replace('.torch', '.pt'))) |
| 46 | |
| 47 | _unordered_depth_paths = set([self._get_depth_path(path) for path in self.src_chunk_paths]) |
| 48 | assert len([glob.glob(path) for path in _unordered_depth_paths]) == len(self.src_chunk_paths), "Depth paths do not match" |
| 49 | |
| 50 | # Pre-compute all frame indices |
| 51 | self._all_frames = [] |
| 52 | self.chunk_data = {} |
| 53 | self.set_len = [] |
| 54 | self.ret_frames = [] |
| 55 | self.key = {} |
| 56 | |
| 57 | for chunk_idx, src_chunk_path in enumerate(self.src_chunk_paths): |
| 58 | src_chunk = torch.load(src_chunk_path) |
| 59 | depth_chunk = pickle.load(open(self._get_depth_path(src_chunk_path), "rb")) |
| 60 | assert [file['key'] for file in src_chunk] == depth_chunk['key'] |
| 61 | depth_chunk = depth_chunk['predicted_depth_norm'] |
| 62 | |
| 63 | # Store chunk data for later access |
| 64 | self.chunk_data[chunk_idx] = { |
| 65 | 'src_chunk': src_chunk, |
| 66 | 'depth_chunk': depth_chunk |
| 67 | } |
| 68 | self.set_len.append(len(src_chunk)) |
| 69 | # Create frame indices for this chunk |
| 70 | for i_vid, src_video_info in enumerate(src_chunk): |
| 71 | num_frames = len(src_video_info["images"]) |
| 72 | self.key[src_video_info["key"]] = (chunk_idx, i_vid) |