| 223 | |
| 224 | |
| 225 | class KITTI(FlowDataset): |
| 226 | def __init__(self, aug_params=None, split='training', root='datasets/KITTI'): |
| 227 | super(KITTI, self).__init__(aug_params, sparse=True) |
| 228 | if split == 'testing': |
| 229 | self.is_test = True |
| 230 | |
| 231 | root = 's3://' |
| 232 | |
| 233 | self.image_list = [] |
| 234 | with open("./flow_dataset/KITTI/KITTI_{}_image.txt".format(split)) as f: |
| 235 | images = f.readlines() |
| 236 | for img1, img2 in zip(images[0::2], images[1::2]): |
| 237 | self.image_list.append([root+img1.strip(), root+img2.strip()]) |
| 238 | |
| 239 | self.extra_info = [] |
| 240 | with open("./flow_dataset/KITTI/KITTI_{}_extra_info.txt".format(split)) as f: |
| 241 | info = f.readlines() |
| 242 | for id in info: |
| 243 | self.extra_info.append([id.strip()]) |
| 244 | |
| 245 | if split == "training": |
| 246 | self.flow_list = [] |
| 247 | with open("./flow_dataset/KITTI/KITTI_{}_flow.txt".format(split)) as f: |
| 248 | flow = f.readlines() |
| 249 | for flo in flow: |
| 250 | self.flow_list.append(root+flo.strip()) |
| 251 | # root = osp.join(root, split) |
| 252 | # images1 = sorted(glob(osp.join(root, 'image_2/*_10.png'))) |
| 253 | # images2 = sorted(glob(osp.join(root, 'image_2/*_11.png'))) |
| 254 | |
| 255 | # for img1, img2 in zip(images1, images2): |
| 256 | # frame_id = img1.split('/')[-1] |
| 257 | # self.extra_info += [ [frame_id] ] |
| 258 | # self.image_list += [ [img1, img2] ] |
| 259 | |
| 260 | # if split == 'training': |
| 261 | # self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png'))) |
| 262 | |
| 263 | class AutoFlow(data.Dataset): |
| 264 | def __init__(self, num_steps, crop_size, log_dir, root='datasets/'): |