Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from typing import Optional, Sequence, Union | |
| import torch | |
| from mmdet.models.task_modules.coders import \ | |
| DistancePointBBoxCoder as MMDET_DistancePointBBoxCoder | |
| from mmdet.structures.bbox import bbox2distance, distance2bbox | |
| from mmyolo.registry import TASK_UTILS | |
| class DistancePointBBoxCoder(MMDET_DistancePointBBoxCoder): | |
| """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. | |
| """ | |
| 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 (left, top, right, bottom). Shape (B, N, 4) | |
| or (N, 4) | |
| stride (Tensor): Featmap stride. | |
| 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, 4) or (B, N, 4) | |
| """ | |
| assert points.size(-2) == pred_bboxes.size(-2) | |
| assert points.size(-1) == 2 | |
| assert pred_bboxes.size(-1) == 4 | |
| if self.clip_border is False: | |
| max_shape = None | |
| pred_bboxes = pred_bboxes * stride[None, :, None] | |
| return distance2bbox(points, pred_bboxes, max_shape) | |
| 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. The rewrite is to support batch | |
| operations. | |
| Args: | |
| points (Tensor): Shape (B, N, 2) or (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 to 16.. | |
| eps (float): a small value to ensure target < max_dis, instead <=. | |
| Default 0.01. | |
| Returns: | |
| Tensor: Box transformation deltas. The shape is (N, 4) or | |
| (B, N, 4). | |
| """ | |
| assert points.size(-2) == gt_bboxes.size(-2) | |
| assert points.size(-1) == 2 | |
| assert gt_bboxes.size(-1) == 4 | |
| return bbox2distance(points, gt_bboxes, max_dis, eps) | |