| 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 |