Spaces:
Runtime error
Runtime error
File size: 2,360 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
# Copyright (c) OpenMMLab. All rights reserved.
import mmcv
import numpy as np
from mmcv.transforms import BaseTransform
from mmdet.registry import TRANSFORMS
from mmdet.structures.bbox import HorizontalBoxes, autocast_box_type
from .transforms import RandomFlip
@TRANSFORMS.register_module()
class GTBoxSubOne_GLIP(BaseTransform):
"""Subtract 1 from the x2 and y2 coordinates of the gt_bboxes."""
def transform(self, results: dict) -> dict:
if 'gt_bboxes' in results:
gt_bboxes = results['gt_bboxes']
if isinstance(gt_bboxes, np.ndarray):
gt_bboxes[:, 2:] -= 1
results['gt_bboxes'] = gt_bboxes
elif isinstance(gt_bboxes, HorizontalBoxes):
gt_bboxes = results['gt_bboxes'].tensor
gt_bboxes[:, 2:] -= 1
results['gt_bboxes'] = HorizontalBoxes(gt_bboxes)
else:
raise NotImplementedError
return results
@TRANSFORMS.register_module()
class RandomFlip_GLIP(RandomFlip):
"""Flip the image & bboxes & masks & segs horizontally or vertically.
When using horizontal flipping, the corresponding bbox x-coordinate needs
to be additionally subtracted by one.
"""
@autocast_box_type()
def _flip(self, results: dict) -> None:
"""Flip images, bounding boxes, and semantic segmentation map."""
# flip image
results['img'] = mmcv.imflip(
results['img'], direction=results['flip_direction'])
img_shape = results['img'].shape[:2]
# flip bboxes
if results.get('gt_bboxes', None) is not None:
results['gt_bboxes'].flip_(img_shape, results['flip_direction'])
# Only change this line
if results['flip_direction'] == 'horizontal':
results['gt_bboxes'].translate_([-1, 0])
# TODO: check it
# flip masks
if results.get('gt_masks', None) is not None:
results['gt_masks'] = results['gt_masks'].flip(
results['flip_direction'])
# flip segs
if results.get('gt_seg_map', None) is not None:
results['gt_seg_map'] = mmcv.imflip(
results['gt_seg_map'], direction=results['flip_direction'])
# record homography matrix for flip
self._record_homography_matrix(results)
|