Upload inference.py

inference.py

@@ -74,9 +74,8 @@ def load_model() -> None:
     """
     global animal_model, MODEL_PATH
 
-    if MODEL_PATH is None:
-        raise RuntimeError("MODEL_PATH not set - must be injected by framework")
-
+    # MODEL_PATH is injected by framework before this function is called
+    # Check that the path exists (framework guarantees it's not None)
     if not MODEL_PATH.exists():
         raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")
 
@@ -181,8 +180,8 @@ def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
         crop: Cropped and preprocessed PIL Image
 
     Returns:
-        List of (class_id, confidence) tuples for ALL classes, sorted by confidence.
-        Example: [("14", 0.99985), ("7", 0.00003), ...]
+        List of (class_name, confidence) tuples for ALL classes, sorted by confidence.
+        Example: [("giraffe", 0.99985), ("cattle", 0.00003), ...]
 
     Raises:
         RuntimeError: If model not loaded or inference fails
@@ -202,15 +201,12 @@ def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
         # Extract probabilities: [0.0001, 0.0002, ..., 0.9998, ...]
         probs = results[0].probs.data.tolist()
 
-        # Build list of (class_id, confidence) tuples
-        # Class IDs are 0-indexed in YOLOv8 but we output as strings
+        # Build list of (class_name, confidence) tuples
+        # YOLOv8 names_dict is already in the correct format: {0: "porcupine", 1: "elephant", ...}
         classifications = []
         for idx, class_name in names_dict.items():
-            # YOLOv8 uses 0-based indexing, but we need 1-based for compatibility
-            # with the taxonomy.csv and expected JSON output
-            class_id_str = str(idx + 1)  # Convert 0-indexed to 1-indexed
             confidence = probs[idx]
-            classifications.append((class_id_str, confidence))
+            classifications.append((class_name, confidence))
 
         # Sort by confidence descending (already sorted by YOLOv8, but ensure it)
         classifications.sort(key=lambda x: x[1], reverse=True)
@@ -238,9 +234,7 @@ def get_class_names() -> dict[str, str]:
     """
     global animal_model, MODEL_DIR
 
-    if MODEL_DIR is None:
-        raise RuntimeError("MODEL_DIR not set - must be injected by framework")
-
+    # MODEL_DIR is injected by framework (guaranteed not None)
     # YOLOv8 models have class names built-in
     # We'll use those directly since they should match taxonomy.csv
     if animal_model is None: