Support re-processing detections and skipping localizer

ami/main/models.py

@@ -31,6 +31,7 @@
 from ami.base.models import BaseModel
 from ami.main import charts
 from ami.main.models_future.projects import ProjectSettingsMixin
+from ami.ml.schemas import BoundingBox
 from ami.users.models import User
 from ami.utils.media import calculate_file_checksum, extract_timestamp
 from ami.utils.schemas import OrderedEnum
@@ -2263,6 +2264,17 @@ class Detection(BaseModel):
     source_image_id: int
     detection_algorithm_id: int
 
+    def get_bbox(self):
+        if self.bbox:
+            return BoundingBox(
+                x1=self.bbox[0],
+                y1=self.bbox[1],
+                x2=self.bbox[2],
+                y2=self.bbox[3],
+            )
+        else:
+            return None
+
     # def bbox(self):
     #     return (
     #         self.bbox_x,

ami/ml/models/pipeline.py

@@ -39,7 +39,9 @@
 from ami.ml.models.algorithm import Algorithm, AlgorithmCategoryMap
 from ami.ml.schemas import (
     AlgorithmConfigResponse,
+    AlgorithmReference,
     ClassificationResponse,
+    DetectionRequest,
     DetectionResponse,
     PipelineRequest,
     PipelineRequestConfigParameters,
@@ -62,6 +64,7 @@ def filter_processed_images(
     Return only images that need to be processed by a given pipeline.
     An image needs processing if:
     1. It has no detections from the pipeline's detection algorithm
+    or
     2. It has detections but they don't have classifications from all the pipeline's classification algorithms
     """
     pipeline_algorithms = pipeline.algorithms.all()
@@ -192,25 +195,45 @@ def process_images(
     task_logger.info(f"Sending {len(images)} images to Pipeline {pipeline}")
     urls = [source_image.public_url() for source_image in images if source_image.public_url()]
 
-    source_images = [
-        SourceImageRequest(
-            id=str(source_image.pk),
-            url=url,
-        )
-        for source_image, url in zip(images, urls)
-        if url
-    ]
+    source_image_requests: list[SourceImageRequest] = []
+    detection_requests: list[DetectionRequest] = []
+
+    for source_image, url in zip(images, urls):
+        if url:
+            source_image_request = SourceImageRequest(
+                id=str(source_image.pk),
+                url=url,
+            )
+            source_image_requests.append(source_image_request)
+            # Re-process all existing detections if they exist
+            for detection in source_image.detections.all():
+                bbox = detection.get_bbox()
+                if bbox and detection.detection_algorithm:
+                    detection_requests.append(
+                        DetectionRequest(
+                            source_image=source_image_request,
+                            bbox=bbox,
+                            crop_image_url=detection.url(),
+                            algorithm=AlgorithmReference(
+                                name=detection.detection_algorithm.name,
+                                key=detection.detection_algorithm.key,
+                            ),
+                        )
+                    )
 
     if not project_id:
         task_logger.warning(f"Pipeline {pipeline} is not associated with a project")
 
     config = pipeline.get_config(project_id=project_id)
     task_logger.info(f"Using pipeline config: {config}")
 
+    task_logger.info(f"Found {len(detection_requests)} existing detections.")
+
     request_data = PipelineRequest(
         pipeline=pipeline.slug,
-        source_images=source_images,
+        source_images=source_image_requests,
         config=config,
+        detections=detection_requests,
     )
 
     session = create_session()
@@ -230,7 +253,8 @@ def process_images(
             pipeline=pipeline.slug,
             total_time=0,
             source_images=[
-                SourceImageResponse(id=source_image.id, url=source_image.url) for source_image in source_images
+                SourceImageResponse(id=source_image_request.id, url=source_image_request.url)
+                for source_image_request in source_image_requests
             ],
             detections=[],
             errors=msg,
@@ -351,23 +375,12 @@ def get_or_create_detection(
     serialized_bbox = list(detection_resp.bbox.dict().values())
     detection_repr = f"Detection {detection_resp.source_image_id} {serialized_bbox}"
 
-    assert detection_resp.algorithm, f"No detection algorithm was specified for detection {detection_repr}"
-    try:
-        detection_algo = algorithms_used[detection_resp.algorithm.key]
-    except KeyError:
-        raise ValueError(
-            f"Detection algorithm {detection_resp.algorithm.key} is not a known algorithm. "
-            "The processing service must declare it in the /info endpoint. "
-            f"Known algorithms: {list(algorithms_used.keys())}"
-        )
-
     assert str(detection_resp.source_image_id) == str(
         source_image.pk
     ), f"Detection belongs to a different source image: {detection_repr}"
 
     existing_detection = Detection.objects.filter(
         source_image=source_image,
-        detection_algorithm=detection_algo,
         bbox=serialized_bbox,
     ).first()
 
@@ -387,6 +400,16 @@ def get_or_create_detection(
         detection = existing_detection
 
     else:
+        assert detection_resp.algorithm, f"No detection algorithm was specified for detection {detection_repr}"
+        try:
+            detection_algo = algorithms_used[detection_resp.algorithm.key]
+        except KeyError:
+            raise ValueError(
+                f"Detection algorithm {detection_resp.algorithm.key} is not a known algorithm. "
+                "The processing service must declare it in the /info endpoint. "
+                f"Known algorithms: {list(algorithms_used.keys())}"
+            )
+
         new_detection = Detection(
             source_image=source_image,
             bbox=serialized_bbox,
@@ -1007,7 +1030,7 @@ def collect_images(
         )
 
     def choose_processing_service_for_pipeline(
-        self, job_id: int, pipeline_name: str, project_id: int
+        self, job_id: int | None, pipeline_name: str, project_id: int
     ) -> ProcessingService:
         # @TODO use the cached `last_checked_latency` and a max age to avoid checking every time
 

ami/ml/schemas.py

@@ -112,16 +112,6 @@ class ClassificationResponse(pydantic.BaseModel):
     timestamp: datetime.datetime
 
 
-class DetectionResponse(pydantic.BaseModel):
-    source_image_id: str
-    bbox: BoundingBox
-    inference_time: float | None = None
-    algorithm: AlgorithmReference
-    timestamp: datetime.datetime
-    crop_image_url: str | None = None
-    classifications: list[ClassificationResponse] = []
-
-
 class SourceImageRequest(pydantic.BaseModel):
     # @TODO bring over new SourceImage & b64 validation from the lepsAI repo
     id: str
@@ -144,6 +134,23 @@ class Config:
 ]
 
 
+class DetectionRequest(pydantic.BaseModel):
+    source_image: SourceImageRequest  # the 'original' image
+    bbox: BoundingBox
+    crop_image_url: str | None = None
+    algorithm: AlgorithmReference
+
+
+class DetectionResponse(pydantic.BaseModel):
+    source_image_id: str
+    bbox: BoundingBox
+    inference_time: float | None = None
+    algorithm: AlgorithmReference
+    timestamp: datetime.datetime
+    crop_image_url: str | None = None
+    classifications: list[ClassificationResponse] = []
+
+
 class PipelineRequestConfigParameters(dict):
     """Parameters used to configure a pipeline request.
 
@@ -166,6 +173,7 @@ class PipelineRequestConfigParameters(dict):
 class PipelineRequest(pydantic.BaseModel):
     pipeline: str
     source_images: list[SourceImageRequest]
+    detections: list[DetectionRequest] | None = None
     config: PipelineRequestConfigParameters | dict | None = None
 
 

ami/ml/tests.py

@@ -158,7 +158,7 @@ def test_created_category_maps(self):
 
     def test_alignment_of_predictions_and_category_map(self):
         # Ensure that the scores and labels are aligned
-        pipeline = self.processing_service_instance.pipelines.all().get(slug="random")
+        pipeline = self.processing_service_instance.pipelines.all().get(slug="random-detection-random-species")
         pipeline_response = pipeline.process_images(self.test_images, project_id=self.project.pk)
         results = save_results(pipeline_response, return_created=True)
         assert results is not None, "Expected results to be returned in a PipelineSaveResults object"
@@ -172,7 +172,7 @@ def test_alignment_of_predictions_and_category_map(self):
 
     def test_top_n_alignment(self):
         # Ensure that the top_n parameter works
-        pipeline = self.processing_service_instance.pipelines.all().get(slug="random")
+        pipeline = self.processing_service_instance.pipelines.all().get(slug="random-detection-random-species")
         pipeline_response = pipeline.process_images(self.test_images, project_id=self.project.pk)
         results = save_results(pipeline_response, return_created=True)
         assert results is not None, "Expecected results to be returned in a PipelineSaveResults object"
@@ -182,6 +182,60 @@ def test_top_n_alignment(self):
             assert classification.score == top_n[0]["score"]
             assert classification.taxon == top_n[0]["taxon"]
 
+    def test_pipeline_reprocessing(self):
+        """
+        Test that reprocessing the same images with differet pipelines does not create duplicate
+        detections. The 2 pipelines used are a random detection + random species classifier, and a
+        constant species classifier.
+        """
+        # Process the images once
+        pipeline = self.processing_service_instance.pipelines.all().get(slug="random-detection-random-species")
+        pipeline_response = pipeline.process_images(self.test_images, project_id=self.project.pk)
+        results = save_results(pipeline_response, return_created=True)
+        assert results is not None, "Expected results to be returned in a PipelineSaveResults object"
+        assert results.detections, "Expected detections to be returned in the results"
+
+        # This particular pipeline produces 2 classifications per detection
+        for det in results.detections:
+            num_classifications = det.classifications.count()
+            assert (
+                num_classifications == 2
+            ), "Expected 2 classifications per detection (random species and random binary classifier)."
+
+        source_images = SourceImage.objects.filter(pk__in=[image.id for image in pipeline_response.source_images])
+        detections = Detection.objects.filter(source_image__in=source_images).select_related(
+            "detection_algorithm",
+            "detection_algorithm__category_map",
+        )
+        initial_detection_ids = sorted([det.pk for det in detections])
+        assert detections.count() > 0
+
+        # Reprocess the same images using a different pipeline
+        pipeline = self.processing_service_instance.pipelines.all().get(slug="constant")
+        pipeline_response = pipeline.process_images(self.test_images, project_id=self.project.pk)
+        reprocessed_results = save_results(pipeline_response, return_created=True)
+        assert reprocessed_results is not None, "Expected results to be returned in a PipelineSaveResults object"
+        assert reprocessed_results.detections, "Expected detections to be returned in the results"
+
+        source_images = SourceImage.objects.filter(pk__in=[image.id for image in pipeline_response.source_images])
+        detections = Detection.objects.filter(source_image__in=source_images).select_related(
+            "detection_algorithm",
+            "detection_algorithm__category_map",
+        )
+
+        # Check detections were re-processed, and not re-created
+        reprocessed_detection_ids = sorted([det.pk for det in detections])
+        assert initial_detection_ids == reprocessed_detection_ids, (
+            "Expected the same detections to be returned after reprocessing with a different pipeline, "
+            f"but found {initial_detection_ids} != {reprocessed_detection_ids}"
+        )
+
+        # The constant pipeline produces 1 classification per detection
+        for detection in detections:
+            assert (
+                detection.classifications.count() == 3
+            ), "Expected 3 classifications per detection (2 random classifiers + constant classifier)."
+
 
 class TestPipeline(TestCase):
     def setUp(self):

processing_services/README.md

@@ -8,13 +8,13 @@ In this directory, we define locally-run processing services as FastAPI apps. A
 - `/info`: returns data about what pipelines and algorithms are supported by the service.
 - `/livez`
 - `/readyz`
-- `/process`: receives source images via a `PipelineRequest` and returns a `PipelineResponse` containing detections
+- `/process`: receives source images and existing detections via a `PipelineRequest` and returns a `PipelineResponse` containing detections
 
 `processing_services` contains 2 apps:
 - `example`: demos how to add custom pipelines/algorithms.
 - `minimal`: a simple ML backend for basic testing of the processing service API. This minimal app also runs within the main Antenna docker compose stack.
 
-If your goal is to run an ML backend locally, simply copy the `example` directory and follow the steps below.
+If your goal is to run an ML backend locally, simply copy the `example` app and follow the steps below.
 
 ## Environment Set Up
 
@@ -38,26 +38,56 @@ If your goal is to run an ML backend locally, simply copy the `example` director
 3. Implement the `run()` function. Some important considerations:
     - Always run `_get_pipeline_response` at the end of `run()` to get a valid `PipelineResultsResponse`
     - Each algorithm/stage in a pipeline should take a list of `SourceImage`s or `Detection`s and produce a list of `Detection`s (with or without classifications). The class member function `_get_detections()` handles this general stage structure; it batchifys the inputs and produces output detections.
-    - 2 example pipelines are already implemented:
-        - `ZeroShotHFClassifierPipeline`: localizer + classifier
-        - `ZeroShotObjectDetectorPipeline`: detector
+    - 4 example pipelines are already implemented. See the table at the end of the README for examples of what detections from each pipeline look like.
+        - `ZeroShotHFClassifierPipeline`
+        - `ZeroShotObjectDetectorPipeline`
+        - `ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline`
+        - `ZeroShotObjectDetectorWithConstantClassifierPipeline`
 
 4. Add `NewPipeline` to `processing_services/example/api/api.py`
 
 ```
-from .pipelines import Pipeline, ZeroShotHFClassifierPipeline, ZeroShotObjectDetectorPipeline, NewPipeline
-
-...
-
-pipelines: list[type[Pipeline]] = [ZeroShotHFClassifierPipeline, ZeroShotObjectDetectorPipeline, NewPipeline ]
-
+from .pipelines import (
+    Pipeline,
+    ZeroShotHFClassifierPipeline,
+    ZeroShotObjectDetectorPipeline,
+    ZeroShotObjectDetectorWithConstantClassifierPipeline,
+    ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline,
+    NewPipeline
+)
 ...
 
+pipelines: list[type[Pipeline]] = [
+    ZeroShotHFClassifierPipeline,
+    ZeroShotObjectDetectorPipeline,
+    ZeroShotObjectDetectorWithConstantClassifierPipeline,
+    ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline,
+    NewPipeline,
+]
 ```
 5. Update `PipelineChoice` in `processing_services/example/api/schemas.py` to include the slug of the new pipeline, as defined in `NewPipeline`'s config.
 
 ```
 PipelineChoice = typing.Literal[
-    "zero-shot-hf-classifier-pipeline", "zero-shot-object-detector-pipeline", "new-pipeline"
+    "zero-shot-hf-classifier-pipeline",
+    "zero-shot-object-detector-pipeline",
+    "zero-shot-object-detector-with-constant-classifier-pipeline",
+    "zero-shot-object-detector-with-random-species-classifier-pipeline",
+    "new-pipeline-slug",
 ]
 ```
+## Demo
+
+## `minimal` Pipelines and Output Images
+
+- `ConstantPipeline` and `RandomDetectionRandomSpeciesPipeline`
+![MinimalReprocessing](images/MinimalReprocessing.png)
+
+
+## `example` Pipelines and Output Images
+
+- `ZeroShotHFClassifierPipeline`
+![ZeroShotHFClassifierPipeline](images/ZeroShotHFClassifierPipeline.png)
+
+- `ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline` and `ZeroShotObjectDetectorWithConstantClassifierPipeline` (using reprocessing, skips the Zero Shot Object Detector if there are existing detections)
+![ZeroShotReprocessing](images/ZeroShotReprocessing.png)