Update src/vision_analyzer.py

src/vision_analyzer.py

@@ -1,8 +1,7 @@
-"""Computer Vision analyzer for coarse vehicle class prediction.
+"""Computer Vision analyzer using transfer-learning ResNet-18 for vehicle brand classification.
 
-The module supports two levels:
-1) A local classifier trained on hand-crafted image features.
-2) A deterministic image-signature fallback so the app always runs.
+The module loads a Hugging Face fine-tuned image classification model from
+models/car-image-classifier/ and provides vehicle brand predictions.
 
 No damage detection and no technical condition estimation is implemented.
 """
@@ -11,63 +10,58 @@ from __future__ import annotations
 
 import json
 from functools import lru_cache
+from pathlib import Path
 from typing import Any
 
-import joblib
 import numpy as np
 from PIL import Image
 
-from src.config import VISION_LABELS_PATH, VISION_METADATA_PATH, VISION_MODEL_PATH
-from src.vision_features import extract_image_features, ensure_pil_image
-
-DEFAULT_CLASSES = [
-    "Audi",
-    "Hyundai Creta",
-    "Mahindra Scorpio",
-    "Rolls Royce",
-    "Swift",
-    "Tata Safari",
-    "Toyota Innova",
-]
-
+from src.config import MODEL_DIR
 
 
+def ensure_pil_image(image: Image.Image | np.ndarray | None) -> Image.Image:
+    """Convert input to PIL Image."""
+    if image is None:
+        raise ValueError("No image provided.")
+    if isinstance(image, np.ndarray):
+        return Image.fromarray(image.astype("uint8"))
+    if not isinstance(image, Image.Image):
+        raise TypeError("Input is not a valid image.")
+    return image
 
 
 @lru_cache(maxsize=1)
-def _load_optional_pytorch_model() -> dict[str, Any] | None:
-    if not VISION_MODEL_PATH.exists():
+def _load_transfer_model() -> dict[str, Any] | None:
+    """Load the transfer-learning model from models/car-image-classifier/.
+    
+    Returns a dict with 'pipeline' and 'metadata' on success, None if model not found.
+    """
+    model_dir = MODEL_DIR / "car-image-classifier"
+    if not model_dir.exists():
         return None
 
     try:
-        model = joblib.load(VISION_MODEL_PATH)
+        from transformers import pipeline
 
-        classes = DEFAULT_CLASSES
-        if VISION_LABELS_PATH.exists():
-            labels_payload = json.loads(VISION_LABELS_PATH.read_text(encoding="utf-8"))
-            if isinstance(labels_payload, list) and labels_payload:
-                classes = [str(x) for x in labels_payload]
+        # Load the image classification pipeline
+        clf_pipeline = pipeline("image-classification", model=str(model_dir), device=-1)
 
+        # Load metadata if it exists
+        metadata_path = model_dir / "vision_metadata.json"
         metadata = {}
-        if VISION_METADATA_PATH.exists():
-            metadata = json.loads(VISION_METADATA_PATH.read_text(encoding="utf-8"))
+        if metadata_path.exists():
+            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
 
-        return {"model": model, "classes": classes, "metadata": metadata}
+        return {"pipeline": clf_pipeline, "metadata": metadata}
     except Exception:
         return None
 
 
-def _fallback_predict(image: Image.Image) -> tuple[str, float, str, list[str]]:
-    notes = [
-        "Fallback classifier active: no trained local CV model found.",
-        "The app falls back to a safe Unknown prediction.",
-        "No damage detection and no technical condition assessment.",
-    ]
-    return "Unknown", 0.0, "fallback", notes
-
-
 def analyze_car_image(image: Image.Image | np.ndarray | None) -> dict[str, Any]:
-    """Analyze uploaded image and return coarse vehicle class/model-group info."""
+    """Analyze uploaded image and return vehicle brand prediction.
+    
+    Uses a transfer-learning ResNet-18 model to classify vehicle brands.
+    """
     if image is None:
         return {
             "predicted_class": "Unknown",
@@ -87,40 +81,42 @@ def analyze_car_image(image: Image.Image | np.ndarray | None) -> dict[str, Any]:
             "notes": ["Input is not a valid image format."],
         }
 
-    model_bundle = _load_optional_pytorch_model()
+    model_bundle = _load_transfer_model()
     if model_bundle is not None:
         try:
-            model = model_bundle["model"]
-            classes = model_bundle["classes"]
+            pipeline = model_bundle["pipeline"]
             metadata = model_bundle.get("metadata", {})
 
-            features = extract_image_features(ensure_pil_image(image))
-            if hasattr(model, "predict_proba"):
-                probs = model.predict_proba([features])[0]
-                idx = int(np.argmax(probs))
-                conf = float(probs[idx])
-            else:
-                idx = int(model.predict([features])[0])
-                conf = 0.5
-
-            predicted = classes[idx] if idx < len(classes) else "Unknown"
-            return {
-                "predicted_class": predicted,
-                "confidence": round(conf, 3),
-                "method": "local_image_classifier",
-                "notes": [
-                    "Coarse image-based class prediction.",
-                    f"Vision test accuracy: {metadata.get('test_accuracy', 'n/a')}.",
-                    "No damage detection or detailed technical diagnostics.",
-                ],
-            }
+            # Run inference
+            pil_image = ensure_pil_image(image)
+            results = pipeline(pil_image, top_k=1)
+
+            if results:
+                top_result = results[0]
+                predicted = top_result["label"]
+                confidence = float(top_result["score"])
+
+                return {
+                    "predicted_class": predicted,
+                    "confidence": round(confidence, 3),
+                    "method": "local_transfer_model",
+                    "notes": [
+                        "Vehicle brand classification using transfer learning (ResNet-18).",
+                        "The classifier can only predict one of the trained classes.",
+                        f"Model accuracy on test set: {metadata.get('accuracy', 'n/a')}.",
+                        "No damage detection or technical condition assessment.",
+                    ],
+                }
         except Exception:
             pass
 
-    predicted, confidence, method, notes = _fallback_predict(ensure_pil_image(image))
+    # Fallback if model not found or inference fails
     return {
-        "predicted_class": predicted,
-        "confidence": confidence,
-        "method": method,
-        "notes": notes,
+        "predicted_class": "Unknown",
+        "confidence": 0.0,
+        "method": "fallback",
+        "notes": [
+            "No trained transfer-learning model found.",
+            "Please train the model using: python -m src.train_vision_model",
+        ],
     }