Keep batch dimension

Author: AntoineSimoulin

torchvision/ops/boxes.py

@@ -318,19 +318,25 @@ def _box_inter_union(boxes1: Tensor, boxes2: Tensor, fmt: str = "xyxy") -> tuple
         raise ValueError(f"Unsupported Box IoU Calculation for given fmt {format}.")
 
     if fmt == "xyxy":
-        lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
-        rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
+        lt = torch.max(boxes1[..., None, :2], boxes2[..., None, :, :2])  # [...,N,M,2]
+        rb = torch.min(boxes1[..., None, 2:], boxes2[..., None, :, 2:])  # [...,N,M,2]
     elif fmt == "xywh":
-        lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
-        rb = torch.min(boxes1[:, None, :2] + boxes1[:, None, 2:], boxes2[:, :2] + boxes2[:, 2:])  # [N,M,2]
+        lt = torch.max(boxes1[..., None, :2], boxes2[..., None, :, :2])  # [...,N,M,2]
+        rb = torch.min(
+            boxes1[..., None, :2] + boxes1[..., None, 2:], boxes2[..., None, :2] + boxes2[..., None, 2:]
+        )  # [N,M,2]
     else:  # fmt == "cxcywh":
-        lt = torch.max(boxes1[:, None, :2] - boxes1[:, None, 2:] / 2, boxes2[:, :2] - boxes2[:, 2:] / 2)  # [N,M,2]
-        rb = torch.min(boxes1[:, None, :2] + boxes1[:, None, 2:] / 2, boxes2[:, :2] + boxes2[:, 2:] / 2)  # [N,M,2]
+        lt = torch.max(
+            boxes1[..., None, :2] - boxes1[..., None, 2:] / 2, boxes2[..., None, :2] - boxes2[..., None, 2:] / 2
+        )  # [N,M,2]
+        rb = torch.min(
+            boxes1[..., None, :2] + boxes1[..., None, 2:] / 2, boxes2[..., None, :2] + boxes2[..., None, 2:] / 2
+        )  # [N,M,2]
 
-    wh = _upcast(rb - lt).clamp(min=0)  # [N,M,2]
-    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+    wh = _upcast(rb - lt).clamp(min=0)  # [N,M,2]	    wh = _upcast(rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[..., 0] * wh[..., 1]  # [N,M]	    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
 
-    union = area1[:, None] + area2 - inter
+    union = area1[..., None] + area2[..., None, :] - inter
 
     return inter, union