| |
| from typing import Union |
|
|
| import torch |
| from torch import Tensor |
|
|
| from mmdet.registry import TASK_UTILS |
| from mmdet.structures.bbox import BaseBoxes, HorizontalBoxes, get_box_tensor |
| from .base_bbox_coder import BaseBBoxCoder |
|
|
|
|
| @TASK_UTILS.register_module() |
| class YOLOBBoxCoder(BaseBBoxCoder): |
| """YOLO BBox coder. |
| |
| Following `YOLO <https://arxiv.org/abs/1506.02640>`_, this coder divide |
| image into grids, and encode bbox (x1, y1, x2, y2) into (cx, cy, dw, dh). |
| cx, cy in [0., 1.], denotes relative center position w.r.t the center of |
| bboxes. dw, dh are the same as :obj:`DeltaXYWHBBoxCoder`. |
| |
| Args: |
| eps (float): Min value of cx, cy when encoding. |
| """ |
|
|
| def __init__(self, eps: float = 1e-6, **kwargs): |
| super().__init__(**kwargs) |
| self.eps = eps |
|
|
| def encode(self, bboxes: Union[Tensor, BaseBoxes], |
| gt_bboxes: Union[Tensor, BaseBoxes], |
| stride: Union[Tensor, int]) -> Tensor: |
| """Get box regression transformation deltas that can be used to |
| transform the ``bboxes`` into the ``gt_bboxes``. |
| |
| Args: |
| bboxes (torch.Tensor or :obj:`BaseBoxes`): Source boxes, |
| e.g., anchors. |
| gt_bboxes (torch.Tensor or :obj:`BaseBoxes`): Target of the |
| transformation, e.g., ground-truth boxes. |
| stride (torch.Tensor | int): Stride of bboxes. |
| |
| Returns: |
| torch.Tensor: Box transformation deltas |
| """ |
| bboxes = get_box_tensor(bboxes) |
| gt_bboxes = get_box_tensor(gt_bboxes) |
| assert bboxes.size(0) == gt_bboxes.size(0) |
| assert bboxes.size(-1) == gt_bboxes.size(-1) == 4 |
| x_center_gt = (gt_bboxes[..., 0] + gt_bboxes[..., 2]) * 0.5 |
| y_center_gt = (gt_bboxes[..., 1] + gt_bboxes[..., 3]) * 0.5 |
| w_gt = gt_bboxes[..., 2] - gt_bboxes[..., 0] |
| h_gt = gt_bboxes[..., 3] - gt_bboxes[..., 1] |
| x_center = (bboxes[..., 0] + bboxes[..., 2]) * 0.5 |
| y_center = (bboxes[..., 1] + bboxes[..., 3]) * 0.5 |
| w = bboxes[..., 2] - bboxes[..., 0] |
| h = bboxes[..., 3] - bboxes[..., 1] |
| w_target = torch.log((w_gt / w).clamp(min=self.eps)) |
| h_target = torch.log((h_gt / h).clamp(min=self.eps)) |
| x_center_target = ((x_center_gt - x_center) / stride + 0.5).clamp( |
| self.eps, 1 - self.eps) |
| y_center_target = ((y_center_gt - y_center) / stride + 0.5).clamp( |
| self.eps, 1 - self.eps) |
| encoded_bboxes = torch.stack( |
| [x_center_target, y_center_target, w_target, h_target], dim=-1) |
| return encoded_bboxes |
|
|
| def decode(self, bboxes: Union[Tensor, BaseBoxes], pred_bboxes: Tensor, |
| stride: Union[Tensor, int]) -> Union[Tensor, BaseBoxes]: |
| """Apply transformation `pred_bboxes` to `boxes`. |
| |
| Args: |
| boxes (torch.Tensor or :obj:`BaseBoxes`): Basic boxes, |
| e.g. anchors. |
| pred_bboxes (torch.Tensor): Encoded boxes with shape |
| stride (torch.Tensor | int): Strides of bboxes. |
| |
| Returns: |
| Union[torch.Tensor, :obj:`BaseBoxes`]: Decoded boxes. |
| """ |
| bboxes = get_box_tensor(bboxes) |
| assert pred_bboxes.size(-1) == bboxes.size(-1) == 4 |
| xy_centers = (bboxes[..., :2] + bboxes[..., 2:]) * 0.5 + ( |
| pred_bboxes[..., :2] - 0.5) * stride |
| whs = (bboxes[..., 2:] - |
| bboxes[..., :2]) * 0.5 * pred_bboxes[..., 2:].exp() |
| decoded_bboxes = torch.stack( |
| (xy_centers[..., 0] - whs[..., 0], xy_centers[..., 1] - |
| whs[..., 1], xy_centers[..., 0] + whs[..., 0], |
| xy_centers[..., 1] + whs[..., 1]), |
| dim=-1) |
|
|
| if self.use_box_type: |
| decoded_bboxes = HorizontalBoxes(decoded_bboxes) |
| return decoded_bboxes |
|
|
