Operations API Reference

oriented_det.ops

Ops module exposing oriented IoU and NMS helpers.

batch_rbox_iou(boxes_a, boxes_b, *, device=None, intersection_backend='auto', use_aabb_prefilter=True)

Produce an IoU matrix for two RBox collections.

This is a CPU Python implementation. For GPU-accelerated batch IoU, use: from oriented_det.ops.gpu_ops import oriented_box_iou_gpu

Parameters:

Name	Type	Description	Default
boxes_a	Sequence[RBox | Sequence[float]]	First collection of RBoxes	required
boxes_b	Sequence[RBox | Sequence[float]]	Second collection of RBoxes	required
device	Optional['torch.device']	Unused, kept for API compatibility	None
intersection_backend	str	Backend for intersection calculation ('auto', 'python', or 'shapely')	'auto'
use_aabb_prefilter	bool	If True (default), use fast AABB pre-filtering to skip expensive polygon intersections for non-overlapping boxes.	True

Returns:

Type	Description
list[list[float]]	IoU matrix as a list of lists [len(boxes_a), len(boxes_b)]

Source code in oriented_det/ops/iou.py

def batch_rbox_iou(
    boxes_a: Sequence[RBox | Sequence[float]],
    boxes_b: Sequence[RBox | Sequence[float]],
    *,
    device: Optional["torch.device"] = None,
    intersection_backend: str = "auto",
    use_aabb_prefilter: bool = True,
) -> list[list[float]]:
    """Produce an IoU matrix for two RBox collections.

    This is a CPU Python implementation. For GPU-accelerated batch IoU, use:
    `from oriented_det.ops.gpu_ops import oriented_box_iou_gpu`

    Args:
        boxes_a: First collection of RBoxes
        boxes_b: Second collection of RBoxes
        device: Unused, kept for API compatibility
        intersection_backend: Backend for intersection calculation 
            ('auto', 'python', or 'shapely')
        use_aabb_prefilter: If True (default), use fast AABB pre-filtering to skip
            expensive polygon intersections for non-overlapping boxes.

    Returns:
        IoU matrix as a list of lists [len(boxes_a), len(boxes_b)]
    """
    # Convert to RBoxes for AABB computation
    rboxes_a = [as_rbox(rb) for rb in boxes_a]
    rboxes_b = [as_rbox(rb) for rb in boxes_b]

    # Pre-compute AABBs for fast pre-filtering
    if use_aabb_prefilter:
        aabbs_a = [rb.axis_aligned_bounds() for rb in rboxes_a]
        aabbs_b = [rb.axis_aligned_bounds() for rb in rboxes_b]

    polys_a = [as_polygon(rb) for rb in rboxes_a]
    polys_b = [as_polygon(rb) for rb in rboxes_b]
    areas_b = [poly.area for poly in polys_b]
    matrix: list[list[float]] = []
    for i, poly_a in enumerate(polys_a):
        row = []
        area_a = poly_a.area
        aabb_a = aabbs_a[i] if use_aabb_prefilter else None

        for j, (poly_b, area_b) in enumerate(zip(polys_b, areas_b)):
            # Fast AABB pre-filter: if axis-aligned boxes don't overlap,
            # rotated boxes cannot overlap either (avoids expensive polygon intersection)
            if use_aabb_prefilter and not _aabb_overlaps(aabb_a, aabbs_b[j]):
                row.append(0.0)
                continue

            inter = polygon_intersection_area(poly_a, poly_b, backend=intersection_backend)
            if inter == 0.0:
                row.append(0.0)
            else:
                union = area_a + area_b - inter
                row.append(inter / union if union > 0 else 0.0)
        matrix.append(row)
    return matrix

generate_oriented_anchors_gpu(image_size, feature_map_sizes, anchor_scales, anchor_ratios, anchor_angles, stride_per_level, device=torch.device('cpu'), octave_base_scale=None, scales_per_octave=None)

Generate oriented anchors using vectorized GPU operations.

Fully vectorized, no Python loops for anchor generation.

Parameters:

Name	Type	Description	Default
image_size	Tuple[int, int]	(height, width) of input image	required
feature_map_sizes	List[Tuple[int, int]]	List of (height, width) for each FPN level	required
anchor_scales	List[float]	List of anchor scales (one per level or shared)	required
anchor_ratios	List[float]	List of aspect ratios	required
anchor_angles	List[float]	List of anchor angles in radians	required
stride_per_level	List[int]	List of strides for each FPN level	required
device	device	Device to create tensors on	device('cpu')

Returns:

Type	Description
List[Tensor]	List of anchor tensors, each [HWA, 5] in format [cx, cy, w, h, angle]

Source code in oriented_det/ops/gpu_ops.py

def generate_oriented_anchors_gpu(
    image_size: Tuple[int, int],
    feature_map_sizes: List[Tuple[int, int]],
    anchor_scales: List[float],
    anchor_ratios: List[float],
    anchor_angles: List[float],
    stride_per_level: List[int],
    device: torch.device = torch.device('cpu'),
    octave_base_scale: Optional[float] = None,
    scales_per_octave: Optional[int] = None,
) -> List[Tensor]:
    """Generate oriented anchors using vectorized GPU operations.

    Fully vectorized, no Python loops for anchor generation.

    Args:
        image_size: (height, width) of input image
        feature_map_sizes: List of (height, width) for each FPN level
        anchor_scales: List of anchor scales (one per level or shared)
        anchor_ratios: List of aspect ratios
        anchor_angles: List of anchor angles in radians
        stride_per_level: List of strides for each FPN level
        device: Device to create tensors on

    Returns:
        List of anchor tensors, each [H*W*A, 5] in format [cx, cy, w, h, angle]
    """
    anchors_per_level = []

    ratios = torch.tensor(anchor_ratios, device=device, dtype=torch.float32)
    angles = torch.tensor(anchor_angles, device=device, dtype=torch.float32)

    for level_idx, (feat_h, feat_w) in enumerate(feature_map_sizes):
        stride = stride_per_level[level_idx]

        # Grid of centers
        shift_x = (torch.arange(feat_w, device=device, dtype=torch.float32) + 0.5) * stride
        shift_y = (torch.arange(feat_h, device=device, dtype=torch.float32) + 0.5) * stride
        grid_y, grid_x = torch.meshgrid(shift_y, shift_x, indexing='ij')

        # Flatten: [H*W]
        grid_x = grid_x.reshape(-1)
        grid_y = grid_y.reshape(-1)
        num_positions = grid_x.shape[0]

        # Anchor dimensions
        sqrt_ratios = torch.sqrt(ratios)
        if octave_base_scale is not None and scales_per_octave is not None:
            factors = torch.tensor(
                [2.0 ** (i / float(scales_per_octave)) for i in range(int(scales_per_octave))],
                device=device,
                dtype=torch.float32,
            )
            base_sizes = stride * float(octave_base_scale) * factors  # [S]
            base_w = (base_sizes.unsqueeze(1) * sqrt_ratios.unsqueeze(0)).reshape(-1)
            base_h = (base_sizes.unsqueeze(1) / sqrt_ratios.unsqueeze(0)).reshape(-1)
        else:
            scale = anchor_scales[level_idx] if level_idx < len(anchor_scales) else anchor_scales[-1]
            effective_scale = scale * stride
            base_w = effective_scale * sqrt_ratios  # [R]
            base_h = effective_scale / sqrt_ratios  # [R]

        num_angles = len(angles)
        num_templates = base_w.numel()
        A = num_templates * num_angles

        # Expand templates x angles -> [A]
        base_w_exp = base_w.unsqueeze(1).expand(-1, num_angles).reshape(-1)
        base_h_exp = base_h.unsqueeze(1).expand(-1, num_angles).reshape(-1)
        angles_exp = angles.unsqueeze(0).expand(num_templates, -1).reshape(-1)

        # Create all anchors: [num_positions, A, 5]
        cx_all = grid_x.unsqueeze(1).expand(-1, A)  # [P, A]
        cy_all = grid_y.unsqueeze(1).expand(-1, A)  # [P, A]
        w_all = base_w_exp.unsqueeze(0).expand(num_positions, -1)  # [P, A]
        h_all = base_h_exp.unsqueeze(0).expand(num_positions, -1)  # [P, A]
        angle_all = angles_exp.unsqueeze(0).expand(num_positions, -1)  # [P, A]

        anchors = torch.stack([cx_all, cy_all, w_all, h_all, angle_all], dim=-1)  # [P, A, 5]
        anchors = anchors.view(-1, 5).detach()  # [P*A, 5]
        anchors.requires_grad_(False)

        anchors_per_level.append(anchors)

    return anchors_per_level

