box_ops.py

# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
# Copyright (c) 2024 Baidu. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Conditional DETR
# Copyright (c) 2021 Microsoft. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Copied from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
# ------------------------------------------------------------------------
"""Utilities for bounding box manipulation and GIoU."""

from __future__ import annotations

import sys

import torch
import torch.nn.functional as F  # noqa: N812
from torch import Tensor
from torchvision.ops.boxes import box_area


def box_cxcywh_to_xyxy(x: Tensor) -> Tensor:
    x_c, y_c, w, h = x.unbind(-1)
    b = [
        (x_c - 0.5 * w.clamp(min=0.0)),
        (y_c - 0.5 * h.clamp(min=0.0)),
        (x_c + 0.5 * w.clamp(min=0.0)),
        (y_c + 0.5 * h.clamp(min=0.0)),
    ]
    return torch.stack(b, dim=-1)


def box_xyxy_to_cxcywh(x: Tensor) -> Tensor:
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


# modified from torchvision to also return the union
def box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
    """Compute pairwise IoU and union for two sets of boxes.

    Returns:
        iou: the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2
        union: the NxM matrix containing the pairwise
            union values for every element in boxes1 and boxes2
    """
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    eps = 1e-7
    # Clamp only the degenerate (union==0) case so identical non-degenerate boxes
    # yield IoU==1.0 exactly; adding eps unconditionally would break that identity.
    iou = inter / union.clamp(min=eps)
    return iou, union


def _assert_equal_length(boxes1: Tensor, boxes2: Tensor) -> None:
    """Reject unequal-length operands so a length-1 side cannot silently broadcast."""
    if boxes1.shape[0] != boxes2.shape[0]:
        raise ValueError(
            "elementwise box ops expect boxes1 and boxes2 to have the same length, "
            f"got {boxes1.shape[0]} and {boxes2.shape[0]}"
        )


def elementwise_box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
    """Compute IoU and union for pre-matched box pairs.

    Unlike ``box_iou``, this avoids materializing the full NxN pairwise matrix
    (of which only the diagonal is used for matched pairs), so peak memory is
    O(N) instead of O(N^2). ``boxes1`` and ``boxes2`` must have the same length
    and be in [x0, y0, x1, y1] format; a length mismatch raises ``ValueError``
    rather than silently broadcasting a length-1 operand.

    The numerics (op order and the ``eps=1e-7`` union clamp) are identical to the
    diagonal of ``box_iou``, so results carry no dtype dependence beyond that of
    the pairwise version under FP16/FP32.

    Returns:
        iou: the [N] tensor of IoU values for each matched pair.
        union: the [N] tensor of union areas for each matched pair.
    """
    _assert_equal_length(boxes1, boxes2)
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, :2], boxes2[:, :2])  # [N,2]
    rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])  # [N,2]

    wh = (rb - lt).clamp(min=0)  # [N,2]
    inter = wh[:, 0] * wh[:, 1]  # [N]

    union = area1 + area2 - inter

    eps = 1e-7
    # Clamp only the degenerate (union==0) case so identical non-degenerate boxes
    # yield IoU==1.0 exactly; adding eps unconditionally would break that identity.
    iou = inter / union.clamp(min=eps)
    return iou, union


def elementwise_generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """Generalized IoU from https://giou.stanford.edu/ for pre-matched box pairs.

    Equivalent to the diagonal of ``generalized_box_iou`` but without building the
    NxN matrix, giving O(N) instead of O(N^2) peak memory. The boxes should be in
    [x0, y0, x1, y1] format, and ``boxes1``/``boxes2`` must have the same length; a
    length mismatch raises ``ValueError`` rather than silently broadcasting.

    Returns a [N] tensor, one GIoU value per matched pair.
    """
    _assert_equal_length(boxes1, boxes2)
    # Degenerate (zero-area) boxes would divide 0/0; eps in the denominators keeps results finite.
    iou, union = elementwise_box_iou(boxes1, boxes2)

    lt = torch.min(boxes1[:, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,2]
    area = wh[:, 0] * wh[:, 1]

    eps = 1e-7
    # Clamp only when enclosing area is zero (degenerate enclosing box) so normal boxes remain exact.
    return iou - (area - union) / area.clamp(min=eps)


def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """Generalized IoU from https://giou.stanford.edu/

    The boxes should be in [x0, y0, x1, y1] format.

    Returns a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2).
    """
    # Degenerate (zero-area) boxes would divide 0/0; eps in the denominators keeps results finite.
    iou, union = box_iou(boxes1, boxes2)

    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    area = wh[:, :, 0] * wh[:, :, 1]

    eps = 1e-7
    # Clamp only when enclosing area is zero (degenerate enclosing box) so normal boxes remain exact.
    return iou - (area - union) / area.clamp(min=eps)


def masks_to_boxes(masks: Tensor) -> Tensor:
    """Compute the bounding boxes around the provided masks.

    The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.

    Returns a [N, 4] tensor, with the boxes in xyxy format.
    """
    if masks.numel() == 0:
        return torch.zeros((0, 4), device=masks.device)

    h, w = masks.shape[-2:]

    y = torch.arange(0, h, dtype=torch.float32, device=masks.device)
    x = torch.arange(0, w, dtype=torch.float32, device=masks.device)
    y, x = torch.meshgrid(y, x, indexing="ij")

    x_mask = masks * x.unsqueeze(0)
    x_max = x_mask.flatten(1).max(-1)[0]
    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    y_mask = masks * y.unsqueeze(0)
    y_max = y_mask.flatten(1).max(-1)[0]
    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]

    boxes = torch.stack([x_min, y_min, x_max, y_max], 1)

    keep = masks.flatten(1).any(-1)
    boxes[~keep] = boxes.new_zeros(4)
    return boxes


def batch_dice_loss(inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute the DICE loss, similar to generalized IOU for masks.

    Args:
        inputs: A float tensor of arbitrary shape. The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
            classification label for each element in inputs (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
    denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
    loss: Tensor = 1 - (numerator + 1) / (denominator + 1)
    return loss


batch_dice_loss_jit = batch_dice_loss if sys.version_info >= (3, 14) else torch.jit.script(batch_dice_loss)


def batch_sigmoid_ce_loss(inputs: Tensor, targets: Tensor) -> Tensor:
    """Compute sigmoid cross-entropy loss for mask predictions.

    Args:
        inputs: A float tensor of arbitrary shape. The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
            classification label for each element in inputs (0 for the negative class and 1 for the positive class).

    Returns:
        Loss tensor.
    """
    hw = inputs.shape[1]

    pos = F.binary_cross_entropy_with_logits(inputs, torch.ones_like(inputs), reduction="none")
    neg = F.binary_cross_entropy_with_logits(inputs, torch.zeros_like(inputs), reduction="none")

    loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum("nc,mc->nm", neg, (1 - targets))

    return loss / hw


batch_sigmoid_ce_loss_jit = (
    batch_sigmoid_ce_loss if sys.version_info >= (3, 14) else torch.jit.script(batch_sigmoid_ce_loss)
)