fix: return category maps from all models in example processing service

Author: mihow

processing_services/example/api/algorithms.py

@@ -42,15 +42,28 @@ def compile(self):
     def run(self, inputs: list[SourceImage] | list[Detection]) -> list[Detection]:
         raise NotImplementedError("Subclasses must implement the run method")
 
-    algorithm_config_response = AlgorithmConfigResponse(
-        name="Base Algorithm",
-        key="base",
-        task_type="base",
-        description="A base class for all algorithms.",
-        version=1,
-        version_name="v1",
-        category_map=None,
-    )
+    def get_category_map(self) -> AlgorithmCategoryMapResponse:
+        return AlgorithmCategoryMapResponse(
+            data=[],
+            labels=[],
+            version="v1",
+            description="A model without labels.",
+            uri=None,
+        )
+
+    def get_algorithm_config_response(self) -> AlgorithmConfigResponse:
+        return AlgorithmConfigResponse(
+            name="Base Algorithm",
+            key="base",
+            task_type="base",
+            description="A base class for all algorithms.",
+            version=1,
+            version_name="v1",
+            category_map=self.get_category_map(),
+        )
+
+    def __init__(self):
+        self.algorithm_config_response = self.get_algorithm_config_response()
 
 
 class ZeroShotObjectDetector(Algorithm):
@@ -141,33 +154,45 @@ def run(self, source_images: list[SourceImage], intermediate=False) -> list[Dete
 
         return detector_responses
 
-    algorithm_config_response = AlgorithmConfigResponse(
-        name="Zero Shot Object Detector",
-        key="zero-shot-object-detector",
-        task_type="detection",
-        description=(
-            "Huggingface Zero Shot Object Detection model."
-            "Produces both a bounding box and a candidate label classification for each detection."
-        ),
-        version=1,
-        version_name="v1",
-        category_map=None,
-    )
+    def get_category_map(self) -> AlgorithmCategoryMapResponse:
+        return AlgorithmCategoryMapResponse(
+            data=[{"index": i, "label": label} for i, label in enumerate(self.candidate_labels)],
+            labels=self.candidate_labels,
+            version="v1",
+            description="Candidate labels used for zero-shot object detection.",
+            uri=None,
+        )
+
+    def get_algorithm_config_response(self) -> AlgorithmConfigResponse:
+        return AlgorithmConfigResponse(
+            name="Zero Shot Object Detector",
+            key="zero-shot-object-detector",
+            task_type="detection",
+            description=(
+                "Huggingface Zero Shot Object Detection model."
+                "Produces both a bounding box and a candidate label classification for each detection."
+            ),
+            version=1,
+            version_name="v1",
+            category_map=self.get_category_map(),
+        )
 
 
 class HFImageClassifier(Algorithm):
     """
     A  local classifier that uses the Hugging Face pipeline to classify images.
     """
 
+    model_name: str = "google/vit-base-patch16-224"  # Vision Transformer model trained on ImageNet-1k
+
     def compile(self):
         saved_models_key = "hf_image_classifier"  # generate a key for each uniquely compiled algorithm
 
         if saved_models_key not in SAVED_MODELS:
             from transformers import pipeline
 
             logger.info(f"Compiling {self.algorithm_config_response.name} from scratch...")
-            self.model = pipeline("image-classification", model="google/vit-base-patch16-224")
+            self.model = pipeline("image-classification", model=self.model_name, device=get_best_device())
             SAVED_MODELS[saved_models_key] = self.model
         else:
             logger.info(f"Using saved model for {self.algorithm_config_response.name}...")
@@ -216,15 +241,55 @@ def run(self, detections: list[Detection]) -> list[Detection]:
 
         return detections_to_return
 
-    algorithm_config_response = AlgorithmConfigResponse(
-        name="HF Image Classifier",
-        key="hf-image-classifier",
-        task_type="classification",
-        description="HF ViT for image classification.",
-        version=1,
-        version_name="v1",
-        category_map=None,
-    )
+    def get_category_map(self) -> AlgorithmCategoryMapResponse:
+        """
+        Extract the category map from the model.
+        Returns an AlgorithmCategoryMapResponse with labels, data, and model information.
+        """
+        from transformers.models.auto.configuration_auto import AutoConfig
+
+        logger.info(f"Loading configuration for {self.model_name}")
+        config = AutoConfig.from_pretrained(self.model_name)
+
+        # Extract label information
+        if not hasattr(config, "id2label") or not config.id2label:
+            raise ValueError(
+                f"Cannot create category map for model {self.model_name}, no id2label mapping found in config"
+            )
+        else:
+            # Sort labels by index
+            # Ensure keys are strings for consistent access
+            id2label: dict[str, str] = {str(k): v for k, v in config.id2label.items()}
+            indices = sorted([int(k) for k in id2label.keys()])
+
+            # Create labels and data
+            labels = [id2label[str(i)] for i in indices]
+            data = [{"label": label, "index": idx} for idx, label in zip(indices, labels)]
+
+        # Build description
+        description_text = (
+            f"Vision Transformer model trained on ImageNet-1k. "
+            f"Contains {len(labels)} object classes. Model: {self.model_name}"
+        )
+
+        return AlgorithmCategoryMapResponse(
+            data=data,
+            labels=labels,
+            version="ImageNet-1k",
+            description=description_text,
+            uri=f"https://huggingface.co/{self.model_name}",
+        )
+
+    def get_algorithm_config_response(self) -> AlgorithmConfigResponse:
+        return AlgorithmConfigResponse(
+            name="HF Image Classifier",
+            key="hf-image-classifier",
+            task_type="classification",
+            description="HF ViT for image classification.",
+            version=1,
+            version_name="v1",
+            category_map=self.get_category_map(),
+        )
 
 
 class RandomSpeciesClassifier(Algorithm):

processing_services/example/api/pipelines.py

@@ -153,8 +153,8 @@ class ZeroShotHFClassifierPipeline(Pipeline):
         description=("Zero Shot Object Detector with HF image classifier."),
         version=1,
         algorithms=[
-            ZeroShotObjectDetector.algorithm_config_response,
-            HFImageClassifier.algorithm_config_response,
+            ZeroShotObjectDetector().algorithm_config_response,
+            HFImageClassifier().algorithm_config_response,
         ],
     )
 
