Spaces:
Runtime error
Runtime error
File size: 3,512 Bytes
3094730 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Sequence, Union
import torch
from mmyolo.registry import TASK_UTILS
try:
from mmrotate.models.task_modules.coders import \
DistanceAnglePointCoder as MMROTATE_DistanceAnglePointCoder
MMROTATE_AVAILABLE = True
except ImportError:
from mmdet.models.task_modules.coders import BaseBBoxCoder
MMROTATE_DistanceAnglePointCoder = BaseBBoxCoder
MMROTATE_AVAILABLE = False
@TASK_UTILS.register_module()
class DistanceAnglePointCoder(MMROTATE_DistanceAnglePointCoder):
"""Distance Angle Point BBox coder.
This coder encodes gt bboxes (x, y, w, h, theta) into (top, bottom, left,
right, theta) and decode it back to the original.
"""
def __init__(self, clip_border=True, angle_version='oc'):
if not MMROTATE_AVAILABLE:
raise ImportError(
'Please run "mim install -r requirements/mmrotate.txt" '
'to install mmrotate first for rotated detection.')
super().__init__(clip_border=clip_border, angle_version=angle_version)
def decode(
self,
points: torch.Tensor,
pred_bboxes: torch.Tensor,
stride: torch.Tensor,
max_shape: Optional[Union[Sequence[int], torch.Tensor,
Sequence[Sequence[int]]]] = None,
) -> torch.Tensor:
"""Decode distance prediction to bounding box.
Args:
points (Tensor): Shape (B, N, 2) or (N, 2).
pred_bboxes (Tensor): Distance from the given point to 4
boundaries and angle (left, top, right, bottom, angle).
Shape (B, N, 5) or (N, 5)
max_shape (Sequence[int] or torch.Tensor or Sequence[
Sequence[int]],optional): Maximum bounds for boxes, specifies
(H, W, C) or (H, W). If priors shape is (B, N, 4), then
the max_shape should be a Sequence[Sequence[int]],
and the length of max_shape should also be B.
Default None.
Returns:
Tensor: Boxes with shape (N, 5) or (B, N, 5)
"""
assert points.size(-2) == pred_bboxes.size(-2)
assert points.size(-1) == 2
assert pred_bboxes.size(-1) == 5
if self.clip_border is False:
max_shape = None
if pred_bboxes.dim() == 2:
stride = stride[:, None]
else:
stride = stride[None, :, None]
pred_bboxes[..., :4] = pred_bboxes[..., :4] * stride
return self.distance2obb(points, pred_bboxes, max_shape,
self.angle_version)
def encode(self,
points: torch.Tensor,
gt_bboxes: torch.Tensor,
max_dis: float = 16.,
eps: float = 0.01) -> torch.Tensor:
"""Encode bounding box to distances.
Args:
points (Tensor): Shape (N, 2), The format is [x, y].
gt_bboxes (Tensor): Shape (N, 5), The format is "xywha"
max_dis (float): Upper bound of the distance. Default None.
eps (float): a small value to ensure target < max_dis, instead <=.
Default 0.1.
Returns:
Tensor: Box transformation deltas. The shape is (N, 5).
"""
assert points.size(-2) == gt_bboxes.size(-2)
assert points.size(-1) == 2
assert gt_bboxes.size(-1) == 5
return self.obb2distance(points, gt_bboxes, max_dis, eps)
|