| 22 | |
| 23 | |
| 24 | class SevenScenes(BaseStereoViewDataset): |
| 25 | def __init__( |
| 26 | self, |
| 27 | num_seq=1, |
| 28 | num_frames=5, |
| 29 | min_thresh=10, |
| 30 | max_thresh=100, |
| 31 | test_id=None, |
| 32 | full_video=False, |
| 33 | tuple_list=None, |
| 34 | seq_id=None, |
| 35 | rebuttal=False, |
| 36 | shuffle_seed=-1, |
| 37 | kf_every=1, |
| 38 | max_frames=None, |
| 39 | *args, |
| 40 | ROOT, |
| 41 | **kwargs, |
| 42 | ): |
| 43 | self.ROOT = ROOT |
| 44 | super().__init__(*args, **kwargs) |
| 45 | self.num_seq = num_seq |
| 46 | self.num_frames = num_frames |
| 47 | self.max_thresh = max_thresh |
| 48 | self.min_thresh = min_thresh |
| 49 | self.test_id = test_id |
| 50 | self.full_video = full_video |
| 51 | self.kf_every = kf_every |
| 52 | self.seq_id = seq_id |
| 53 | self.rebuttal = rebuttal |
| 54 | self.shuffle_seed = shuffle_seed |
| 55 | self.max_frames = max_frames |
| 56 | |
| 57 | # load all scenes |
| 58 | self.load_all_tuples(tuple_list) |
| 59 | self.load_all_scenes(ROOT) |
| 60 | |
| 61 | def __len__(self): |
| 62 | if self.tuple_list is not None: |
| 63 | return len(self.tuple_list) |
| 64 | return len(self.scene_list) * self.num_seq |
| 65 | |
| 66 | def load_all_tuples(self, tuple_list): |
| 67 | if tuple_list is not None: |
| 68 | self.tuple_list = tuple_list |
| 69 | # with open(tuple_path) as f: |
| 70 | # self.tuple_list = f.read().splitlines() |
| 71 | |
| 72 | else: |
| 73 | self.tuple_list = None |
| 74 | |
| 75 | def load_all_scenes(self, base_dir): |
| 76 | |
| 77 | if self.tuple_list is not None: |
| 78 | # Use pre-defined simplerecon scene_ids |
| 79 | self.scene_list = [ |
| 80 | "stairs/seq-06", |
| 81 | "stairs/seq-02", |