|
|
|
|
|
from typing import Optional, Sequence, Union |
|
|
|
|
|
from torch import Tensor |
|
|
|
|
|
from mmdet.registry import TASK_UTILS |
|
|
from mmdet.structures.bbox import (BaseBoxes, HorizontalBoxes, bbox2distance, |
|
|
distance2bbox, get_box_tensor) |
|
|
from .base_bbox_coder import BaseBBoxCoder |
|
|
|
|
|
|
|
|
@TASK_UTILS.register_module() |
|
|
class DistancePointBBoxCoder(BaseBBoxCoder): |
|
|
"""Distance Point BBox coder. |
|
|
|
|
|
This coder encodes gt bboxes (x1, y1, x2, y2) into (top, bottom, left, |
|
|
right) and decode it back to the original. |
|
|
|
|
|
Args: |
|
|
clip_border (bool, optional): Whether clip the objects outside the |
|
|
border of the image. Defaults to True. |
|
|
""" |
|
|
|
|
|
def __init__(self, clip_border: Optional[bool] = True, **kwargs) -> None: |
|
|
super().__init__(**kwargs) |
|
|
self.clip_border = clip_border |
|
|
|
|
|
def encode(self, |
|
|
points: Tensor, |
|
|
gt_bboxes: Union[Tensor, BaseBoxes], |
|
|
max_dis: Optional[float] = None, |
|
|
eps: float = 0.1) -> Tensor: |
|
|
"""Encode bounding box to distances. |
|
|
|
|
|
Args: |
|
|
points (Tensor): Shape (N, 2), The format is [x, y]. |
|
|
gt_bboxes (Tensor or :obj:`BaseBoxes`): Shape (N, 4), The format |
|
|
is "xyxy" |
|
|
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, 4). |
|
|
""" |
|
|
gt_bboxes = get_box_tensor(gt_bboxes) |
|
|
assert points.size(0) == gt_bboxes.size(0) |
|
|
assert points.size(-1) == 2 |
|
|
assert gt_bboxes.size(-1) == 4 |
|
|
return bbox2distance(points, gt_bboxes, max_dis, eps) |
|
|
|
|
|
def decode( |
|
|
self, |
|
|
points: Tensor, |
|
|
pred_bboxes: Tensor, |
|
|
max_shape: Optional[Union[Sequence[int], Tensor, |
|
|
Sequence[Sequence[int]]]] = None |
|
|
) -> Union[Tensor, BaseBoxes]: |
|
|
"""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 (left, top, right, bottom). Shape (B, N, 4) |
|
|
or (N, 4) |
|
|
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: |
|
|
Union[Tensor, :obj:`BaseBoxes`]: Boxes with shape (N, 4) or |
|
|
(B, N, 4) |
|
|
""" |
|
|
assert points.size(0) == pred_bboxes.size(0) |
|
|
assert points.size(-1) == 2 |
|
|
assert pred_bboxes.size(-1) == 4 |
|
|
if self.clip_border is False: |
|
|
max_shape = None |
|
|
bboxes = distance2bbox(points, pred_bboxes, max_shape) |
|
|
|
|
|
if self.use_box_type: |
|
|
bboxes = HorizontalBoxes(bboxes) |
|
|
return bboxes |
|
|
|
