82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
def _box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
|
|
"""
|
|
Converts bounding boxes from (cx, cy, w, h) format to (x1, y1, x2, y2) format.
|
|
(cx, cy) refers to center of bounding box
|
|
(w, h) are width and height of bounding box
|
|
Args:
|
|
boxes (Tensor[N, 4]): boxes in (cx, cy, w, h) format which will be converted.
|
|
|
|
Returns:
|
|
boxes (Tensor(N, 4)): boxes in (x1, y1, x2, y2) format.
|
|
"""
|
|
# We need to change all 4 of them so some temporary variable is needed.
|
|
cx, cy, w, h = boxes.unbind(-1)
|
|
x1 = cx - 0.5 * w
|
|
y1 = cy - 0.5 * h
|
|
x2 = cx + 0.5 * w
|
|
y2 = cy + 0.5 * h
|
|
|
|
boxes = torch.stack((x1, y1, x2, y2), dim=-1)
|
|
|
|
return boxes
|
|
|
|
|
|
def _box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor:
|
|
"""
|
|
Converts bounding boxes from (x1, y1, x2, y2) format to (cx, cy, w, h) format.
|
|
(x1, y1) refer to top left of bounding box
|
|
(x2, y2) refer to bottom right of bounding box
|
|
Args:
|
|
boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format which will be converted.
|
|
|
|
Returns:
|
|
boxes (Tensor(N, 4)): boxes in (cx, cy, w, h) format.
|
|
"""
|
|
x1, y1, x2, y2 = boxes.unbind(-1)
|
|
cx = (x1 + x2) / 2
|
|
cy = (y1 + y2) / 2
|
|
w = x2 - x1
|
|
h = y2 - y1
|
|
|
|
boxes = torch.stack((cx, cy, w, h), dim=-1)
|
|
|
|
return boxes
|
|
|
|
|
|
def _box_xywh_to_xyxy(boxes: Tensor) -> Tensor:
|
|
"""
|
|
Converts bounding boxes from (x, y, w, h) format to (x1, y1, x2, y2) format.
|
|
(x, y) refers to top left of bounding box.
|
|
(w, h) refers to width and height of box.
|
|
Args:
|
|
boxes (Tensor[N, 4]): boxes in (x, y, w, h) which will be converted.
|
|
|
|
Returns:
|
|
boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format.
|
|
"""
|
|
x, y, w, h = boxes.unbind(-1)
|
|
boxes = torch.stack([x, y, x + w, y + h], dim=-1)
|
|
return boxes
|
|
|
|
|
|
def _box_xyxy_to_xywh(boxes: Tensor) -> Tensor:
|
|
"""
|
|
Converts bounding boxes from (x1, y1, x2, y2) format to (x, y, w, h) format.
|
|
(x1, y1) refer to top left of bounding box
|
|
(x2, y2) refer to bottom right of bounding box
|
|
Args:
|
|
boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) which will be converted.
|
|
|
|
Returns:
|
|
boxes (Tensor[N, 4]): boxes in (x, y, w, h) format.
|
|
"""
|
|
x1, y1, x2, y2 = boxes.unbind(-1)
|
|
w = x2 - x1 # x2 - x1
|
|
h = y2 - y1 # y2 - y1
|
|
boxes = torch.stack((x1, y1, w, h), dim=-1)
|
|
return boxes
|