hbb_iou_gpu(boxes1_xyxy, boxes2_xyxy)

Compute axis-aligned (horizontal) IoU matrix between two sets of xyxy boxes.

Parameters:

Name	Type	Description	Default
boxes1_xyxy	Tensor	[N, 4] (x1, y1, x2, y2)	required
boxes2_xyxy	Tensor	[M, 4] (x1, y1, x2, y2)	required

Returns:

Type	Description
Tensor	[N, M] IoU matrix

Source code in oriented_det/ops/gpu_ops.py

def hbb_iou_gpu(boxes1_xyxy: Tensor, boxes2_xyxy: Tensor) -> Tensor:
    """Compute axis-aligned (horizontal) IoU matrix between two sets of xyxy boxes.

    Args:
        boxes1_xyxy: [N, 4] (x1, y1, x2, y2)
        boxes2_xyxy: [M, 4] (x1, y1, x2, y2)

    Returns:
        [N, M] IoU matrix
    """
    N = boxes1_xyxy.shape[0]
    M = boxes2_xyxy.shape[0]
    device = boxes1_xyxy.device
    if N == 0 or M == 0:
        return torch.zeros((N, M), device=device)
    # Broadcast: [N, 1, 4] vs [1, M, 4]
    x1_1 = boxes1_xyxy[:, 0].unsqueeze(1)   # [N, 1]
    y1_1 = boxes1_xyxy[:, 1].unsqueeze(1)
    x2_1 = boxes1_xyxy[:, 2].unsqueeze(1)
    y2_1 = boxes1_xyxy[:, 3].unsqueeze(1)
    x1_2 = boxes2_xyxy[:, 0].unsqueeze(0)   # [1, M]
    y1_2 = boxes2_xyxy[:, 1].unsqueeze(0)
    x2_2 = boxes2_xyxy[:, 2].unsqueeze(0)
    y2_2 = boxes2_xyxy[:, 3].unsqueeze(0)
    inter_x1 = torch.maximum(x1_1, x1_2)
    inter_y1 = torch.maximum(y1_1, y1_2)
    inter_x2 = torch.minimum(x2_1, x2_2)
    inter_y2 = torch.minimum(y2_1, y2_2)
    inter_w = (inter_x2 - inter_x1).clamp(min=0)
    inter_h = (inter_y2 - inter_y1).clamp(min=0)
    inter_area = inter_w * inter_h
    area1 = (x2_1 - x1_1) * (y2_1 - y1_1)  # [N, 1]
    area2 = (x2_2 - x1_2) * (y2_2 - y1_2)  # [1, M]
    union = area1 + area2 - inter_area
    iou = inter_area / (union + 1e-8)
    return iou

kfiou_loss(pred_decode, targets_decode, *, pred=None, target=None, fun=None, beta=1.0 / 9.0, eps=1e-06, reduction='mean', weight=None)

Reduced KFIoU loss (scalar unless reduction="none").

Source code in oriented_det/ops/kfiou.py

def kfiou_loss(
    pred_decode: Tensor,
    targets_decode: Tensor,
    *,
    pred: Optional[Tensor] = None,
    target: Optional[Tensor] = None,
    fun: Optional[str] = None,
    beta: float = 1.0 / 9.0,
    eps: float = 1e-6,
    reduction: str = "mean",
    weight: Optional[Tensor] = None,
) -> Tensor:
    """Reduced KFIoU loss (scalar unless ``reduction=\"none\"``)."""
    loss = kfiou_loss_per_box(
        pred_decode,
        targets_decode,
        pred=pred,
        target=target,
        fun=fun,
        beta=beta,
        eps=eps,
    )
    if weight is not None:
        if weight.dim() > 1:
            weight = weight.mean(dim=-1)
        loss = loss * weight
    if reduction == "none":
        return loss
    if reduction == "sum":
        return loss.sum()
    if reduction == "mean":
        if weight is not None:
            denom = weight.sum().clamp(min=eps)
            return loss.sum() / denom
        return loss.mean()
    raise ValueError(f"reduction must be none|mean|sum, got {reduction!r}")

kfiou_loss_per_box(pred_decode, targets_decode, *, pred=None, target=None, fun=None, beta=1.0 / 9.0, eps=1e-06)

KFIoU loss terms per box (MMRotate kfiou_loss without reduction).

Parameters:

Name	Type	Description	Default
pred_decode	Tensor	[N, 5] predicted decoded boxes.	required
targets_decode	Tensor	[N, 5] target decoded boxes.	required
pred	Optional[Tensor]	Optional [N, *] whose first two columns are predicted centers; defaults to pred_decode.	None
target	Optional[Tensor]	Optional [N, *] whose first two columns are GT centers; defaults to targets_decode.	None
fun	Optional[str]	None / "none" → 1 - KFIoU; "ln" → -log(KFIoU + eps); "exp" → exp(1 - KFIoU) - 1.	None
beta	float	Smooth L1 beta for the center term.	1.0 / 9.0
eps	float	Numerical stabilizer.	1e-06

Returns:

Type	Description
Tensor	[N] tensor of per-instance losses (non-negative).

Source code in oriented_det/ops/kfiou.py

def kfiou_loss_per_box(
    pred_decode: Tensor,
    targets_decode: Tensor,
    *,
    pred: Optional[Tensor] = None,
    target: Optional[Tensor] = None,
    fun: Optional[str] = None,
    beta: float = 1.0 / 9.0,
    eps: float = 1e-6,
) -> Tensor:
    """KFIoU loss terms per box (MMRotate ``kfiou_loss`` without reduction).

    Args:
        pred_decode: ``[N, 5]`` predicted decoded boxes.
        targets_decode: ``[N, 5]`` target decoded boxes.
        pred: Optional ``[N, *]`` whose first two columns are predicted centers;
            defaults to ``pred_decode``.
        target: Optional ``[N, *]`` whose first two columns are GT centers;
            defaults to ``targets_decode``.
        fun: ``None`` / ``\"none\"`` → ``1 - KFIoU``; ``\"ln\"`` → ``-log(KFIoU + eps)``;
            ``\"exp\"`` → ``exp(1 - KFIoU) - 1``.
        beta: Smooth L1 beta for the center term.
        eps: Numerical stabilizer.

    Returns:
        ``[N]`` tensor of per-instance losses (non-negative).
    """
    fun = _normalize_kfiou_fun(fun)
    if pred is None:
        pred = pred_decode
    if target is None:
        target = targets_decode
    xy_p = pred[:, :2]
    xy_t = target[:, :2]

    diff = torch.abs(xy_p - xy_t)
    xy_loss = torch.where(
        diff < beta,
        0.5 * diff * diff / beta,
        diff - 0.5 * beta,
    ).sum(dim=-1)

    kfiou = kfiou_overlap_ratio(pred_decode, targets_decode, eps=eps)
    if fun == "ln":
        kf_loss = -torch.log(kfiou + eps)
    elif fun == "exp":
        kf_loss = torch.exp(1.0 - kfiou) - 1.0
    else:
        kf_loss = 1.0 - kfiou

    return (xy_loss + kf_loss).clamp(min=0.0)

kfiou_overlap_ratio(pred_decode, targets_decode, *, eps=1e-06)

Per-row KFIoU overlap ratio Vb / (Vb_p + Vb_t - Vb + eps) in [0, 1] (typically).

Parameters:

Name	Type	Description	Default
pred_decode	Tensor	[N, 5] predicted decoded boxes (differentiable).	required
targets_decode	Tensor	[N, 5] matched GT boxes (usually detached).	required
eps	float	Small constant for numerical stability in the denominator.	1e-06

Returns:

Type	Description
Tensor	Tensor [N] overlap ratios. Always returned in float32 so that
Tensor	downstream -log(kfiou + eps) / exp(1 - kfiou) stay numerically
Tensor	stable under AMP (eps=1e-6 underflows to 0 in fp16). Autograd
Tensor	casts gradients back to the input dtype for backward.

Source code in oriented_det/ops/kfiou.py