@@ -167,7 +167,7 @@ def get_stages(self) -> list[Algorithm]:
             zero_shot_object_detector.candidate_labels = self.request_config["candidate_labels"]
         self.config.algorithms = [
             zero_shot_object_detector.algorithm_config_response,
-            HFImageClassifier.algorithm_config_response,
+            HFImageClassifier().algorithm_config_response,
         ]
 
         return [zero_shot_object_detector, HFImageClassifier()]
@@ -212,7 +212,7 @@ class ZeroShotObjectDetectorPipeline(Pipeline):
         slug="zero-shot-object-detector-pipeline",
         description=("Zero shot object detector (bbox and classification)."),
         version=1,
-        algorithms=[ZeroShotObjectDetector.algorithm_config_response],
+        algorithms=[ZeroShotObjectDetector().algorithm_config_response],
     )
 
     def get_stages(self) -> list[Algorithm]:
@@ -254,8 +254,8 @@ class ZeroShotObjectDetectorWithRandomSpeciesClassifierPipeline(Pipeline):
         description=("HF zero shot object detector with random species classifier."),
         version=1,
         algorithms=[
-            ZeroShotObjectDetector.algorithm_config_response,
-            RandomSpeciesClassifier.algorithm_config_response,
+            ZeroShotObjectDetector().algorithm_config_response,
+            RandomSpeciesClassifier().algorithm_config_response,
         ],
     )
 
@@ -266,7 +266,7 @@ def get_stages(self) -> list[Algorithm]:
 
         self.config.algorithms = [
             zero_shot_object_detector.algorithm_config_response,
-            RandomSpeciesClassifier.algorithm_config_response,
+            RandomSpeciesClassifier().algorithm_config_response,
         ]
 
         return [zero_shot_object_detector, RandomSpeciesClassifier()]
@@ -307,8 +307,8 @@ class ZeroShotObjectDetectorWithConstantClassifierPipeline(Pipeline):
         description=("HF zero shot object detector with constant classifier."),
         version=1,
         algorithms=[
-            ZeroShotObjectDetector.algorithm_config_response,
-            ConstantClassifier.algorithm_config_response,
+            ZeroShotObjectDetector().algorithm_config_response,
+            ConstantClassifier().algorithm_config_response,
         ],
     )
 
@@ -319,7 +319,7 @@ def get_stages(self) -> list[Algorithm]:
 
         self.config.algorithms = [
             zero_shot_object_detector.algorithm_config_response,
-            ConstantClassifier.algorithm_config_response,
+            ConstantClassifier().algorithm_config_response,
         ]
 
         return [zero_shot_object_detector, ConstantClassifier()]