| 316 | |
| 317 | |
| 318 | class Stack(object): |
| 319 | |
| 320 | def __init__(self, roll=False): |
| 321 | self.roll = roll |
| 322 | |
| 323 | def __call__(self, img_dict): |
| 324 | img_group = img_dict['gt_frames'] |
| 325 | |
| 326 | mode = img_group[0].mode |
| 327 | if mode == '1': |
| 328 | img_group = [img.convert('L') for img in img_group] |
| 329 | mode = 'L' |
| 330 | |
| 331 | if mode == 'L': |
| 332 | img_group = np.stack([np.expand_dims(x, 2) for x in img_group], axis=2) |
| 333 | elif mode == 'RGB': |
| 334 | if self.roll: |
| 335 | img_group = np.stack([np.array(x)[:, :, ::-1] for x in img_group], axis=2) |
| 336 | else: |
| 337 | img_group = np.stack(img_group, axis=2) |
| 338 | else: |
| 339 | raise NotImplementedError(f"Image mode {mode}") |
| 340 | |
| 341 | img_dict['gt_frames'] = img_group |
| 342 | img_dict['flow_forward'] = np.stack(img_dict['flow_forward'], axis=2) |
| 343 | img_dict['flow_backward'] = np.stack(img_dict['flow_backward'], axis=2) |
| 344 | img_dict['flowmask_forward'] = np.stack([np.expand_dims(x, 2) for x in img_dict['flowmask_forward']], axis=2) |
| 345 | img_dict['flowmask_backward'] = np.stack([np.expand_dims(x, 2) for x in img_dict['flowmask_backward']], axis=2) |
| 346 | return img_dict |
| 347 | |
| 348 | |
| 349 | class ToTorchFormatTensor(object): |
no outgoing calls
no test coverage detected