test: ensure pipeline requests respect the classification_num_predictions config

Author: mihow

trapdata/api/tests/test_api.py

@@ -5,6 +5,7 @@
 from fastapi.testclient import TestClient
 
 from trapdata.api.api import (
+    PIPELINE_CHOICES,
     PipelineChoice,
     PipelineConfig,
     PipelineRequest,
@@ -13,6 +14,7 @@
     app,
 )
 from trapdata.api.tests.image_server import StaticFileTestServer
+from trapdata.ml.models.classification import SpeciesClassifier
 from trapdata.tests import TEST_IMAGES_BASE_PATH
 
 logging.basicConfig(level=logging.INFO)
@@ -22,7 +24,7 @@
 class TestInferenceAPI(TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.test_images_dir = pathlib.Path(TEST_IMAGES_BASE_PATH) / "vermont"
+        cls.test_images_dir = pathlib.Path(TEST_IMAGES_BASE_PATH)
         if not cls.test_images_dir.exists():
             raise FileNotFoundError(
                 f"Test images directory not found: {cls.test_images_dir}"
@@ -34,25 +36,81 @@ def setUpClass(cls):
     def setUp(self):
         self.file_server = StaticFileTestServer(self.test_images_dir)
 
+    def get_test_images(self, subdir: str = "vermont", num: int = 2):
+        images_dir = self.test_images_dir / subdir
+        source_image_urls = [
+            self.file_server.get_url(f.relative_to(images_dir))
+            for f in self.test_images_dir.glob("*.jpg")
+        ][:num]
+        source_images = [
+            SourceImageRequest(id=str(i), url=url)
+            for i, url in enumerate(source_image_urls)
+        ]
+        return source_images
+
+    def get_test_pipeline(self, slug: str = "quebec_vermont_moths_2023"):
+        pipeline = PIPELINE_CHOICES[slug]
+        return pipeline
+
     def test_pipeline_request(self):
+        """
+        Ensure that the pipeline accepts a valid request and returns a valid response.
+        """
+        pipeline_request = PipelineRequest(
+            pipeline=PipelineChoice["quebec_vermont_moths_2023"],
+            source_images=self.get_test_images(num=2),
+        )
         with self.file_server:
-            num_images = 2
-            source_image_urls = [
-                self.file_server.get_url(f.relative_to(self.test_images_dir))
-                for f in self.test_images_dir.glob("*.jpg")
-            ][:num_images]
-            source_images = [
-                SourceImageRequest(id=str(i), url=url)
-                for i, url in enumerate(source_image_urls)
-            ]
-            pipeline_request = PipelineRequest(
-                pipeline=PipelineChoice["quebec_vermont_moths_2023"],
-                source_images=source_images,
-                config=PipelineConfig(classification_num_predictions=1),
-            )
             response = self.client.post(
                 "/pipeline/process", json=pipeline_request.dict()
             )
             assert response.status_code == 200
+            PipelineResponse(**response.json())
+
+    def test_config_num_classification_predictions(self):
+        """
+        Test that the pipeline respects the `classification_num_predictions` configuration.
+
+        If the configuration is set to a number, the pipeline should return that number of labels/scores per prediction.
+        If the configuration is set to `None`, the pipeline should return all labels/scores per prediction.
+        """
+        test_images = self.get_test_images(num=1)
+        test_pipeline_slug = "quebec_vermont_moths_2023"
+        terminal_classifier: SpeciesClassifier = self.get_test_pipeline(
+            test_pipeline_slug
+        )
+
+        def _send_request(classification_num_predictions: int | None):
+            config = PipelineConfig(
+                classification_num_predictions=classification_num_predictions
+            )
+            pipeline_request = PipelineRequest(
+                pipeline=PipelineChoice[test_pipeline_slug],
+                source_images=test_images,
+                config=config,
+            )
+            with self.file_server:
+                response = self.client.post(
+                    "/pipeline/process", json=pipeline_request.dict()
+                )
+            assert response.status_code == 200
             pipeline_response = PipelineResponse(**response.json())
-            assert len(pipeline_response.detections) > 0
+            terminal_classifications = [
+                classification
+                for detection in pipeline_response.detections
+                for classification in detection.classifications
+                if classification.terminal
+            ]
+            for classification in terminal_classifications:
+                if classification_num_predictions is None:
+                    # Ensure that a score is returned for every possible class
+                    assert len(classification.labels) == terminal_classifier.num_classes
+                    assert len(classification.scores) == terminal_classifier.num_classes
+                else:
+                    # Ensure that the number of predictions is limited to the number specified
+                    # There may be fewer predictions than the number specified if there are fewer classes.
+                    assert len(classification.labels) <= classification_num_predictions
+                    assert len(classification.scores) <= classification_num_predictions
+
+        _send_request(classification_num_predictions=1)
+        _send_request(classification_num_predictions=None)