def kfiou_overlap_ratio(
    pred_decode: Tensor,
    targets_decode: Tensor,
    *,
    eps: float = 1e-6,
) -> Tensor:
    """Per-row KFIoU overlap ratio ``Vb / (Vb_p + Vb_t - Vb + eps)`` in ``[0, 1]`` (typically).

    Args:
        pred_decode: ``[N, 5]`` predicted decoded boxes (differentiable).
        targets_decode: ``[N, 5]`` matched GT boxes (usually detached).
        eps: Small constant for numerical stability in the denominator.

    Returns:
        Tensor ``[N]`` overlap ratios. Always returned in ``float32`` so that
        downstream ``-log(kfiou + eps)`` / ``exp(1 - kfiou)`` stay numerically
        stable under AMP (``eps=1e-6`` underflows to 0 in fp16). Autograd
        casts gradients back to the input dtype for backward.
    """
    with _fp32_no_autocast_ctx(pred_decode):
        # Cast inputs (not just sigmas) so that wh.square() in
        # xy_wh_r_to_xy_sigma cannot overflow fp16 for ~1024 px boxes.
        pred_f = pred_decode.float()
        targets_f = targets_decode.float()
        _, sigma_p = xy_wh_r_to_xy_sigma(pred_f)
        _, sigma_t = xy_wh_r_to_xy_sigma(targets_f)

        det_p = sigma_p.det().clamp(min=eps)
        det_t = sigma_t.det().clamp(min=eps)
        vb_p = 4.0 * det_p.sqrt()
        vb_t = 4.0 * det_t.sqrt()

        sigma_sum = sigma_p + sigma_t
        # Kalman gain K = sigma_p @ sigma_sum^{-1}. The sum can be singular or
        # badly conditioned when both covariances are (near) rank-1 with the same
        # kernel (degenerate skinny boxes, same orientation) or under unstable
        # predictions (e.g. LR finder at very high LR). Diagonal jitter helps
        # conditioning; ``pinv`` avoids ``linalg.solve`` hard failures on CUDA when
        # LAPACK still reports a singular/invalid factorization.
        solve_jitter = max(float(eps), 1e-8)
        eye2 = torch.eye(2, device=sigma_sum.device, dtype=sigma_sum.dtype)
        eye2 = eye2.reshape(*([1] * (sigma_sum.dim() - 2)), 2, 2).expand_as(sigma_sum)
        sigma_sum_reg = sigma_sum + solve_jitter * eye2
        pinv_sum = torch.linalg.pinv(sigma_sum_reg)
        k_mat = torch.matmul(sigma_p, pinv_sum)
        sigma_fused = sigma_p - k_mat.bmm(sigma_p)
        det_f = sigma_fused.det()
        vb = 4.0 * det_f.clamp(min=0.0).sqrt()
        vb = torch.where(torch.isfinite(vb), vb, torch.zeros_like(vb))

        denom = vb_p + vb_t - vb + eps
        kfiou = vb / denom.clamp(min=eps)
        kfiou = kfiou.clamp(0.0, 1.0)
    return kfiou

match_anchors_to_gt_gpu(anchors, gt_boxes, positive_iou_threshold=0.7, negative_iou_threshold=0.3, use_hbb_for_assignment=False, min_pos_iou=0.3, match_low_quality=True)

Match anchors to GT using GPU-accelerated IoU.

Parameters:

Name	Type	Description	Default
anchors	Tensor	[N, 5] anchors	required
gt_boxes	Tensor	[M, 5] ground truth boxes	required
positive_iou_threshold	float	IoU threshold for positive	0.7
negative_iou_threshold	float	IoU threshold for negative	0.3
use_hbb_for_assignment	bool	If True, use axis-aligned (HBB) overlap for assignment instead of oriented IoU. Each box is converted to its horizontal bounding box (min/max of corners) then standard IoU is computed. Can yield more positive anchors when GT is rotated.	False

Returns:

Name	Type	Description
labels	Tensor	[N] with -1=ignore, 0=background, 1=foreground
matched_gt_indices	Tensor	[N] index of matched GT (-1 if none)

Source code in oriented_det/ops/gpu_ops.py

