| 179 | |
| 180 | |
| 181 | class HD1K(FlowDataset): |
| 182 | def __init__(self, aug_params=None, root='datasets/HD1k'): |
| 183 | super(HD1K, self).__init__(aug_params, sparse=True) |
| 184 | |
| 185 | seq_ix = 0 |
| 186 | while 1: |
| 187 | flows = sorted(glob(os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix))) |
| 188 | images = sorted(glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix))) |
| 189 | |
| 190 | if len(flows) == 0: |
| 191 | break |
| 192 | |
| 193 | for i in range(len(flows)-1): |
| 194 | self.flow_list += [flows[i]] |
| 195 | self.image_list += [ [images[i], images[i+1]] ] |
| 196 | |
| 197 | seq_ix += 1 |
| 198 | |
| 199 | |
| 200 | def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'): |