Droid210/FleetVision

1"""Model B wrapper for multi-angle damage inspection."""2from __future__ import annotations3 4import asyncio5from typing import Dict, List6 7from models.model_b.grad_cam import generate_damage_heatmap8from models.model_b.inference import classify_damage9 10ANGLE_ORDER = ["Front", "Back", "Left", "Right"]11 12 13async def inspect_angle(angle: str, image_path: str, model_path: str) -> dict:14    """Classify a single image and generate Grad-CAM if damaged."""15 16    def _run_sync() -> dict:17        status, confidence = classify_damage(image_path, model_path=model_path)18        heatmap_path = None19        if status == "Damaged":20            heatmap_path, _ = generate_damage_heatmap(21                image_path=image_path,22                model_path=model_path,23                device_type="cpu",24            )25        return {26            "angle": angle,27            "status": status,28            "confidence": round(confidence, 4),29            "heatmap_path": heatmap_path,30        }31 32    return await asyncio.to_thread(_run_sync)33 34 35async def inspect_all_angles(image_paths: Dict[str, str], model_path: str = "weights/model b/best_damage_detector.pth") -> List[dict]:36    """Run Model B in parallel over 5 required angles."""37    missing = [name for name in ANGLE_ORDER if name not in image_paths]38    if missing:39        raise ValueError(f"Missing angles: {missing}")40 41    tasks = [inspect_angle(angle, image_paths[angle], model_path) for angle in ANGLE_ORDER]42    results = await asyncio.gather(*tasks)43    return results44