src/detector.py

"""
CarDentIQ — inference engine
Developer: Saksham Pathak (github.com/parthmax2)

Exposes:
  CLASS_NAMES         dict[int, str]
  CLASS_COLORS        dict[int, tuple]
  DEFAULT_THRESHOLDS  dict[int, float]
  load_model()        → YOLO
  detect()            → dict
"""

import cv2
import numpy as np
from ultralytics import YOLO

# ── class registry ───────────────────────────────────────────────
CLASS_NAMES: dict[int, str] = {
    0: "no damage",
    1: "lost parts",
    2: "torn",
    3: "dent",
    4: "paint scratch",
    5: "hole",
    6: "broken glass",
    7: "broken lamp",
}

CLASS_COLORS: dict[int, tuple] = {
    0: (128, 128, 128),   # grey
    1: (0,   128, 255),   # orange
    2: (255,   0, 255),   # magenta
    3: (0,   255,   0),   # green
    4: (0,   140, 255),   # deep orange
    5: (0,   255, 255),   # yellow
    6: (0,     0, 255),   # red
    7: (255, 165,   0),   # blue-ish
}

DEFAULT_THRESHOLDS: dict[int, float] = {
    0: 0.90, 1: 0.26, 2: 0.05, 3: 0.05,
    4: 0.05, 5: 0.16, 6: 0.33, 7: 0.05,
}

# relative repair-impact weight per class, used only for the severity summary
CLASS_SEVERITY_WEIGHT: dict[int, float] = {
    0: 0,   # no_damage
    4: 1,   # paint_scratch — cosmetic
    3: 2,   # dent
    7: 3,   # broken_lamp
    6: 3,   # broken_glass
    5: 4,   # hole
    2: 4,   # torn
    1: 5,   # lost_parts — missing component, most severe
}

_model: YOLO | None = None


# ── model loader (singleton) ─────────────────────────────────────
def load_model(weights: str = "best.pt") -> YOLO:
    global _model
    if _model is None:
        print(f"[detector] Loading model from {weights} …")
        _model = YOLO(weights)
    return _model


# ── helpers ──────────────────────────────────────────────────────
def _xyxy_to_yolo(x1: float, y1: float, x2: float, y2: float,
                   img_w: int, img_h: int) -> tuple:
    cx = (x1 + x2) / 2.0 / img_w
    cy = (y1 + y2) / 2.0 / img_h
    w  = (x2 - x1) / img_w
    h  = (y2 - y1) / img_h
    return cx, cy, w, h


def _draw_box(img: np.ndarray, x1: int, y1: int, x2: int, y2: int,
              label: str, color: tuple) -> None:
    cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
    cv2.putText(img, label, (x1, max(y1 - 10, 10)),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2, cv2.LINE_AA)


def _compute_summary(detections: list[dict]) -> dict:
    by_class: dict[int, dict] = {}
    raw_score = 0.0
    total_defects = 0

    for d in detections:
        cid = d["class_id"]
        if cid == 0:
            continue  # no_damage is not a defect
        total_defects += 1
        raw_score += CLASS_SEVERITY_WEIGHT.get(cid, 2) * d["confidence"]

        entry = by_class.setdefault(cid, {
            "class_id": cid, "class_name": d["class_name"],
            "color": CLASS_COLORS.get(cid, (255, 255, 255)),
            "count": 0, "confidence_sum": 0.0,
        })
        entry["count"] += 1
        entry["confidence_sum"] += d["confidence"]

    by_class_list = [
        {
            "class_id": e["class_id"], "class_name": e["class_name"], "color": e["color"],
            "count": e["count"], "avg_confidence": round(e["confidence_sum"] / e["count"], 3),
        }
        for e in by_class.values()
    ]
    by_class_list.sort(key=lambda x: x["count"], reverse=True)

    score = round(min(raw_score, 10), 1)
    if total_defects == 0:
        label = "No Damage Detected"
    elif score <= 3:
        label = "Minor"
    elif score <= 7:
        label = "Moderate"
    else:
        label = "Severe"

    return {
        "severity_label": label,
        "severity_score": score,
        "total_defects": total_defects,
        "by_class": by_class_list,
    }


# ── main inference function ───────────────────────────────────────
def detect(
    image_rgb: np.ndarray,
    resize: bool,
    thresholds: dict[int, float],
    weights: str = "best.pt",
) -> dict:
    """
    Parameters
    ----------
    image_rgb   : np.ndarray (H, W, 3) — RGB image
    resize      : bool — cap image at 1024 px before inference
    thresholds  : per-class confidence floor {class_id: min_conf}
    weights     : path to model weights file

    Returns
    -------
    dict with keys:
      annotated  : np.ndarray (H, W, 3) RGB — image with boxes drawn
      original   : np.ndarray (H, W, 3) RGB — image at inference resolution, no boxes
      detections : list[dict] — class_id, class_name, confidence, box_xyxy, box_yolo
      summary    : dict — severity_label, severity_score, total_defects, by_class
    """
    img = image_rgb

    if resize and max(img.shape[:2]) > 1024:
        img = cv2.resize(img, (1024, 1024))

    model = load_model(weights)
    preds = model(img, augment=True)

    raw_boxes, confs, class_ids = [], [], []
    for r in preds:
        raw_boxes.extend(r.boxes.xyxy.cpu().numpy().tolist())
        confs.extend(r.boxes.conf.cpu().numpy().tolist())
        class_ids.extend(r.boxes.cls.cpu().numpy().astype(int).tolist())

    img_h, img_w = img.shape[:2]
    annotated  = img.copy()
    detections = []

    for box, conf, cls in zip(raw_boxes, confs, class_ids):
        if conf < thresholds.get(cls, 0.25):
            continue

        x1, y1, x2, y2 = map(int, box)
        color = CLASS_COLORS.get(cls, (255, 255, 255))
        label = f"{CLASS_NAMES[cls]} {conf:.2f}"
        _draw_box(annotated, x1, y1, x2, y2, label, color)

        cx, cy, w, h = _xyxy_to_yolo(*box, img_w, img_h)
        detections.append({
            "class_id":    cls,
            "class_name":  CLASS_NAMES[cls],
            "confidence":  round(conf, 3),
            "box_xyxy":    [x1, y1, x2, y2],
            "box_yolo":    [round(cx, 6), round(cy, 6), round(w, 6), round(h, 6)],
        })

    return {
        "annotated": annotated,
        "original": img,
        "detections": detections,
        "summary": _compute_summary(detections),
    }