Extract ensemble prediction logic to a new class to improve readability

Author: tomergadot

Diff for: speciesnet/ensemble.py

@@ -20,16 +20,16 @@
 
 import json
 import time
-from typing import Any, Optional
+from typing import Any, Callable
 
 from absl import logging
 from humanfriendly import format_timespan
 import PIL.ExifTags
 import PIL.Image
 
 from speciesnet.constants import Classification
-from speciesnet.constants import Detection
 from speciesnet.constants import Failure
+from speciesnet.ensemble_prediction_combiner import combine_predictions_for_single_item
 from speciesnet.geofence_utils import geofence_animal_classification
 from speciesnet.geofence_utils import roll_up_labels_to_first_matching_level
 from speciesnet.utils import ModelInfo
@@ -44,7 +44,12 @@
 class SpeciesNetEnsemble:
     """Ensemble component of SpeciesNet."""
 
-    def __init__(self, model_name: str, geofence: bool = True) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        geofence: bool = True,
+        prediction_combiner: Callable = combine_predictions_for_single_item,
+    ) -> None:
         """Loads the ensemble resources.
 
         Args:
@@ -62,6 +67,7 @@ def __init__(self, model_name: str, geofence: bool = True) -> None:
         self.enable_geofence = geofence
         self.taxonomy_map = self.load_taxonomy()
         self.geofence_map = self.load_geofence()
+        self.prediction_combiner = prediction_combiner
 
         end_time = time.time()
         logging.info(
@@ -101,153 +107,6 @@ def load_geofence(self):
             geofence_map = json.load(fp)
         return geofence_map
 
-    def _combine_predictions_for_single_item(
-        self,
-        classifications: dict[str, list],
-        detections: list[dict],
-        country: Optional[str],
-        admin1_region: Optional[str],
-    ) -> PredictionType:
-        """Ensembles classifications and detections for a single image.
-
-        This operation leverages multiple heuristics to make the most of the classifier
-        and the detector predictions through a complex set of decisions. It introduces
-        various thresholds to identify humans, vehicles, blanks, animals at species
-        level, animals at higher taxonomy levels and even unknowns.
-
-        Args:
-            classifications:
-                Dict of classification results. "classes" and "scores" are expected to
-                be provided among the dict keys.
-            detections:
-                List of detection results, sorted in decreasing order of their
-                confidence score. Each detection is expected to be a dict providing
-                "label" and "conf" among its keys.
-            country:
-                Country (in ISO 3166-1 alpha-3 format) associated with predictions.
-                Optional.
-            admin1_region:
-                First-level administrative division (in ISO 3166-2 format) associated
-                with predictions. Optional.
-
-        Returns:
-            A tuple of <label, score, prediction_source> describing the ensemble result.
-        """
-
-        top_classification_class = classifications["classes"][0]
-        top_classification_score = classifications["scores"][0]
-        top_detection_class = detections[0]["label"] if detections else Detection.ANIMAL
-        top_detection_score = detections[0]["conf"] if detections else 0.0
-
-        if top_detection_class == Detection.HUMAN:
-            # Threshold #1a: high-confidence HUMAN detections.
-            if top_detection_score > 0.7:
-                return Classification.HUMAN, top_detection_score, "detector"
-
-            # Threshold #1b: mid-confidence HUMAN detections + high-confidence
-            # HUMAN/VEHICLE classifications.
-            if (
-                top_detection_score > 0.2
-                and top_classification_class
-                in {Classification.HUMAN, Classification.VEHICLE}
-                and top_classification_score > 0.5
-            ):
-                return Classification.HUMAN, top_classification_score, "classifier"
-
-        if top_detection_class == Detection.VEHICLE:
-            # Threshold #2a: mid-confidence VEHICLE detections + high-confidence HUMAN
-            # classifications.
-            if (
-                top_detection_score > 0.2
-                and top_classification_class == Classification.HUMAN
-                and top_classification_score > 0.5
-            ):
-                return Classification.HUMAN, top_classification_score, "classifier"
-
-            # Threshold #2b: high-confidence VEHICLE detections.
-            if top_detection_score > 0.7:
-                return Classification.VEHICLE, top_detection_score, "detector"
-
-            # Threshold #2c: mid-confidence VEHICLE detections + high-confidence VEHICLE
-            # classifications.
-            if (
-                top_detection_score > 0.2
-                and top_classification_class == Classification.VEHICLE
-                and top_classification_score > 0.4
-            ):
-                return Classification.VEHICLE, top_classification_score, "classifier"
-
-        # Threshold #3a: high-confidence BLANK "detections" + high-confidence BLANK
-        # classifications.
-        if (
-            top_detection_score < 0.2
-            and top_classification_class == Classification.BLANK
-            and top_classification_score > 0.5
-        ):
-            return Classification.BLANK, top_classification_score, "classifier"
-
-        # Threshold #3b: extra-high-confidence BLANK classifications.
-        if (
-            top_classification_class == Classification.BLANK
-            and top_classification_score > 0.99
-        ):
-            return Classification.BLANK, top_classification_score, "classifier"
-
-        if top_classification_class not in {
-            Classification.BLANK,
-            Classification.HUMAN,
-            Classification.VEHICLE,
-        }:
-            # Threshold #4a: extra-high-confidence ANIMAL classifications.
-            if top_classification_score > 0.8:
-                return geofence_animal_classification(
-                    labels=classifications["classes"],
-                    scores=classifications["scores"],
-                    country=country,
-                    admin1_region=admin1_region,
-                    taxonomy_map=self.taxonomy_map,
-                    geofence_map=self.geofence_map,
-                    enable_geofence=self.enable_geofence,
-                )
-
-            # Threshold #4b: high-confidence ANIMAL classifications + mid-confidence
-            # ANIMAL detections.
-            if (
-                top_classification_score > 0.65
-                and top_detection_class == Detection.ANIMAL
-                and top_detection_score > 0.2
-            ):
-                return geofence_animal_classification(
-                    labels=classifications["classes"],
-                    scores=classifications["scores"],
-                    country=country,
-                    admin1_region=admin1_region,
-                    taxonomy_map=self.taxonomy_map,
-                    geofence_map=self.geofence_map,
-                    enable_geofence=self.enable_geofence,
-                )
-
-        # Threshold #5a: high-confidence ANIMAL rollups.
-        rollup = roll_up_labels_to_first_matching_level(
-            labels=classifications["classes"],
-            scores=classifications["scores"],
-            country=country,
-            admin1_region=admin1_region,
-            target_taxonomy_levels=["genus", "family", "order", "class", "kingdom"],
-            non_blank_threshold=0.65,
-            taxonomy_map=self.taxonomy_map,
-            geofence_map=self.geofence_map,
-            enable_geofence=self.enable_geofence,
-        )
-        if rollup:
-            return rollup
-
-        # Threshold #5b: mid-confidence ANIMAL detections.
-        if top_detection_class == Detection.ANIMAL and top_detection_score > 0.5:
-            return Classification.ANIMAL, top_detection_score, "detector"
-
-        return Classification.UNKNOWN, top_classification_score, "classifier"
-
     def combine(  # pylint: disable=too-many-positional-arguments
         self,
         filepaths: list[str],
@@ -331,11 +190,16 @@ def combine(  # pylint: disable=too-many-positional-arguments
 
             # Most importantly, ensemble everything into a single prediction.
             if classifications is not None and detections is not None:
-                prediction, score, source = self._combine_predictions_for_single_item(
+                prediction, score, source = self.prediction_combiner(
                     classifications=classifications,
                     detections=detections,
                     country=geolocation.get("country"),
                     admin1_region=geolocation.get("admin1_region"),
+                    taxonomy_map=self.taxonomy_map,
+                    geofence_map=self.geofence_map,
+                    enable_geofence=self.enable_geofence,
+                    geofence_fn=geofence_animal_classification,
+                    roll_up_fn=roll_up_labels_to_first_matching_level,
                 )
                 result["prediction"] = (
                     prediction.value