16031aae96
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
175 lines
8.4 KiB
Python
175 lines
8.4 KiB
Python
# ------------------------------------------------------------------------
|
|
# 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 decoding in PostProcess."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from rfdetr.models.postprocess import PostProcess
|
|
|
|
|
|
def test_postprocess_keypoints_shape_and_scores() -> None:
|
|
"""PostProcess should emit keypoints and raw precision parameters for top detections."""
|
|
postprocess = PostProcess(num_select=2, num_keypoints_per_class=[17])
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5], [0.4, 0.6, 0.2, 0.3]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 2, 17, 8), dtype=torch.float32),
|
|
}
|
|
outputs["pred_keypoints"][0, :, :, 0] = 0.5
|
|
outputs["pred_keypoints"][0, :, :, 1] = 0.25
|
|
outputs["pred_keypoints"][0, :, :, 2] = 3.0
|
|
outputs["pred_keypoints"][0, :, :, 4] = 0.25
|
|
outputs["pred_keypoints"][0, :, :, 5] = 0.5
|
|
outputs["pred_keypoints"][0, :, :, 6] = -0.25
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
keypoints = results[0]["keypoints"]
|
|
keypoint_precision = results[0]["keypoint_precision_cholesky"]
|
|
|
|
assert keypoints.shape == (2, 17, 3)
|
|
assert torch.allclose(keypoints[:, :, 0], torch.full((2, 17), 100.0))
|
|
assert torch.allclose(keypoints[:, :, 1], torch.full((2, 17), 25.0))
|
|
assert torch.all((keypoints[:, :, 2] > 0) & (keypoints[:, :, 2] < 1))
|
|
assert keypoint_precision.shape == (2, 17, 3)
|
|
torch.testing.assert_close(keypoint_precision[:, :, 0], torch.full((2, 17), 0.25))
|
|
torch.testing.assert_close(keypoint_precision[:, :, 1], torch.full((2, 17), 0.5))
|
|
torch.testing.assert_close(keypoint_precision[:, :, 2], torch.full((2, 17), -0.25))
|
|
|
|
|
|
def test_postprocess_keypoints_class_filtering() -> None:
|
|
"""Class-specific keypoint slots should be selected from padded per-class keypoint tensors."""
|
|
postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1])
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[0.0, 10.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
|
|
}
|
|
# class 0 slots: [0, 1], class 1 slots: [2, 3]
|
|
outputs["pred_keypoints"][0, 0, 2, 0] = 0.25
|
|
outputs["pred_keypoints"][0, 0, 2, 1] = 0.4
|
|
outputs["pred_keypoints"][0, 0, 2, 2] = 2.0
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
keypoints = results[0]["keypoints"]
|
|
keypoint_precision = results[0]["keypoint_precision_cholesky"]
|
|
|
|
assert keypoints.shape == (1, 2, 3)
|
|
assert torch.allclose(keypoints[0, 0, 0], torch.tensor(50.0))
|
|
assert torch.allclose(keypoints[0, 0, 1], torch.tensor(40.0))
|
|
assert 0.0 < keypoints[0, 0, 2].item() < 1.0
|
|
torch.testing.assert_close(keypoints[0, 1], torch.zeros(3))
|
|
torch.testing.assert_close(keypoint_precision[0, 1], torch.full((3,), float("nan")), equal_nan=True)
|
|
|
|
|
|
def test_postprocess_keypoints_trace_alpha_rescores_active_keypoints_only() -> None:
|
|
"""Trace fusion should use active keypoints for the predicted class and ignore padded slots."""
|
|
postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1], trace_alpha=1.0)
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[-10.0, 0.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
|
|
}
|
|
# class 1 has one active slot at flat index 2 and one padded inactive slot at flat index 3.
|
|
outputs["pred_keypoints"][0, 0, 2, 2] = 10.0
|
|
outputs["pred_keypoints"][0, 0, 2, 4] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 2, 5] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 2, 6] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 3, 2] = 10.0
|
|
outputs["pred_keypoints"][0, 0, 3, 4] = -2.0
|
|
outputs["pred_keypoints"][0, 0, 3, 6] = -2.0
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
|
|
expected_score = torch.tensor([0.2], dtype=torch.float32)
|
|
torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
|
|
|
|
|
|
def test_postprocess_keypoints_trace_alpha_normalizes_large_fused_scores() -> None:
|
|
"""Trace-fused keypoint scores should be bounded after empirical normalization."""
|
|
postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[10.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
|
|
}
|
|
outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
|
|
outputs["pred_keypoints"][0, 0, 0, 4] = 2.0
|
|
outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 0, 6] = 2.0
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
|
|
original_score = torch.sigmoid(torch.tensor([10.0], dtype=torch.float32))
|
|
mean_trace = torch.tensor([2.0], dtype=torch.float32) * torch.exp(torch.tensor([-4.0], dtype=torch.float32))
|
|
fused_score = original_score * mean_trace.pow(-1.0)
|
|
expected_score = fused_score / (1.0 + fused_score)
|
|
assert fused_score.item() > 1.0
|
|
assert 0.0 < results[0]["scores"].item() < 1.0
|
|
torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
|
|
|
|
|
|
def test_postprocess_keypoints_trace_alpha_clamps_overflowing_fused_scores() -> None:
|
|
"""Trace fusion should stay finite and strictly below 1.0 when the raw fused score overflows."""
|
|
postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
|
|
}
|
|
outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
|
|
outputs["pred_keypoints"][0, 0, 0, 4] = 50.0
|
|
outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 0, 6] = 50.0
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
|
|
score = results[0]["scores"]
|
|
expected_score = torch.nextafter(torch.ones_like(score), torch.zeros_like(score))
|
|
assert torch.isfinite(score).all()
|
|
assert 0.0 < score.item() < 1.0
|
|
torch.testing.assert_close(score, expected_score, rtol=0.0, atol=0.0)
|
|
|
|
|
|
def test_postprocess_keypoints_trace_alpha_uses_log_space_for_extreme_trace() -> None:
|
|
"""Trace fusion should stay finite for extreme covariance terms."""
|
|
postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1])
|
|
outputs = {
|
|
"pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
|
|
"pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
|
|
"pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
|
|
}
|
|
outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
|
|
outputs["pred_keypoints"][0, 0, 0, 4] = -50.0
|
|
outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
|
|
outputs["pred_keypoints"][0, 0, 0, 6] = 0.0
|
|
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
results = postprocess(outputs, target_sizes)
|
|
|
|
expected_score = torch.tensor([0.5], dtype=torch.float32) * torch.exp(torch.tensor([-20.0], dtype=torch.float32))
|
|
torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-12)
|
|
|
|
|
|
def test_postprocess_validate_outputs_raises_when_masks_and_keypoints_both_present() -> None:
|
|
"""PostProcess should raise ValueError when both pred_masks and pred_keypoints are present."""
|
|
postprocess = PostProcess(num_select=10)
|
|
outputs = {
|
|
"pred_logits": torch.zeros((1, 2, 2)),
|
|
"pred_boxes": torch.zeros((1, 2, 4)),
|
|
"pred_masks": torch.zeros((1, 2, 4, 4)),
|
|
"pred_keypoints": torch.zeros((1, 2, 17, 8)),
|
|
}
|
|
target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
|
|
|
|
with pytest.raises(ValueError, match="cannot be used together"):
|
|
postprocess(outputs, target_sizes)
|