Randomly horizontally flips the given PIL.Image with a probability of 0.5
| 69 | |
| 70 | |
| 71 | class GroupRandomHorizontalFlip(object): |
| 72 | """Randomly horizontally flips the given PIL.Image with a probability of 0.5""" |
| 73 | def __init__(self, prob): |
| 74 | self.prob = prob |
| 75 | |
| 76 | def flip_flow(self, flow): |
| 77 | # flow: [H, W, 2] |
| 78 | flow = flow[:, ::-1, :] |
| 79 | flow[..., 0] = - flow[..., 0] |
| 80 | return flow |
| 81 | |
| 82 | def __call__(self, img_dict, is_flow=False): |
| 83 | v = random.random() |
| 84 | if v < self.prob: |
| 85 | img_dict['gt_frames'] = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_dict['gt_frames']] |
| 86 | img_dict['flow_forward'] = [self.flip_flow(img) for img in img_dict['flow_forward']] |
| 87 | img_dict['flow_backward'] = [self.flip_flow(img) for img in img_dict['flow_backward']] |
| 88 | img_dict['flowmask_forward'] = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_dict['flowmask_forward']] |
| 89 | img_dict['flowmask_backward'] = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_dict['flowmask_backward']] |
| 90 | |
| 91 | # flow_forward = img_dict['flow_forward'][0] |
| 92 | # flow_forward = computeColor(flow_forward[:,:,0], flow_forward[:,:,1]) |
| 93 | # flow_backward = img_dict['flow_backward'][0] |
| 94 | # flow_backward = computeColor(flow_backward[:,:,0], flow_backward[:,:,1]) |
| 95 | # gt_frames = img_dict['gt_frames'][0] |
| 96 | # flowmask_forward = img_dict['flowmask_forward'][0] |
| 97 | # flowmask_backward = img_dict['flowmask_backward'][0] |
| 98 | |
| 99 | # gt_frames.save('gt_frames.png') |
| 100 | # flowmask_forward.save('flowmask_forward.png') |
| 101 | # flowmask_backward.save('flowmask_backward.png') |
| 102 | # cv2.imwrite('flow_forward.png', flow_forward) |
| 103 | # cv2.imwrite('flow_backward.png', flow_backward) |
| 104 | # print('hi') |
| 105 | # exit() |
| 106 | |
| 107 | return img_dict |
| 108 | else: |
| 109 | return img_dict |
| 110 | |
| 111 | |
| 112 | class GroupNormalize(object): |