Droid210/FleetVision

1"""Single image inference for damage detection."""2from pathlib import Path3from typing import Tuple4 5import torch6from PIL import Image7from transformers import AutoImageProcessor8 9from .config import MODEL_ID10 11 12def load_trained_model(13    model_path: Path,14    device: torch.device,15):16    """Load trained image classification model from checkpoint.17 18    Args:19        model_path: Path to model checkpoint.20        device: Torch device.21 22    Returns:23        Tuple of (model, processor).24    """25    from .model import build_model26 27    checkpoint = torch.load(model_path, map_location=device, weights_only=False)28    model = build_model()29    model.load_state_dict(checkpoint["model_state_dict"])30    model.to(device)31    model.eval()32 33    processor = AutoImageProcessor.from_pretrained(MODEL_ID)34    return model, processor35 36 37def classify_damage(38    image_path: str,39    model_path: str = "weights/model b/best_damage_detector.pth",40    device: str | None = None,41) -> Tuple[str, float]:42    """Classify if car is damaged.43 44    Args:45        image_path: Path to car image.46        model_path: Path to trained model.47        device: Device to use ('cuda', 'cpu', or None for auto with fallback).48 49    Returns:50        Tuple of (class_name, confidence).51    """52    # Device selection with fallback53    if device is None:54        device_obj = torch.device("cuda" if torch.cuda.is_available() else "cpu")55    else:56        device_obj = torch.device(device)57 58    try:59        model, processor = load_trained_model(Path(model_path), device_obj)60    except RuntimeError as e:61        if "CUDA" in str(e) and device_obj.type == "cuda":62            print("WARNING: CUDA error during model loading. Falling back to CPU...")63            device_obj = torch.device("cpu")64            model, processor = load_trained_model(Path(model_path), device_obj)65        else:66            raise67 68    # Load and process image69    image = Image.open(image_path).convert("RGB")70    processed = processor(image, return_tensors="pt")71    pixel_values = processed["pixel_values"].to(device_obj)72 73    try:74        with torch.no_grad():75            outputs = model(pixel_values)76            logits = outputs.logits77            probs = torch.softmax(logits, dim=1)78            confidence, pred_idx = torch.max(probs, dim=1)79    except RuntimeError as e:80        if "CUDA" in str(e) and device_obj.type == "cuda":81            print("WARNING: CUDA error during inference. Retrying on CPU...")82            device_obj = torch.device("cpu")83            model.to(device_obj)84            pixel_values = pixel_values.to(device_obj)85            with torch.no_grad():86                outputs = model(pixel_values)87                logits = outputs.logits88                probs = torch.softmax(logits, dim=1)89                confidence, pred_idx = torch.max(probs, dim=1)90        else:91            raise92 93    class_name = "Damaged" if pred_idx.item() == 1 else "Whole"94    confidence_score = confidence.item()95 96    return class_name, confidence_score97