Upload inference.py

inference.py

@@ -195,14 +195,15 @@ def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
         # Run YOLOv8 classification (verbose=False suppresses progress bar)
         results = animal_model(crop, verbose=False)
 
-        # Extract class names dict: {0: "porcupine", 1: "elephant", ...}
+        # Extract class names dict (YOLOv8 uses alphabetical order)
+        # Example: {0: "aardwolf", 1: "african wild cat", ..., 13: "giraffe", ...}
         names_dict = results[0].names
 
         # Extract probabilities: [0.0001, 0.0002, ..., 0.9998, ...]
         probs = results[0].probs.data.tolist()
 
         # Build list of (class_name, confidence) tuples
-        # YOLOv8 names_dict is already in the correct format: {0: "porcupine", 1: "elephant", ...}
+        # Return YOLOv8's class names (which will be mapped to taxonomy IDs later)
         classifications = []
         for idx, class_name in names_dict.items():
             confidence = probs[idx]
@@ -219,32 +220,31 @@ def get_classification(crop: Image.Image) -> list[tuple[str, float]]:
 
 def get_class_names() -> dict[str, str]:
     """
-    Get mapping of class IDs to species names from taxonomy.csv.
+    Get mapping of class IDs to species names from YOLOv8 model.
 
-    Reads taxonomy.csv from the model directory and extracts the model_class
-    (common name) for each species.
+    YOLOv8 stores class names in alphabetical order internally. This function
+    extracts those names and creates a 1-indexed mapping for the JSON format.
+
+    NOTE: taxonomy.csv is NOT used here - it's only for UI taxonomy tree display.
+    The class IDs here are YOLOv8's alphabetical indices (0-based) + 1.
 
     Returns:
         Dict mapping class ID (1-indexed string) to common name
-        Example: {"1": "porcupine", "2": "elephant", ...}
+        Example: {"1": "aardwolf", "2": "african wild cat", ..., "14": "giraffe", ...}
 
     Raises:
-        FileNotFoundError: If taxonomy.csv not found
-        RuntimeError: If parsing fails
+        RuntimeError: If model not loaded
     """
-    global animal_model, MODEL_DIR
+    global animal_model
 
-    # 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:
         raise RuntimeError("Model not loaded - call load_model() first")
 
     try:
-        # YOLOv8 names dict: {0: "porcupine", 1: "elephant", ...}
+        # YOLOv8 names dict (alphabetical order): {0: "aardwolf", 1: "african wild cat", ...}
         yolo_names = animal_model.names
 
-        # Convert to 1-indexed dict for compatibility
+        # Convert to 1-indexed dict for JSON compatibility
         class_names = {}
         for idx, name in yolo_names.items():
             class_id_str = str(idx + 1)  # 1-indexed
@@ -253,4 +253,4 @@ def get_class_names() -> dict[str, str]:
         return class_names
 
     except Exception as e:
-        raise RuntimeError(f"Failed to extract class names: {e}") from e
+        raise RuntimeError(f"Failed to extract class names from model: {e}") from e