Spaces:
Sleeping
Sleeping
File size: 4,057 Bytes
bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 2c527f4 bcbe62f 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 | """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/.
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:
from transformers import pipeline
# 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 metadata_path.exists():
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
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",
],
}
|