| 110 | return target_pose.reshape(poses.shape[0],12) |
| 111 | |
| 112 | class Cambridge2(data.Dataset): |
| 113 | def __init__(self, scene, data_path, train, transform=None, |
| 114 | target_transform=None, mode=0, seed=7, |
| 115 | skip_images=False, df=2., trainskip=1, testskip=1, hwf=[480,854,744.], |
| 116 | ret_idx=False, fix_idx=False, ret_hist=False, hist_bin=10): |
| 117 | """ |
| 118 | :param scene: scene name ['chess', 'pumpkin', ...] |
| 119 | :param data_path: root 7scenes data directory. |
| 120 | Usually '../data/deepslam_data/7Scenes' |
| 121 | :param train: if True, return the training images. If False, returns the |
| 122 | testing images |
| 123 | :param transform: transform to apply to the images |
| 124 | :param target_transform: transform to apply to the poses |
| 125 | :param mode: (Obsolete) 0: just color image, 1: color image in NeRF 0-1 and resized |
| 126 | :param skip_images: If True, skip loading images and return None instead |
| 127 | :param df: downscale factor |
| 128 | :param trainskip: due to 7scenes are so big, now can use less training sets # of trainset = 1/trainskip |
| 129 | :param testskip: skip part of testset, # of testset = 1/testskip |
| 130 | :param hwf: H,W,Focal from COLMAP |
| 131 | """ |
| 132 | |
| 133 | self.transform = transform |
| 134 | self.target_transform = target_transform |
| 135 | self.df = df |
| 136 | |
| 137 | self.H, self.W, self.focal = hwf |
| 138 | self.H = int(self.H) |
| 139 | self.W = int(self.W) |
| 140 | np.random.seed(seed) |
| 141 | |
| 142 | self.train = train |
| 143 | self.ret_idx = ret_idx |
| 144 | self.fix_idx = fix_idx |
| 145 | self.ret_hist = ret_hist |
| 146 | self.hist_bin = hist_bin # histogram bin size |
| 147 | |
| 148 | if self.train: |
| 149 | root_dir = osp.join(data_path, scene) + '/train' |
| 150 | else: |
| 151 | root_dir = osp.join(data_path, scene) + '/test' |
| 152 | |
| 153 | rgb_dir = root_dir + '/rgb/' |
| 154 | |
| 155 | pose_dir = root_dir + '/poses/' |
| 156 | |
| 157 | world_setup_fn = osp.join(data_path, scene) + '/world_setup.json' |
| 158 | |
| 159 | # collect poses and image names |
| 160 | self.rgb_files = os.listdir(rgb_dir) |
| 161 | self.rgb_files = [rgb_dir + f for f in self.rgb_files] |
| 162 | self.rgb_files.sort() |
| 163 | |
| 164 | self.pose_files = os.listdir(pose_dir) |
| 165 | self.pose_files = [pose_dir + f for f in self.pose_files] |
| 166 | self.pose_files.sort() |
| 167 | |
| 168 | # remove some abnormal data, need to fix later |
| 169 | if scene == 'ShopFacade' and self.train: |
no outgoing calls
no test coverage detected