detection.py

"""
detection.py — Core YOLO inference engine.
  - run_detection_single()   : one model on one image
  - run_detection_combined() : ALL 6 models, keep only the single
                                highest-confidence box overall
  - extract_video_frames()   : samples frames every N seconds from a video
"""

import os
import cv2
import random
import logging
from pathlib import Path
from typing import Dict, List, Any

from model_loader import load_model, get_model_info, get_severity, MODEL_REGISTRY

logger = logging.getLogger(__name__)

SEVERITY_COLOR_BGR = {
    "Critical": (0, 0, 220),
    "High":     (0, 100, 255),
    "Medium":   (0, 200, 255),
    "Low":      (0, 200, 0),
}

VIDEO_FRAME_INTERVAL_SECONDS = 2


def run_detection_single(image_path: str, model_key: str, output_dir: str, image_id: str, conf_threshold: float = 0.45) -> Dict[str, Any]:
    model = load_model(model_key)
    if model is None:
        detections = _mock_detections(model_key, conf_threshold)
    else:
        detections = _run_yolo(model, model_key, image_path, conf_threshold)
    return _finalize_result(image_path, detections, output_dir, image_id, model_key)


def run_detection_combined(image_path: str, output_dir: str, image_id: str, conf_threshold: float = 0.45) -> Dict[str, Any]:
    all_detections: List[Dict] = []
    for model_key in MODEL_REGISTRY.keys():
        model = load_model(model_key)
        if model is None:
            dets = _mock_detections(model_key, conf_threshold)
        else:
            dets = _run_yolo(model, model_key, image_path, conf_threshold)
        all_detections.extend(dets)

    best = [max(all_detections, key=lambda d: d["confidence"])] if all_detections else []
    return _finalize_result(image_path, best, output_dir, image_id, "combined (todos los modelos)")


def extract_video_frames(video_path: str, output_dir: str, video_id: str) -> List[Dict]:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video file: {video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    frame_interval = max(1, int(fps * VIDEO_FRAME_INTERVAL_SECONDS))

    frames_dir = Path(output_dir) / "video_frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    results = []
    frame_idx = 0
    saved_idx = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_interval == 0:
            timestamp_sec = round(frame_idx / fps, 1)
            frame_id = f"{video_id}_frame{saved_idx:03d}"
            frame_path = frames_dir / f"{frame_id}.jpg"
            cv2.imwrite(str(frame_path), frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
            results.append({"frame_path": str(frame_path), "frame_id": frame_id, "timestamp_sec": timestamp_sec})
            saved_idx += 1
        frame_idx += 1

    cap.release()
    logger.info("Extracted %d frames from video %s", len(results), video_id)
    return results


def _run_yolo(model, model_key: str, image_path: str, conf_threshold: float) -> List[Dict]:
    results = model(image_path, conf=conf_threshold, iou=0.45, imgsz=640, device="cpu", verbose=False)
    detections = []
    for r in results:
        boxes = r.boxes
        if boxes is None:
            continue
        for box in boxes:
            cls_id = int(box.cls[0])
            cls_name = model.names.get(cls_id, f"class_{cls_id}")
            conf = float(box.conf[0])
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            detections.append({
                "class": cls_name, "class_name": cls_name,
                "confidence": round(conf, 4),
                "severity": get_severity(model_key, cls_name),
                "bbox": [round(x1), round(y1), round(x2), round(y2)],
                "model_used": model_key,
            })
    return detections


def _mock_detections(model_key: str, conf_threshold: float) -> List[Dict]:
    info = get_model_info(model_key) or {}
    classes = info.get("classes", ["defect"])
    n = random.randint(1, 3)
    detections = []
    for _ in range(n):
        cls = random.choice(classes)
        conf = round(random.uniform(max(conf_threshold, 0.4), 0.95), 3)
        x1, y1 = random.randint(20, 200), random.randint(20, 200)
        x2, y2 = x1 + random.randint(80, 250), y1 + random.randint(80, 250)
        detections.append({
            "class": cls, "class_name": cls, "confidence": conf,
            "severity": get_severity(model_key, cls),
            "bbox": [x1, y1, x2, y2], "model_used": model_key,
        })
    return detections


def _finalize_result(image_path: str, detections: List[Dict], output_dir: str, image_id: str, model_label: str) -> Dict[str, Any]:
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"Cannot read image: {image_path}")

    annotated = img.copy()
    for det in detections:
        x1, y1, x2, y2 = det["bbox"]
        color = SEVERITY_COLOR_BGR.get(det["severity"], (255, 255, 255))
        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)
        label = f"{det['class_name'].replace('_', ' ')} {det['confidence']:.0%}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)
        cv2.rectangle(annotated, (x1, y1 - th - 8), (x1 + tw + 6, y1), color, -1)
        cv2.putText(annotated, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2)

    out_name = f"{image_id}_annotated.jpg"
    out_path = Path(output_dir) / out_name
    cv2.imwrite(str(out_path), annotated, [cv2.IMWRITE_JPEG_QUALITY, 92])

    severity_summary = {"Critical": 0, "High": 0, "Medium": 0, "Low": 0}
    for d in detections:
        severity_summary[d["severity"]] = severity_summary.get(d["severity"], 0) + 1

    overall_severity = "Low"
    for sev in ["Critical", "High", "Medium", "Low"]:
        if severity_summary.get(sev, 0) > 0:
            overall_severity = sev
            break

    return {
        "detections": detections,
        "total_defects": len(detections),
        "severity_summary": severity_summary,
        "overall_severity": overall_severity,
        "annotated_image": f"/outputs/{out_name}",
        "model_used": model_label,
    }