|
|
|
from typing import Optional |
|
|
|
import torch |
|
from torch import Tensor |
|
|
|
|
|
def yolov5_bbox_decoder(priors: Tensor, bbox_preds: Tensor, |
|
stride: Tensor) -> Tensor: |
|
bbox_preds = bbox_preds.sigmoid() |
|
|
|
x_center = (priors[..., 0] + priors[..., 2]) * 0.5 |
|
y_center = (priors[..., 1] + priors[..., 3]) * 0.5 |
|
w = priors[..., 2] - priors[..., 0] |
|
h = priors[..., 3] - priors[..., 1] |
|
|
|
x_center_pred = (bbox_preds[..., 0] - 0.5) * 2 * stride + x_center |
|
y_center_pred = (bbox_preds[..., 1] - 0.5) * 2 * stride + y_center |
|
w_pred = (bbox_preds[..., 2] * 2)**2 * w |
|
h_pred = (bbox_preds[..., 3] * 2)**2 * h |
|
|
|
decoded_bboxes = torch.stack( |
|
[x_center_pred, y_center_pred, w_pred, h_pred], dim=-1) |
|
|
|
return decoded_bboxes |
|
|
|
|
|
def rtmdet_bbox_decoder(priors: Tensor, bbox_preds: Tensor, |
|
stride: Optional[Tensor]) -> Tensor: |
|
stride = stride[None, :, None] |
|
bbox_preds *= stride |
|
tl_x = (priors[..., 0] - bbox_preds[..., 0]) |
|
tl_y = (priors[..., 1] - bbox_preds[..., 1]) |
|
br_x = (priors[..., 0] + bbox_preds[..., 2]) |
|
br_y = (priors[..., 1] + bbox_preds[..., 3]) |
|
decoded_bboxes = torch.stack([tl_x, tl_y, br_x, br_y], -1) |
|
return decoded_bboxes |
|
|
|
|
|
def yolox_bbox_decoder(priors: Tensor, bbox_preds: Tensor, |
|
stride: Optional[Tensor]) -> Tensor: |
|
stride = stride[None, :, None] |
|
xys = (bbox_preds[..., :2] * stride) + priors |
|
whs = bbox_preds[..., 2:].exp() * stride |
|
decoded_bboxes = torch.cat([xys, whs], -1) |
|
return decoded_bboxes |
|
|
