Spaces:
Sleeping
Sleeping
File size: 4,150 Bytes
c36af5b 2c527f4 c36af5b b99d10a c36af5b 2c527f4 c36af5b 2c527f4 c36af5b 2c527f4 b99d10a c36af5b b99d10a c36af5b b99d10a c36af5b b99d10a c36af5b 2c527f4 c36af5b b99d10a c36af5b 2c527f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Computer Vision analyzer using transfer-learning ResNet-18 for vehicle brand classification.
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.
"""
from __future__ import annotations
import json
from functools import lru_cache
from pathlib import Path
from typing import Any
import numpy as np
from PIL import Image
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_transfer_model() -> dict[str, Any] | None:
"""Load the transfer-learning model from models/car-image-classifier/ or from HF Hub.
Returns a dict with 'pipeline' and 'metadata' on success, None if model not found.
"""
from transformers import pipeline
model_dir = MODEL_DIR / "car-image-classifier"
model_source = None
metadata = {}
# Try to load locally first
if model_dir.exists():
model_source = str(model_dir)
metadata_path = model_dir / "vision_metadata.json"
if metadata_path.exists():
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
else:
# Fall back to Hugging Face Hub
model_source = "ochsncon/car-image-classifier"
try:
clf_pipeline = pipeline("image-classification", model=model_source, device=-1)
return {"pipeline": clf_pipeline, "metadata": metadata}
except Exception:
return None
def analyze_car_image(image: Image.Image | np.ndarray | None) -> dict[str, Any]:
"""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",
"confidence": 0.0,
"method": "no_image",
"notes": ["No image was provided."],
}
if isinstance(image, np.ndarray):
image = Image.fromarray(image.astype("uint8"))
if not isinstance(image, Image.Image):
return {
"predicted_class": "Unknown",
"confidence": 0.0,
"method": "invalid_input",
"notes": ["Input is not a valid image format."],
}
model_bundle = _load_transfer_model()
if model_bundle is not None:
try:
pipeline = model_bundle["pipeline"]
metadata = model_bundle.get("metadata", {})
# 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
# Fallback if model not found or inference fails
return {
"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",
],
}
|