| import mmcv |
| import numpy as np |
|
|
| from mmcv.transforms import RandomFlip as MMCV_RandomFlip |
|
|
|
|
| from mmdet.registry import TRANSFORMS |
| from mmdet.structures.bbox import HorizontalBoxes, autocast_box_type |
|
|
|
|
| @TRANSFORMS.register_module() |
| class RandomFlipPiexlTarget(MMCV_RandomFlip): |
| """Flip the image & bbox & mask & segmentation map. Added or Updated keys: |
| flip, flip_direction, img, gt_bboxes, and gt_seg_map. There are 3 flip |
| modes: |
| |
| - ``prob`` is float, ``direction`` is string: the image will be |
| ``direction``ly flipped with probability of ``prob`` . |
| E.g., ``prob=0.5``, ``direction='horizontal'``, |
| then image will be horizontally flipped with probability of 0.5. |
| - ``prob`` is float, ``direction`` is list of string: the image will |
| be ``direction[i]``ly flipped with probability of |
| ``prob/len(direction)``. |
| E.g., ``prob=0.5``, ``direction=['horizontal', 'vertical']``, |
| then image will be horizontally flipped with probability of 0.25, |
| vertically with probability of 0.25. |
| - ``prob`` is list of float, ``direction`` is list of string: |
| given ``len(prob) == len(direction)``, the image will |
| be ``direction[i]``ly flipped with probability of ``prob[i]``. |
| E.g., ``prob=[0.3, 0.5]``, ``direction=['horizontal', |
| 'vertical']``, then image will be horizontally flipped with |
| probability of 0.3, vertically with probability of 0.5. |
| |
| |
| Required Keys: |
| |
| - img |
| - gt_bboxes (BaseBoxes[torch.float32]) (optional) |
| - gt_masks (BitmapMasks | PolygonMasks) (optional) |
| - gt_seg_map (np.uint8) (optional) |
| |
| Modified Keys: |
| |
| - img |
| - gt_bboxes |
| - gt_masks |
| - gt_seg_map |
| |
| Added Keys: |
| |
| - flip |
| - flip_direction |
| - homography_matrix |
| |
| |
| Args: |
| prob (float | list[float], optional): The flipping probability. |
| Defaults to None. |
| direction(str | list[str]): The flipping direction. Options |
| If input is a list, the length must equal ``prob``. Each |
| element in ``prob`` indicates the flip probability of |
| corresponding direction. Defaults to 'horizontal'. |
| """ |
|
|
| def _record_homography_matrix(self, results: dict) -> None: |
| """Record the homography matrix for the RandomFlip.""" |
| cur_dir = results['flip_direction'] |
| h, w = results['img'].shape[:2] |
|
|
| if cur_dir == 'horizontal': |
| homography_matrix = np.array([[-1, 0, w], [0, 1, 0], [0, 0, 1]], |
| dtype=np.float32) |
| elif cur_dir == 'vertical': |
| homography_matrix = np.array([[1, 0, 0], [0, -1, h], [0, 0, 1]], |
| dtype=np.float32) |
| elif cur_dir == 'diagonal': |
| homography_matrix = np.array([[-1, 0, w], [0, -1, h], [0, 0, 1]], |
| dtype=np.float32) |
| else: |
| homography_matrix = np.eye(3, dtype=np.float32) |
|
|
| if results.get('homography_matrix', None) is None: |
| results['homography_matrix'] = homography_matrix |
| else: |
| results['homography_matrix'] = homography_matrix @ results[ |
| 'homography_matrix'] |
|
|
| @autocast_box_type() |
| def _flip(self, results: dict) -> None: |
| """Flip images, bounding boxes, and semantic segmentation map.""" |
| |
| results['img'] = mmcv.imflip( |
| results['img'], direction=results['flip_direction']) |
|
|
| img_shape = results['img'].shape[:2] |
|
|
| |
| if results.get('gt_bboxes', None) is not None: |
| results['gt_bboxes'].flip_(img_shape, results['flip_direction']) |
|
|
| |
| if results.get('gt_masks', None) is not None: |
| results['gt_masks'] = results['gt_masks'].flip( |
| results['flip_direction']) |
|
|
|
|
| if results.get('gt_seg', None) is not None: |
| results['gt_seg'] = mmcv.imflip( |
| results['gt_seg'], direction=results['flip_direction']) |
| if results.get('gt_abu', None) is not None: |
| results['gt_abu'] = mmcv.imflip( |
| results['gt_abu'], direction=results['flip_direction']) |
| |
| self._record_homography_matrix(results) |
