| 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/'): |
| 265 | super(AutoFlow, self).__init__() |
| 266 | |
| 267 | root = 's3://' |
| 268 | self.image_list = [] |
| 269 | with open("./flow_dataset/AutoFlow/AutoFlow_image.txt") as f: |
| 270 | images = f.readlines() |
| 271 | for img1, img2 in zip(images[0::2], images[1::2]): |
| 272 | self.image_list.append([root+img1.strip(), root+img2.strip()]) |
| 273 | self.flow_list = [] |
| 274 | with open("./flow_dataset/AutoFlow/AutoFlow_flow.txt") as f: |
| 275 | flows = f.readlines() |
| 276 | for flow in flows: |
| 277 | self.flow_list.append(root+flow.strip()) |
| 278 | |
| 279 | self.crop_size = crop_size |
| 280 | self.log_dir = log_dir |
| 281 | self.num_steps = num_steps |
| 282 | self.scale = 1 |
| 283 | self.order = 1 |
| 284 | self.black = False |
| 285 | self.noise = 0 |
| 286 | self.is_test = False |
| 287 | self.init_seed = False |
| 288 | |
| 289 | self.iter_counts = 0 |
| 290 | |
| 291 | def __rmul__(self, v): |
| 292 | self.flow_list = v * self.flow_list |
| 293 | self.image_list = v * self.image_list |
| 294 | return self |
| 295 | |
| 296 | def __len__(self): |
| 297 | return len(self.image_list) * 100 |
| 298 | |
| 299 | def __getitem__(self, index): |
| 300 | #print(self.flow_list[index]) |
| 301 | if self.is_test: |
| 302 | img1 = frame_utils.read_gen(self.image_list[index][0], test=self.is_test) |
| 303 | img2 = frame_utils.read_gen(self.image_list[index][1], test=self.is_test) |
| 304 | img1 = np.array(img1).astype(np.uint8)[..., :3] |
| 305 | img2 = np.array(img2).astype(np.uint8)[..., :3] |
| 306 | img1 = torch.from_numpy(img1).permute(2, 0, 1).float() |
| 307 | img2 = torch.from_numpy(img2).permute(2, 0, 1).float() |
| 308 | return img1, img2, self.extra_info[index] |
| 309 | |
| 310 | if not self.init_seed: |
| 311 | worker_info = torch.utils.data.get_worker_info() |
| 312 | if worker_info is not None: |
| 313 | torch.manual_seed(worker_info.id) |
| 314 | np.random.seed(worker_info.id) |
| 315 | random.seed(worker_info.id) |
| 316 | self.init_seed = True |
| 317 | index = index % len(self.image_list) |
| 318 | valid = None |
| 319 | |
| 320 | flow = frame_utils.read_gen(self.flow_list[index]) |