Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| 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 | |