(self, index)
| 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]) |
| 321 | |
| 322 | img1 = frame_utils.read_gen(self.image_list[index][0]) |
| 323 | img2 = frame_utils.read_gen(self.image_list[index][1]) |
| 324 | |
| 325 | flow = np.array(flow).astype(np.float32) |
| 326 | # For PWC-style augmentation, pixel values are in [0, 1] |
| 327 | img1 = np.array(img1).astype(np.uint8) / 255.0 |
| 328 | img2 = np.array(img2).astype(np.uint8) / 255.0 |
| 329 | |
| 330 | # grayscale images |
| 331 | if len(img1.shape) == 2: |
| 332 | img1 = np.tile(img1[...,None], (1, 1, 3)) |
| 333 | img2 = np.tile(img2[...,None], (1, 1, 3)) |
| 334 | else: |
| 335 | img1 = img1[..., :3] |
| 336 | img2 = img2[..., :3] |
| 337 | |
| 338 | iter_counts = self.iter_counts |
| 339 | self.iter_counts = self.iter_counts + 1 |
| 340 | print(self.iter_counts) |
| 341 | th, tw = self.crop_size |
| 342 | schedule = [0.5, 1., self.num_steps] # initial coeff, final_coeff, half life |
| 343 | schedule_coeff = schedule[0] + (schedule[1] - schedule[0]) * \ |
| 344 | (2/(1+np.exp(-1.0986*iter_counts/schedule[2])) - 1) |
| 345 | |
| 346 | co_transform = flow_transforms.Compose([ |
| 347 | flow_transforms.Scale(self.scale, order=self.order), |
| 348 | flow_transforms.SpatialAug([th,tw],scale=[0.4,0.03,0.2], |
| 349 | rot=[0.4,0.03], |
| 350 | trans=[0.4,0.03], |
| 351 | squeeze=[0.3,0.], schedule_coeff=schedule_coeff, order=self.order, black=self.black), |
| 352 | flow_transforms.PCAAug(schedule_coeff=schedule_coeff), |
| 353 | flow_transforms.ChromaticAug( schedule_coeff=schedule_coeff, noise=self.noise), |
| 354 | ]) |
| 355 | |
| 356 | flow = np.concatenate([flow, np.ones((flow.shape[0], flow.shape[1], 1))], axis=-1) |
nothing calls this directly
no outgoing calls
no test coverage detected