chore: import upstream snapshot with attribution

tests/models/test_criterion_keypoints.py

@@ -0,0 +1,161 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for keypoint losses in SetCriterion."""
+
+import torch
+
+from rfdetr.models.criterion import SetCriterion
+
+
+class _MatcherStub:
+    """Matcher stub used to avoid depending on Hungarian matching internals."""
+
+    def __call__(self, outputs, targets, group_detr=1):
+        indices = []
+        for target in targets:
+            num_targets = int(target["labels"].shape[0])
+            idx = torch.arange(num_targets, dtype=torch.int64)
+            indices.append((idx, idx))
+        return indices
+
+
+def _make_outputs(
+    batch_size: int,
+    num_queries: int,
+    num_keypoints: int,
+) -> dict[str, torch.Tensor]:
+    return {
+        "pred_logits": torch.zeros(batch_size, num_queries, 2),
+        "pred_boxes": torch.rand(batch_size, num_queries, 4).clamp(0.05, 0.95),
+        "pred_keypoints": torch.randn(batch_size, num_queries, num_keypoints, 8),
+    }
+
+
+def test_loss_keypoints_list_of_dicts_targets() -> None:
+    """Keypoint loss should consume list-of-dicts targets used by public training."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.cat(
+                [
+                    torch.rand(1, 17, 2),
+                    torch.full((1, 17, 1), 2.0),
+                ],
+                dim=-1,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert "loss_keypoints_l1" in losses
+    assert "loss_keypoints_findable" in losses
+    assert "loss_keypoints_visible" in losses
+    assert "loss_keypoints_nll" in losses
+    assert all(torch.isfinite(value) for value in losses.values())
+
+
+def test_loss_keypoints_empty_targets() -> None:
+    """Empty target batches should produce finite zero-valued keypoint losses."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.zeros((0,), dtype=torch.int64),
+            "boxes": torch.zeros((0, 4), dtype=torch.float32),
+            "keypoints": torch.zeros((0, 17, 3), dtype=torch.float32),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].item() == 0.0
+    assert losses["loss_keypoints_findable"].item() == 0.0
+    assert losses["loss_keypoints_visible"].item() == 0.0
+    assert losses["loss_keypoints_nll"].item() == 0.0
+
+
+def test_loss_keypoints_person_schema_shape() -> None:
+    """Person-only schema `[17]` should be consumed without shape mismatches."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=2, num_queries=2, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+
+
+def test_loss_keypoints_multiclass_schema_kmax_targets() -> None:
+    """Heterogeneous keypoint classes should consume Kmax-padded targets."""
+    criterion = SetCriterion(
+        num_classes=3,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[2, 1],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=2, num_keypoints=4)
+    targets = [
+        {
+            "labels": torch.tensor([0, 1], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4], [0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.tensor(
+                [
+                    [[0.2, 0.3, 2.0], [0.4, 0.5, 2.0]],
+                    [[0.6, 0.7, 2.0], [0.0, 0.0, 0.0]],
+                ],
+                dtype=torch.float32,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+    assert all(torch.isfinite(value) for value in losses.values())