def match_anchors_to_gt_gpu(
    anchors: Tensor,
    gt_boxes: Tensor,
    positive_iou_threshold: float = 0.7,
    negative_iou_threshold: float = 0.3,
    use_hbb_for_assignment: bool = False,
    min_pos_iou: float = 0.3,
    match_low_quality: bool = True,
) -> Tuple[Tensor, Tensor]:
    """Match anchors to GT using GPU-accelerated IoU.

    Args:
        anchors: [N, 5] anchors
        gt_boxes: [M, 5] ground truth boxes
        positive_iou_threshold: IoU threshold for positive
        negative_iou_threshold: IoU threshold for negative
        use_hbb_for_assignment: If True, use axis-aligned (HBB) overlap for assignment
                               instead of oriented IoU. Each box is converted to its
                               horizontal bounding box (min/max of corners) then
                               standard IoU is computed. Can yield more positive
                               anchors when GT is rotated.

    Returns:
        labels: [N] with -1=ignore, 0=background, 1=foreground
        matched_gt_indices: [N] index of matched GT (-1 if none)
    """
    device = anchors.device
    N = anchors.shape[0]
    M = gt_boxes.shape[0]

    labels = torch.full((N,), -1, dtype=torch.long, device=device)
    matched_gt_indices = torch.full((N,), -1, dtype=torch.long, device=device)

    if M == 0:
        labels.fill_(0)
        return labels, matched_gt_indices

    if use_hbb_for_assignment:
        # HBB assignment only needs the best GT per anchor and best anchor per GT.
        # Process large chunks against all GTs at once so dense P2 grids do not
        # create thousands of tiny GPU launches or materialize a full N x M matrix.
        anchor_xyxy = obb_to_xyxy_gpu(anchors)
        gt_xyxy = obb_to_xyxy_gpu(gt_boxes)

        max_iou_per_anchor = torch.zeros((N,), device=device)
        best_gt_per_anchor = torch.zeros((N,), dtype=torch.long, device=device)
        max_iou_per_gt = torch.zeros((M,), device=device)
        best_anchor_per_gt = torch.zeros((M,), dtype=torch.long, device=device)

        target_pairs_per_chunk = 8_000_000
        chunk_size = max(1, min(N, target_pairs_per_chunk // max(1, M)))
        for start in range(0, N, chunk_size):
            end = min(start + chunk_size, N)
            ious = hbb_iou_gpu(anchor_xyxy[start:end], gt_xyxy)

            chunk_anchor_iou, chunk_best_gt = ious.max(dim=1)
            max_iou_per_anchor[start:end] = chunk_anchor_iou
            best_gt_per_anchor[start:end] = chunk_best_gt

            chunk_gt_iou, chunk_best_anchor = ious.max(dim=0)
            better_gt = chunk_gt_iou > max_iou_per_gt
            if better_gt.any():
                max_iou_per_gt[better_gt] = chunk_gt_iou[better_gt]
                best_anchor_per_gt[better_gt] = chunk_best_anchor[better_gt] + start

        # Mark anchors above positive threshold as positive
        positive_mask = max_iou_per_anchor >= positive_iou_threshold
        labels[positive_mask] = 1
        matched_gt_indices[positive_mask] = best_gt_per_anchor[positive_mask]

        # Best anchor per GT - mark as positive when match_low_quality and IoU >= min_pos_iou (MMRotate-style)
        best_anchor_mask = torch.zeros(N, dtype=torch.bool, device=device)
        if match_low_quality:
            for gt_idx in range(M):
                if max_iou_per_gt[gt_idx] >= min_pos_iou:
                    anchor_idx = best_anchor_per_gt[gt_idx]
                    labels[anchor_idx] = 1
                    matched_gt_indices[anchor_idx] = gt_idx
                    best_anchor_mask[anchor_idx] = True

        # Mark anchors below negative threshold as negative
        negative_mask = (max_iou_per_anchor < negative_iou_threshold) & (~best_anchor_mask)
        labels[negative_mask] = 0

        return labels, matched_gt_indices
    else:
        # Oriented IoU in chunks to manage memory
        chunk_size = 5000
        iou_matrix = torch.zeros((N, M), device=device)
        for start in range(0, N, chunk_size):
            end = min(start + chunk_size, N)
            iou_matrix[start:end] = oriented_box_iou_gpu(anchors[start:end], gt_boxes)

    # Best GT per anchor
    max_iou_per_anchor, best_gt_per_anchor = iou_matrix.max(dim=1)

    # Mark anchors above positive threshold as positive
    positive_mask = max_iou_per_anchor >= positive_iou_threshold
    labels[positive_mask] = 1
    matched_gt_indices[positive_mask] = best_gt_per_anchor[positive_mask]

    # Best anchor per GT - mark as positive when match_low_quality and IoU >= min_pos_iou (MMRotate-style)
    max_iou_per_gt, best_anchor_per_gt = iou_matrix.max(dim=0)
    best_anchor_mask = torch.zeros(N, dtype=torch.bool, device=device)
    if match_low_quality:
        for gt_idx in range(M):
            if max_iou_per_gt[gt_idx] >= min_pos_iou:
                anchor_idx = best_anchor_per_gt[gt_idx]
                labels[anchor_idx] = 1
                matched_gt_indices[anchor_idx] = gt_idx
                best_anchor_mask[anchor_idx] = True

    # Mark anchors below negative threshold as negative
    # BUT protect "best anchor per GT" from being overwritten
    negative_mask = (max_iou_per_anchor < negative_iou_threshold) & (~best_anchor_mask)
    labels[negative_mask] = 0

    return labels, matched_gt_indices

mean_auxiliary_box_reg_loss(decoded_boxes, matched_gt, *, loss_type='riou', kfiou_fun=None, probiou_mode=None)

Scalar auxiliary loss on decoded positives (mean over boxes).

Parameters:

Name	Type	Description	Default
decoded_boxes	Tensor	[N, 5] predictions (typically require grad).	required
matched_gt	Tensor	[N, 5] matched ground truth (typically detached).	required
loss_type	str	"riou" for mean(1 - rIoU) (sampling/GPU backend via :func:rotated_ops.pairwise_rotated_iou); "kfiou" for reduced KFIoU loss (Gaussian overlap + center Smooth L1); "probiou" for mean ProbIoU (see :func:probiou.probiou_loss).	'riou'
kfiou_fun	Optional[str]	When loss_type="kfiou", optional ln / exp (see :func:kfiou_loss); none uses 1 - KFIoU.	None
probiou_mode	Optional[str]	When loss_type="probiou", "l1" (default) or "l2".	None

Source code in oriented_det/ops/kfiou.py

def mean_auxiliary_box_reg_loss(
    decoded_boxes: Tensor,
    matched_gt: Tensor,
    *,
    loss_type: str = "riou",
    kfiou_fun: Optional[str] = None,
    probiou_mode: Optional[str] = None,
) -> Tensor:
    """Scalar auxiliary loss on decoded positives (mean over boxes).

    Args:
        decoded_boxes: ``[N, 5]`` predictions (typically require grad).
        matched_gt: ``[N, 5]`` matched ground truth (typically detached).
        loss_type: ``\"riou\"`` for ``mean(1 - rIoU)`` (sampling/GPU backend via
            :func:`rotated_ops.pairwise_rotated_iou`); ``\"kfiou\"`` for reduced
            KFIoU loss (Gaussian overlap + center Smooth L1); ``\"probiou\"`` for
            mean ProbIoU (see :func:`probiou.probiou_loss`).
        kfiou_fun: When ``loss_type=\"kfiou\"``, optional ``ln`` / ``exp`` (see
            :func:`kfiou_loss`); ``none`` uses ``1 - KFIoU``.
        probiou_mode: When ``loss_type=\"probiou\"``, ``\"l1\"`` (default) or ``\"l2\"``.
    """
    lt = (loss_type or "riou").strip().lower()
    if lt == "riou":
        from .rotated_ops import pairwise_rotated_iou

        iou_matrix = pairwise_rotated_iou(decoded_boxes, matched_gt)
        pairwise_iou = torch.diagonal(iou_matrix, offset=0).clamp(0.0, 1.0)
        return 1.0 - pairwise_iou.mean()
    if lt == "kfiou":
        return kfiou_loss(
            decoded_boxes,
            matched_gt,
            fun=_normalize_kfiou_fun(kfiou_fun),
            reduction="mean",
        )
    if lt == "probiou":
        from .probiou import _normalize_probiou_mode, probiou_loss

        return probiou_loss(
            decoded_boxes,
            matched_gt,
            mode=_normalize_probiou_mode(probiou_mode),
            reduction="mean",
        )
    raise ValueError(
        f"Unknown loss_type {loss_type!r} for auxiliary box loss; "
        "use 'riou', 'kfiou', or 'probiou'"
    )

obb_to_xyxy_gpu(boxes)

Convert oriented boxes [N, 5] (cx, cy, w, h, angle) to axis-aligned xyxy [N, 4].

Uses the 4 corner vertices of each OBB and takes min/max to get the HBB. Fully vectorized and GPU-compatible.

Parameters:

Name	Type	Description	Default
boxes	Tensor	[N, 5] in format [cx, cy, w, h, angle]	required

Returns:

Type	Description
Tensor	[N, 4] in format [x1, y1, x2, y2]

Source code in oriented_det/ops/gpu_ops.py

def obb_to_xyxy_gpu(boxes: Tensor) -> Tensor:
    """Convert oriented boxes [N, 5] (cx, cy, w, h, angle) to axis-aligned xyxy [N, 4].

    Uses the 4 corner vertices of each OBB and takes min/max to get the HBB.
    Fully vectorized and GPU-compatible.

    Args:
        boxes: [N, 5] in format [cx, cy, w, h, angle]

    Returns:
        [N, 4] in format [x1, y1, x2, y2]
    """
    verts = _box_vertices(boxes)  # [N, 4, 2]
    xy_min = verts.min(dim=1).values   # [N, 2]
    xy_max = verts.max(dim=1).values   # [N, 2]
    return torch.cat([xy_min, xy_max], dim=1)  # [N, 4]

oriented_box_hbb_iou_gpu(anchors, gt_boxes)

Compute HBB (axis-aligned) IoU matrix between oriented boxes.

Each OBB is converted to its axis-aligned bounding box (min/max of corners), then standard xyxy IoU is computed. Fully GPU-compatible.

Parameters:

Name	Type	Description	Default
anchors	Tensor	[N, 5] (cx, cy, w, h, angle)	required
gt_boxes	Tensor	[M, 5] (cx, cy, w, h, angle)	required

Returns:

Type	Description
Tensor	[N, M] IoU matrix

Source code in oriented_det/ops/gpu_ops.py

def oriented_box_hbb_iou_gpu(anchors: Tensor, gt_boxes: Tensor) -> Tensor:
    """Compute HBB (axis-aligned) IoU matrix between oriented boxes.

    Each OBB is converted to its axis-aligned bounding box (min/max of corners),
    then standard xyxy IoU is computed. Fully GPU-compatible.

    Args:
        anchors: [N, 5] (cx, cy, w, h, angle)
        gt_boxes: [M, 5] (cx, cy, w, h, angle)

    Returns:
        [N, M] IoU matrix
    """
    xyxy_anchors = obb_to_xyxy_gpu(anchors)
    xyxy_gt = obb_to_xyxy_gpu(gt_boxes)
    return hbb_iou_gpu(xyxy_anchors, xyxy_gt)

oriented_box_iou_gpu(boxes1, boxes2, num_samples=None)

Compute oriented IoU matrix using sampling-based approximation.

This is a fully vectorized GPU implementation that approximates IoU by sampling points inside boxes and checking containment.

Parameters:

Name	Type	Description	Default
boxes1	Tensor	[N, 5] boxes in format [cx, cy, w, h, angle]	required
boxes2	Tensor	[M, 5] boxes in format [cx, cy, w, h, angle]	required
num_samples	Optional[int]	Samples per box (perfect square, e.g. 100 = 10×10). If None, sample count follows geometry-based sizing (default): grid side from each box w/h and aspect ratio with target spacing ORIENTED_DET_GPU_ORIENTED_IOU_TARGET_SPACING_PX (default 2 px), clamped to [ORIENTED_DET_GPU_ORIENTED_IOU_MIN_SAMPLES, ORIENTED_DET_GPU_ORIENTED_IOU_MAX_SAMPLES] (defaults 25 … 1024). Disable geometry with ORIENTED_DET_GPU_ORIENTED_IOU_SAMPLE_BY_MAX_SIDE=0 (flat 100-sample grid, debug only).	None

Returns:

Type	Description
Tensor	[N, M] IoU matrix

Source code in oriented_det/ops/gpu_ops.py

def oriented_box_iou_gpu(
    boxes1: Tensor,
    boxes2: Tensor,
    num_samples: Optional[int] = None,
) -> Tensor:
    """Compute oriented IoU matrix using sampling-based approximation.

    This is a fully vectorized GPU implementation that approximates IoU
    by sampling points inside boxes and checking containment.

    Args:
        boxes1: [N, 5] boxes in format [cx, cy, w, h, angle]
        boxes2: [M, 5] boxes in format [cx, cy, w, h, angle]
        num_samples: Samples per box (perfect square, e.g. 100 = 10×10).
            If ``None``, sample count follows geometry-based sizing (default):
            grid side from each box ``w``/``h`` and aspect ratio with target spacing
            ``ORIENTED_DET_GPU_ORIENTED_IOU_TARGET_SPACING_PX`` (default ``2`` px),
            clamped to ``[ORIENTED_DET_GPU_ORIENTED_IOU_MIN_SAMPLES, ORIENTED_DET_GPU_ORIENTED_IOU_MAX_SAMPLES]``
            (defaults ``25`` … ``1024``). Disable geometry with ``ORIENTED_DET_GPU_ORIENTED_IOU_SAMPLE_BY_MAX_SIDE=0``
            (flat ``100``-sample grid, debug only).

    Returns:
        [N, M] IoU matrix
    """
    device = boxes1.device
    N = boxes1.shape[0]
    M = boxes2.shape[0]

    if N == 0 or M == 0:
        return torch.zeros((N, M), device=device)

    num_samples = _iou_num_samples(boxes1, boxes2, num_samples)

    # Compute areas
    area1 = boxes1[:, 2] * boxes1[:, 3]  # [N]
    area2 = boxes2[:, 2] * boxes2[:, 3]  # [M]

    # AABB pre-filtering
    verts1 = _box_vertices(boxes1)  # [N, 4, 2]
    verts2 = _box_vertices(boxes2)  # [M, 4, 2]

    aabb1_min = verts1.min(dim=1).values  # [N, 2]
    aabb1_max = verts1.max(dim=1).values  # [N, 2]
    aabb2_min = verts2.min(dim=1).values  # [M, 2]
    aabb2_max = verts2.max(dim=1).values  # [M, 2]

    # Check AABB overlap
    aabb_overlap = (
        (aabb1_min[:, 0:1] <= aabb2_max[None, :, 0]) &
        (aabb1_max[:, 0:1] >= aabb2_min[None, :, 0]) &
        (aabb1_min[:, 1:2] <= aabb2_max[None, :, 1]) &
        (aabb1_max[:, 1:2] >= aabb2_min[None, :, 1])
    )  # [N, M]

    # Initialize IoU matrix
    iou_matrix = torch.zeros((N, M), device=device)

    # Find overlapping pairs
    overlap_mask = aabb_overlap
    if not overlap_mask.any():
        return iou_matrix

    # Generate samples for boxes1
    samples1 = _generate_box_samples(boxes1, num_samples)  # [N, S, 2]
    S = samples1.shape[1]

    # Also generate samples for boxes2 (for symmetric IoU)
    samples2 = _generate_box_samples(boxes2, num_samples)  # [M, S, 2]

    # For each box in boxes1, count how many of its samples are in each box of boxes2.
    # This is done in chunks to manage memory. _points_in_boxes_batch builds [P, M, 2]
    # float tensors; with P = chunk_n*S and M = num boxes, we need chunk_n*S*M*2*4 bytes.
    # Cap so that allocation stays under ~1.5 GB to avoid CUDA OOM when M is large.
    target_max_bytes = 1.5 * (1024 ** 3)
    max_elements = int(target_max_bytes / (4 * 2))  # float32, 2 coords per element
    max_inner_chunk = max(1, max_elements // (S * M))
    chunk_size = min(1000, max_inner_chunk)

    # Track total samples inside for debugging
    _total_samples_in_box2 = 0
    _total_samples_in_box1 = 0
    _num_chunks = 0

    for i_start in range(0, N, chunk_size):
        i_end = min(i_start + chunk_size, N)
        chunk_samples = samples1[i_start:i_end]  # [chunk, S, 2]
        chunk_n = chunk_samples.shape[0]

        # Flatten samples: [chunk * S, 2]
        flat_samples = chunk_samples.view(-1, 2)

        # Check which samples are in each box2: [chunk * S, M]
        in_boxes2 = _points_in_boxes_batch(flat_samples, verts2)

        # Reshape: [chunk, S, M]
        in_boxes2 = in_boxes2.view(chunk_n, S, M)

        # Count samples in intersection: [chunk, M]
        count_in_box2 = in_boxes2.sum(dim=1).float()
        _total_samples_in_box2 += int(count_in_box2.sum().item())

        # Approximate intersection area from boxes1 perspective
        # intersection ≈ (count_in_box2 / S) * area1
        inter_approx1 = (count_in_box2 / S) * area1[i_start:i_end, None]

        # Also compute intersection from boxes2 perspective for symmetry
        # For each box2, check how many of its samples are in boxes1[i_start:i_end]
        flat_samples2 = samples2.view(-1, 2)  # [M * S, 2]
        in_boxes1_chunk = _points_in_boxes_batch(flat_samples2, verts1[i_start:i_end])  # [M*S, chunk]
        in_boxes1_chunk = in_boxes1_chunk.view(M, S, chunk_n)  # [M, S, chunk]
        count_in_box1 = in_boxes1_chunk.sum(dim=1).float().T  # [chunk, M]
        _total_samples_in_box1 += int(count_in_box1.sum().item())

        inter_approx2 = (count_in_box1 / S) * area2[None, :]

        # Use MAX of both estimates (not geometric mean)
        # Geometric mean fails when one box is much smaller than the other:
        # - Small proposal inside large GT: inter_approx1 is good, inter_approx2 is ~0
        # - sqrt(good * 0) = 0, which is wrong!
        # MAX gives a better estimate in asymmetric size scenarios
        inter_approx = torch.maximum(inter_approx1, inter_approx2)

        # Clamp intersection to theoretical maximum (can't exceed smaller box area)
        # This prevents IoU > 1.0 due to sampling approximation errors
        max_possible_inter = torch.minimum(area1[i_start:i_end, None], area2[None, :])
        inter_approx = torch.minimum(inter_approx, max_possible_inter)

        _num_chunks += 1

        # Compute union (ensure union >= intersection to prevent IoU > 1.0)
        union_approx = area1[i_start:i_end, None] + area2[None, :] - inter_approx
        union_approx = torch.maximum(union_approx, inter_approx + 1e-8)

        # Compute IoU
        iou_chunk = inter_approx / (union_approx + 1e-8)

        # Final clamp to [0, 1] to handle any remaining numerical errors
        iou_chunk = torch.clamp(iou_chunk, 0.0, 1.0)

        # Apply AABB mask
        iou_chunk = iou_chunk * overlap_mask[i_start:i_end].float()

        iou_matrix[i_start:i_end] = iou_chunk

    return iou_matrix

oriented_nms(boxes, scores, iou_threshold=0.5, max_detections=None, labels=None, backend=None)

Perform oriented NMS and return kept indices.

This is a CPU Python implementation. For GPU-accelerated NMS, use: from oriented_det.ops.gpu_ops import oriented_nms_gpu

Parameters:

Name	Type	Description	Default
boxes	Sequence[RBox | Sequence[float]]	Sequence of RBoxes or 5-tuples (cx, cy, w, h, angle).	required
scores	Sequence[float]	Confidence scores aligned with boxes.	required
iou_threshold	float	Overlap threshold to suppress boxes.	0.5
max_detections	int | None	Optional cap on number of returns.	None
labels	Optional[Sequence[int]]	Optional sequence of class labels aligned with boxes. If provided, NMS only suppresses boxes of the same class (class-aware NMS). Boxes of different classes will not suppress each other even if they overlap.	None
backend	Optional[str]	Optional backend selector ('python' or 'torch'); for API compatibility.	None

Returns:

Type	Description
List[int]	List of indices of boxes to keep after NMS.

Examples:

>>> boxes = [RBox(0, 0, 2, 2, 0), RBox(0.5, 0.5, 2, 2, 0)]
>>> scores = [0.9, 0.8]
>>> keep = oriented_nms(boxes, scores, iou_threshold=0.3)

>>> # Class-aware NMS
>>> labels = [1, 2]  # Different classes
>>> keep = oriented_nms(boxes, scores, labels=labels, iou_threshold=0.3)
>>> # Both boxes kept since they're different classes

Source code in oriented_det/ops/nms.py

def oriented_nms(
    boxes: Sequence[RBox | Sequence[float]],
    scores: Sequence[float],
    iou_threshold: float = 0.5,
    max_detections: int | None = None,
    labels: Optional[Sequence[int]] = None,
    backend: Optional[str] = None,
) -> List[int]:
    """Perform oriented NMS and return kept indices.

    This is a CPU Python implementation. For GPU-accelerated NMS, use:
    `from oriented_det.ops.gpu_ops import oriented_nms_gpu`

    Args:
        boxes: Sequence of RBoxes or 5-tuples ``(cx, cy, w, h, angle)``.
        scores: Confidence scores aligned with `boxes`.
        iou_threshold: Overlap threshold to suppress boxes.
        max_detections: Optional cap on number of returns.
        labels: Optional sequence of class labels aligned with `boxes`. If provided,
                NMS only suppresses boxes of the same class (class-aware NMS).
                Boxes of different classes will not suppress each other even if they overlap.
        backend: Optional backend selector ('python' or 'torch'); for API compatibility.

    Returns:
        List of indices of boxes to keep after NMS.

    Examples:
        >>> boxes = [RBox(0, 0, 2, 2, 0), RBox(0.5, 0.5, 2, 2, 0)]
        >>> scores = [0.9, 0.8]
        >>> keep = oriented_nms(boxes, scores, iou_threshold=0.3)

        >>> # Class-aware NMS
        >>> labels = [1, 2]  # Different classes
        >>> keep = oriented_nms(boxes, scores, labels=labels, iou_threshold=0.3)
        >>> # Both boxes kept since they're different classes
    """
    if backend is not None:
        if backend == "invalid":
            raise ValueError("Invalid backend: invalid")
        if backend == "torch":
            from . import utils as ops_utils
            if getattr(ops_utils, "TORCH_NMS_ROTATED", None) is None:
                raise RuntimeError("torch backend for oriented_nms is not available")
    if len(boxes) != len(scores):
        raise ValueError("boxes and scores must have the same length.")
    if labels is not None and len(labels) != len(boxes):
        raise ValueError("labels must have the same length as boxes.")
    if iou_threshold < 0 or iou_threshold > 1:
        raise ValueError("iou_threshold must be between 0 and 1.")

    rboxes = [as_rbox(b) for b in boxes]
    return _python_nms(rboxes, scores, iou_threshold, max_detections, labels)

oriented_nms_gpu(boxes, scores, iou_threshold=0.5, max_detections=None)

GPU-accelerated oriented NMS using sampling-based IoU.

Uses ORIENTED_DET_GPU_NMS_IOU_SAMPLES as an optional floor when geometry sizing is enabled; otherwise a flat count (default 100). Geometry uses the same rules as oriented IoU (see ops/README.md).

Parameters:

Name	Type	Description	Default
boxes	Tensor	[N, 5] boxes	required
scores	Tensor	[N] confidence scores	required
iou_threshold	float	Suppression threshold	0.5
max_detections	Optional[int]	Maximum detections to keep	None

Returns:

Type	Description
Tensor	Indices of kept boxes

Source code in oriented_det/ops/gpu_ops.py

def oriented_nms_gpu(
    boxes: Tensor,
    scores: Tensor,
    iou_threshold: float = 0.5,
    max_detections: Optional[int] = None,
) -> Tensor:
    """GPU-accelerated oriented NMS using sampling-based IoU.

    Uses ``ORIENTED_DET_GPU_NMS_IOU_SAMPLES`` as an optional floor when geometry
    sizing is enabled; otherwise a flat count (default ``100``). Geometry uses the
    same rules as oriented IoU (see ``ops/README.md``).

    Args:
        boxes: [N, 5] boxes
        scores: [N] confidence scores
        iou_threshold: Suppression threshold
        max_detections: Maximum detections to keep

    Returns:
        Indices of kept boxes
    """
    device = boxes.device
    N = boxes.shape[0]

    if N == 0:
        return torch.tensor([], dtype=torch.long, device=device)

    # Sort by score
    _, order = scores.sort(descending=True)

    # ALWAYS limit boxes before NxN IoU computation to prevent OOM.
    # Rotated RetinaNet can produce 100k+ proposals per class; full pairwise IoU would need ~1TB.
    max_candidates = (
        min(max_detections * 3, MAX_BOXES_FOR_NMS) if max_detections is not None
        else MAX_BOXES_FOR_NMS
    )
    order = order[:min(len(order), max_candidates)]

    ordered_boxes = boxes[order]
    M = ordered_boxes.shape[0]

    if M == 0:
        return torch.tensor([], dtype=torch.long, device=device)

    # Candidate pair pruning: rotated IoU is only needed where it can plausibly exceed
    # the threshold. Upper-bound IoU with the AABB intersection over the rotated-box
    # union: iou_rot <= I_aabb / (area_i + area_j - I_aabb). A 0.5 safety factor on the
    # threshold absorbs sampling-grid noise in the estimate below.
    verts = _box_vertices(ordered_boxes)  # [M, 4, 2]
    xy_min = verts.min(dim=1).values
    xy_max = verts.max(dim=1).values
    inter_wh = (
        torch.minimum(xy_max.unsqueeze(1), xy_max.unsqueeze(0))
        - torch.maximum(xy_min.unsqueeze(1), xy_min.unsqueeze(0))
    ).clamp(min=0)  # [M, M, 2]
    inter_aabb = inter_wh[..., 0] * inter_wh[..., 1]  # [M, M]
    areas = ordered_boxes[:, 2] * ordered_boxes[:, 3]  # [M]
    union_lb = (areas.unsqueeze(1) + areas.unsqueeze(0) - inter_aabb).clamp(min=1e-8)
    iou_upper_bound = inter_aabb / union_lb
    # Strictly upper triangular: boxes are score-sorted, suppression flows i -> j (i < j).
    cand = torch.triu(iou_upper_bound > iou_threshold * 0.5, diagonal=1)
    pair_i, pair_j = cand.nonzero(as_tuple=True)  # [P], [P]

    suppress_matrix = torch.zeros((M, M), dtype=torch.bool, device=device)
    if pair_i.numel() > 0:
        # Sampled rotated IoU computed per pair (O(P*S)) instead of as a dense
        # [M, M] matrix (O(M^2 * S)) — same estimator as oriented_box_iou_gpu:
        # max of both containment perspectives, clamped by the smaller box area.
        num_s = _nms_num_samples_for_boxes(ordered_boxes)
        target_max_bytes = 1.5 * (1024 ** 3)
        pair_chunk = max(1, int(target_max_bytes / (4 * 2 * num_s)))
        for p_start in range(0, pair_i.numel(), pair_chunk):
            p_end = min(p_start + pair_chunk, pair_i.numel())
            pi = pair_i[p_start:p_end]
            pj = pair_j[p_start:p_end]
            boxes_i = ordered_boxes[pi]
            boxes_j = ordered_boxes[pj]
            area_i = areas[pi]
            area_j = areas[pj]

            samples_i = _generate_box_samples(boxes_i, num_s)  # [p, S, 2]
            samples_j = _generate_box_samples(boxes_j, num_s)
            in_j = _points_in_paired_boxes(samples_i, verts[pj])  # [p, S]
            in_i = _points_in_paired_boxes(samples_j, verts[pi])

            inter_1 = in_j.sum(dim=1).float() / float(num_s) * area_i
            inter_2 = in_i.sum(dim=1).float() / float(num_s) * area_j
            inter = torch.maximum(inter_1, inter_2)
            inter = torch.minimum(inter, torch.minimum(area_i, area_j))
            union = torch.maximum(area_i + area_j - inter, inter + 1e-8)
            iou = torch.clamp(inter / (union + 1e-8), 0.0, 1.0)

            sup = iou > iou_threshold
            suppress_matrix[pi[sup], pj[sup]] = True

    # Greedy NMS via iterative matrix suppression (Cluster-NMS, Zheng et al. 2020).
    # Iterating keep[j] = not any_i(suppress_matrix[i, j] and keep[i]) converges to the
    # unique fixpoint, which equals the sequential greedy NMS result. Convergence takes
    # at most the longest suppression-chain depth (typically < 10 iterations), so this
    # replaces M kernel launches + M host syncs of the previous per-row Python loop
    # (the dominant validation cost for large weakly-suppressing candidate pools).
    keep_mask = torch.ones(M, dtype=torch.bool, device=device)
    for _ in range(M):
        new_keep = ~(suppress_matrix & keep_mask.unsqueeze(1)).any(dim=0)
        if torch.equal(new_keep, keep_mask):
            break
        keep_mask = new_keep

    kept_indices = order[keep_mask]

    # Apply max_detections limit after processing all candidates
    if max_detections is not None:
        kept_indices = kept_indices[:max_detections]

    return kept_indices

polygon_iou(poly_a, poly_b, *, intersection_backend='auto')

Compute IoU between two polygon-like objects.

Parameters:

Name	Type	Description	Default
poly_a		First polygon-like object	required
poly_b		Second polygon-like object	required
intersection_backend	str	Backend for intersection calculation ('auto', 'python', or 'shapely')	'auto'

Returns:

Type	Description
float	IoU value as a float

Source code in oriented_det/ops/iou.py

def polygon_iou(poly_a, poly_b, *, intersection_backend: str = "auto") -> float:
    """Compute IoU between two polygon-like objects.

    Args:
        poly_a: First polygon-like object
        poly_b: Second polygon-like object
        intersection_backend: Backend for intersection calculation ('auto', 'python', or 'shapely')

    Returns:
        IoU value as a float
    """
    pa = as_polygon(poly_a)
    pb = as_polygon(poly_b)
    inter = polygon_intersection_area(pa, pb, backend=intersection_backend)
    if inter == 0.0:
        return 0.0
    union = pa.area + pb.area - inter
    if union <= 0.0:
        return 0.0
    return inter / union

probiou_loss(pred, target, *, eps=0.001, mode='l1', reduction='mean', weight=None)

Reduced ProbIoU loss (scalar unless reduction="none").

Source code in oriented_det/ops/probiou.py

def probiou_loss(
    pred: Tensor,
    target: Tensor,
    *,
    eps: float = 1e-3,
    mode: ProbIoUMode = "l1",
    reduction: str = "mean",
    weight: Optional[Tensor] = None,
) -> Tensor:
    """Reduced ProbIoU loss (scalar unless ``reduction=\"none\"``)."""
    loss = probiou_loss_per_box(pred, target, eps=eps, mode=mode)
    if weight is not None:
        if weight.dim() > 1:
            weight = weight.mean(dim=-1)
        loss = loss * weight
    if reduction == "none":
        return loss
    if reduction == "sum":
        return loss.sum()
    if reduction == "mean":
        if weight is not None:
            denom = weight.sum().clamp(min=eps)
            return loss.sum() / denom
        return loss.mean()
    raise ValueError(f"reduction must be none|mean|sum, got {reduction!r}")

probiou_loss_per_box(pred, target, *, eps=0.001, mode='l1')

Per-instance ProbIoU loss [N] (lower is better overlap).

Parameters:

Name	Type	Description	Default
pred	Tensor	[N, 5] predicted boxes (cx, cy, w, h, angle_rad).	required
target	Tensor	[N, 5] target boxes (same layout).	required
eps	float	Stabilizer for divisions and logarithms.	0.001
mode	ProbIoUMode	"l1" → bounded metric in [0, 1]; "l2" → unbounded -log(1 - l1^2) variant from the paper.	'l1'

Source code in oriented_det/ops/probiou.py

def probiou_loss_per_box(
    pred: Tensor,
    target: Tensor,
    *,
    eps: float = 1e-3,
    mode: ProbIoUMode = "l1",
) -> Tensor:
    """Per-instance ProbIoU loss ``[N]`` (lower is better overlap).

    Args:
        pred: ``[N, 5]`` predicted boxes ``(cx, cy, w, h, angle_rad)``.
        target: ``[N, 5]`` target boxes (same layout).
        eps: Stabilizer for divisions and logarithms.
        mode: ``\"l1\"`` → bounded metric in ``[0, 1]``; ``\"l2\"`` → unbounded
            ``-log(1 - l1^2)`` variant from the paper.
    """
    if pred.shape != target.shape or pred.shape[-1] != 5:
        raise ValueError(f"Expected matching [N, 5] tensors, got {pred.shape} vs {target.shape}")
    mode_norm = (mode or "l1").strip().lower()
    if mode_norm not in ("l1", "l2"):
        raise ValueError(f"mode must be 'l1' or 'l2', got {mode!r}")

    with _fp32_no_autocast_ctx(pred):
        pred_f = pred.float()
        target_f = target.float()

        gbboxes1 = gbb_form(pred_f)
        gbboxes2 = gbb_form(target_f)

        x1, y1, a1_, b1_, c1_ = (
            gbboxes1[:, 0],
            gbboxes1[:, 1],
            gbboxes1[:, 2],
            gbboxes1[:, 3],
            gbboxes1[:, 4],
        )
        x2, y2, a2_, b2_, c2_ = (
            gbboxes2[:, 0],
            gbboxes2[:, 1],
            gbboxes2[:, 2],
            gbboxes2[:, 3],
            gbboxes2[:, 4],
        )

        a1, b1, c1 = rotated_form(a1_, b1_, c1_)
        a2, b2, c2 = rotated_form(a2_, b2_, c2_)

        denom_base = (a1 + a2) * (b1 + b2) - torch.pow(c1 + c2, 2) + eps
        t1 = (
            ((a1 + a2) * torch.pow(y1 - y2, 2) + (b1 + b2) * torch.pow(x1 - x2, 2))
            / denom_base
        ) * 0.25
        t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / denom_base) * 0.5
        det1 = (a1 * b1 - torch.pow(c1, 2)).clamp(min=eps)
        det2 = (a2 * b2 - torch.pow(c2, 2)).clamp(min=eps)
        t3 = (
            torch.log(
                denom_base
                / (4.0 * torch.sqrt(det1 * det2) + eps)
                + eps
            )
            * 0.5
        )

        b_d = torch.clamp(t1 + t2 + t3, min=eps, max=100.0)
        l1 = torch.sqrt(1.0 - torch.exp(-b_d) + eps)
        if mode_norm == "l1":
            probiou = l1
        else:
            l_i = torch.pow(l1, 2.0)
            probiou = -torch.log(1.0 - l_i + eps)

    return probiou.to(dtype=pred.dtype)

qbox_iou(box_a, box_b, *, intersection_backend='auto')

Compute IoU between two quadrilateral boxes.

Parameters:

Name	Type	Description	Default
box_a	QBox | Iterable[Sequence[float]]	First quadrilateral box	required
box_b	QBox | Iterable[Sequence[float]]	Second quadrilateral box	required
intersection_backend	str	Backend for intersection calculation ('auto', 'python', or 'shapely')	'auto'

Returns:

Type	Description
float	IoU value as a float

Source code in oriented_det/ops/iou.py

def qbox_iou(
    box_a: QBox | Iterable[Sequence[float]], 
    box_b: QBox | Iterable[Sequence[float]],
    *,
    intersection_backend: str = "auto"
) -> float:
    """Compute IoU between two quadrilateral boxes.

    Args:
        box_a: First quadrilateral box
        box_b: Second quadrilateral box
        intersection_backend: Backend for intersection calculation ('auto', 'python', or 'shapely')

    Returns:
        IoU value as a float
    """
    return polygon_iou(box_a, box_b, intersection_backend=intersection_backend)

rbox_iou(box_a, box_b, *, intersection_backend='auto', use_aabb_prefilter=True, backend=None)

Compute IoU between two rotated boxes.

This is a CPU Python implementation. For GPU-accelerated IoU, use: from oriented_det.ops.gpu_ops import oriented_box_iou_gpu

Parameters:

Name	Type	Description	Default
box_a	RBox | Sequence[float]	First rotated box	required
box_b	RBox | Sequence[float]	Second rotated box	required
intersection_backend	str	Backend for intersection calculation ('auto', 'python', or 'shapely')	'auto'
use_aabb_prefilter	bool	If True (default), use fast AABB pre-filtering to skip expensive polygon intersection for non-overlapping boxes.	True
backend	Optional[str]	Alias for intersection_backend (for API compatibility).	None

Returns:

Type	Description
float	IoU value as a float

Source code in oriented_det/ops/iou.py

def rbox_iou(
    box_a: RBox | Sequence[float],
    box_b: RBox | Sequence[float],
    *,
    intersection_backend: str = "auto",
    use_aabb_prefilter: bool = True,
    backend: Optional[str] = None,
) -> float:
    """Compute IoU between two rotated boxes.

    This is a CPU Python implementation. For GPU-accelerated IoU, use:
    `from oriented_det.ops.gpu_ops import oriented_box_iou_gpu`

    Args:
        box_a: First rotated box
        box_b: Second rotated box
        intersection_backend: Backend for intersection calculation 
            ('auto', 'python', or 'shapely')
        use_aabb_prefilter: If True (default), use fast AABB pre-filtering to skip
            expensive polygon intersection for non-overlapping boxes.
        backend: Alias for intersection_backend (for API compatibility).

    Returns:
        IoU value as a float
    """
    if backend is not None:
        if backend == "invalid":
            raise ValueError("Invalid backend: invalid")
        if backend == "torch":
            from ..ops import utils as ops_utils
            if getattr(ops_utils, "TORCH_BOX_IOU_ROTATED", None) is None:
                raise RuntimeError("torch backend for rbox_iou is not available")
        intersection_backend = backend
    # Fast AABB pre-filter: if axis-aligned boxes don't overlap,
    # rotated boxes cannot overlap either (avoids expensive polygon intersection)
    if use_aabb_prefilter:
        rb_a = as_rbox(box_a)
        rb_b = as_rbox(box_b)
        aabb_a = rb_a.axis_aligned_bounds()
        aabb_b = rb_b.axis_aligned_bounds()
        if not _aabb_overlaps(aabb_a, aabb_b):
            return 0.0  # No overlap possible

    return polygon_iou(box_a, box_b, intersection_backend=intersection_backend)

xy_wh_r_to_xy_sigma(xywhr)

Map oriented boxes to Gaussian mean xy and covariance sigma.

Parameters:

Name	Type	Description	Default
xywhr	Tensor	[..., 5] with (cx, cy, w, h, angle).	required

Returns:

Type	Description
Tuple[Tensor, Tensor]	xy with shape [..., 2], sigma with shape [..., 2, 2].

Source code in oriented_det/ops/kfiou.py

def xy_wh_r_to_xy_sigma(xywhr: Tensor) -> Tuple[Tensor, Tensor]:
    """Map oriented boxes to Gaussian mean ``xy`` and covariance ``sigma``.

    Args:
        xywhr: ``[..., 5]`` with ``(cx, cy, w, h, angle)``.

    Returns:
        ``xy`` with shape ``[..., 2]``, ``sigma`` with shape ``[..., 2, 2]``.
    """
    shape = xywhr.shape
    if shape[-1] != 5:
        raise ValueError(f"Expected last dim 5 (xywhr), got {shape[-1]}")
    xy = xywhr[..., :2]
    wh = xywhr[..., 2:4].clamp(min=1e-7, max=1e7).reshape(-1, 2)
    r = xywhr[..., 4].reshape(-1)
    cos_r = torch.cos(r)
    sin_r = torch.sin(r)
    rot = torch.stack((cos_r, -sin_r, sin_r, cos_r), dim=-1).reshape(-1, 2, 2)
    s = 0.5 * torch.diag_embed(wh)
    sigma = rot.bmm(s.square()).bmm(rot.transpose(-2, -1)).reshape(shape[:-1] + (2, 2))
    return xy, sigma

Important Notes

NMS Performance

Important: torchvision.ops.nms_rotated does not exist in current torchvision versions (tested with 0.22.0+).

Current Implementation: - The framework uses a Python-based NMS implementation with optimizations: - AABB (axis-aligned bounding box) pre-filtering before expensive rotated IoU computation - Batch IoU computation for better performance - Class-aware NMS support (boxes of different classes don't suppress each other) - Performance: ~2.6x faster than naive implementation (11.9s for 2000 boxes vs 31.5s) - Future: GPU-accelerated kernels via custom CUDA implementation planned

Optimization Details: - AABB pre-filtering eliminates ~80-90% of rotated IoU computations - Pre-computed AABBs avoid redundant calculations - Batch processing reduces Python overhead

Backend Selection

For CPU-based IoU, use iou.rbox_iou() and nms.oriented_nms() with intersection_backend ("auto", "python", or "shapely"). For GPU-accelerated operations, use oriented_det.ops.gpu_ops (e.g., oriented_box_iou_gpu, oriented_nms_gpu). Models automatically use GPU kernels when available.

Examples

Basic IoU Computation

from oriented_det.geometry import RBox
from oriented_det.ops import iou

boxes = [
    RBox(0, 0, 2, 1, 0),
    RBox(0.5, 0.1, 2, 1, 0),
    RBox(4, 0, 2, 1, 0),
]

# Compute IoU between two boxes
overlap = iou.rbox_iou(boxes[0], boxes[1])

# Batch IoU computation (more efficient)
iou_matrix = iou.batch_rbox_iou(boxes, boxes)

Basic NMS

from oriented_det.geometry import RBox
from oriented_det.ops import nms

boxes = [
    RBox(0, 0, 2, 1, 0),
    RBox(0.5, 0.1, 2, 1, 0),  # Overlaps with first
    RBox(4, 0, 2, 1, 0),  # No overlap
]
scores = [0.9, 0.8, 0.6]

# Apply NMS
keep = nms.oriented_nms(boxes, scores, iou_threshold=0.3)

Class-Aware NMS

# Boxes of different classes don't suppress each other
labels = [0, 0, 1]  # First two boxes are class 0, third is class 1
keep = nms.oriented_nms(boxes, scores, iou_threshold=0.3, labels=labels)

GPU Operations

For maximum performance with large batches, use the GPU-accelerated operations:

GPU IoU Computation

import torch
from oriented_det.ops.gpu_ops import oriented_box_iou_gpu

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Input format: [N, 5] tensors with [cx, cy, w, h, angle]
boxes1 = torch.tensor([
    [100, 100, 50, 30, 0.0],
    [200, 150, 40, 60, 0.5],
], device=device)

boxes2 = torch.tensor([
    [105, 105, 50, 30, 0.0],
    [300, 300, 50, 50, 0.0],
], device=device)

# Compute IoU matrix: [N, M]
iou_matrix = oriented_box_iou_gpu(boxes1, boxes2, num_samples=100)

GPU Anchor Generation

import torch
from oriented_det.ops.gpu_ops import generate_oriented_anchors_gpu
import math

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Generate anchors for FPN levels
anchors_per_level = generate_oriented_anchors_gpu(
    image_size=(800, 800),
    feature_map_sizes=[(200, 200), (100, 100), (50, 50)],
    anchor_scales=[8],
    anchor_ratios=[0.5, 1.0, 2.0],
    anchor_angles=[-math.pi/2, 0, math.pi/2],
    stride_per_level=[4, 8, 16],
    device=device
)
# Returns list of [H*W*A, 5] anchor tensors

GPU Anchor Matching

import torch
from oriented_det.ops.gpu_ops import match_anchors_to_gt_gpu

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

anchors = torch.randn(1000, 5, device=device)  # [N, 5]
gt_boxes = torch.randn(10, 5, device=device)  # [M, 5]

labels, matched_gt_indices = match_anchors_to_gt_gpu(
    anchors=anchors,
    gt_boxes=gt_boxes,
    positive_iou_threshold=0.7,
    negative_iou_threshold=0.3,
)
# labels: [N] with -1=ignore, 0=background, 1=foreground
# matched_gt_indices: [N] index of matched GT (-1 if none)

GPU NMS

import torch
from oriented_det.ops.gpu_ops import oriented_nms_gpu

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

boxes = torch.randn(1000, 5, device=device)  # [N, 5]
scores = torch.rand(1000, device=device)  # [N]

keep_indices = oriented_nms_gpu(
    boxes=boxes,
    scores=scores,
    iou_threshold=0.5,
    max_detections=100,
)
# Returns tensor of kept box indices

See Also: - GPU Operations Guide - Detailed usage guide