Spaces:
Running on Zero
Running on Zero
File size: 2,636 Bytes
d6ea92f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from __future__ import annotations
import torch
from torch import Tensor
def box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
cx, cy, width, height = boxes.unbind(-1)
return torch.stack(
(cx - 0.5 * width, cy - 0.5 * height, cx + 0.5 * width, cy + 0.5 * height),
dim=-1,
)
def box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor:
x0, y0, x1, y1 = boxes.unbind(-1)
return torch.stack(
((x0 + x1) * 0.5, (y0 + y1) * 0.5, x1 - x0, y1 - y0), dim=-1
)
def box_area(boxes: Tensor) -> Tensor:
return (boxes[..., 2] - boxes[..., 0]).clamp(min=0) * (
boxes[..., 3] - boxes[..., 1]
).clamp(min=0)
def box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
area1 = box_area(boxes1)
area2 = box_area(boxes2)
top_left = torch.maximum(boxes1[:, None, :2], boxes2[:, :2])
bottom_right = torch.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
union = area1[:, None] + area2 - intersection
return intersection / union.clamp(min=1e-7), union
def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""Pairwise generalized IoU for boxes in x0, y0, x1, y1 format."""
iou, union = box_iou(boxes1, boxes2)
top_left = torch.minimum(boxes1[:, None, :2], boxes2[:, :2])
bottom_right = torch.maximum(boxes1[:, None, 2:], boxes2[:, 2:])
enclosing = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
return iou - (enclosing - union) / enclosing.clamp(min=1e-7)
def generalized_box_iou_batched(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""Batched pairwise generalized IoU, boxes in x0, y0, x1, y1 format.
boxes1: [B, N, 4], boxes2: [B, M, 4] -> [B, N, M]
"""
area1 = box_area(boxes1)
area2 = box_area(boxes2)
top_left = torch.maximum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
bottom_right = torch.minimum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
union = area1[:, :, None] + area2[:, None, :] - intersection
iou = intersection / union.clamp(min=1e-7)
enc_top_left = torch.minimum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
enc_bottom_right = torch.maximum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
enclosing = (enc_bottom_right - enc_top_left).clamp(min=0).prod(dim=-1)
return iou - (enclosing - union) / enclosing.clamp(min=1e-7)
def inverse_sigmoid(value: Tensor, eps: float = 1e-5) -> Tensor:
value = value.clamp(min=0.0, max=1.0)
return torch.log(value.clamp(min=eps) / (1.0 - value).clamp(min=eps))
|