Add inference tests for RF-DETR+ XLarge and 2XLarge models with parameterized resolutions

Author: Borda

tests/test_version.py

@@ -8,15 +8,22 @@
 import numpy as np
 import pytest
 
-from rfdetr_plus import RFDETRXLarge
+from rfdetr_plus import RFDETR2XLarge, RFDETRXLarge
 
 
 @pytest.mark.gpu
-def test_model_inference() -> None:
-    """Test that we can instantiate RF-DETR+ XLarge and run inference."""
+@pytest.mark.parametrize(
+    ("model_class", "resolution"),
+    [
+        (RFDETRXLarge, 700),
+        (RFDETR2XLarge, 880),
+    ],
+)
+def test_model_inference(model_class, resolution) -> None:
+    """Test that we can instantiate RF-DETR+ models and run inference."""
     # Instantiate and run inference
-    rf_detr = RFDETRXLarge()
-    dummy_image = np.random.randint(0, 255, (700, 700, 3), dtype=np.uint8)
+    rf_detr = model_class()
+    dummy_image = np.random.randint(0, 255, (resolution, resolution, 3), dtype=np.uint8)
 
     # Run inference - this verifies the model can be instantiated and used
     predictions = rf_detr.predict(dummy_image, conf_threshold=0.1)