detection.py

"""
YOLO-based facade frame detector for Takeoff AI Annotator.

Model resolution order:
1. FACADE_MODEL_PATH — explicit local weights path.
2. FACADE_MODEL_REPO — HF Hub model repo (e.g. 'Infin8-AI/estimat8-vm-v1');
   downloads 'best.pt' at FACADE_MODEL_REVISION (default 'main') using
   HF_TOKEN for private repos. Pinning a revision is the rollback mechanism.
3. 'best.pt' next to this file.

Full drawing sheets must be detected with tiled inference (the model is
trained on 640x640 tiles; whole-sheet inference shrinks mullions below
detectable size). detect_rectangles() routes through detect_tiled()
automatically for images larger than one tile.
"""

from __future__ import annotations

import os
from pathlib import Path

import cv2
import numpy as np

# Tiling geometry — must match the training dataset builder.
TILE_SIZE = 640
TILE_STRIDE = 480
TILE_MERGE_IOU = 0.5
DETECTION_DPI = 200

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
_MODEL = None
_MODEL_PATH: str | None = None
_DEVICE: str | None = None


def _resolve_model_path() -> str:
    explicit = os.environ.get("FACADE_MODEL_PATH")
    if explicit:
        return explicit

    repo = os.environ.get("FACADE_MODEL_REPO")
    if repo:
        try:
            from huggingface_hub import hf_hub_download

            return hf_hub_download(
                repo_id=repo,
                filename="best.pt",
                revision=os.environ.get("FACADE_MODEL_REVISION") or "main",
                token=os.environ.get("HF_TOKEN") or None,
            )
        except Exception as exc:
            print(f"[detection] WARNING: HF Hub download failed ({exc}). Trying local best.pt.")

    return str(Path(__file__).parent / "best.pt")


def get_device() -> str:
    """'cuda' when available (local training boxes), else 'cpu' (HF Space)."""
    global _DEVICE
    if _DEVICE is None:
        try:
            import torch

            _DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
        except Exception:
            _DEVICE = "cpu"
    return _DEVICE


def _get_model():
    global _MODEL, _MODEL_PATH

    path = _resolve_model_path()

    if _MODEL is None or path != _MODEL_PATH:
        try:
            from ultralytics import YOLO

            _MODEL = YOLO(path)
            _MODEL_PATH = path
            print(f"[detection] Loaded YOLO model from: {path} (device={get_device()})")
        except Exception as exc:
            print(f"[detection] WARNING: Could not load YOLO model ({exc}). Falling back to OpenCV contour detection.")
            _MODEL = None

    return _MODEL


def get_model_version() -> str:
    """Return the version tag from versions.json if present, else 'unversioned'."""
    versions_path = Path(__file__).parent / "versions.json"
    if versions_path.exists():
        try:
            import json

            data = json.loads(versions_path.read_text())
            if isinstance(data, list) and data:
                return data[-1].get("version", "unversioned")
            if isinstance(data, dict):
                return data.get("version", "unversioned")
        except Exception:
            pass
    return "unversioned"


# ---------------------------------------------------------------------------
# Detection — returns list[dict] with keys x1, y1, x2, y2, conf
# ---------------------------------------------------------------------------


def detect_rectangles(
    image: np.ndarray,
    conf: float = 0.25,
    iou: float = 0.45,
    # Legacy parameters — only used when YOLO is not available (OpenCV fallback)
    min_area: int = 500,
    max_area: int = 500_000,
    epsilon_factor: float = 0.02,
    threshold: int = 127,
) -> list[dict]:
    """
    Detect facade frame panels in a construction plan image.

    Tries YOLO first; falls back to OpenCV contour detection if the model
    is unavailable.

    Returns a list of dicts: {x1, y1, x2, y2, conf} in image pixel coordinates.
    conf is the model confidence (0-1). OpenCV fallback sets conf=1.0.

    Images larger than one tile are detected via detect_tiled() — whole-sheet
    single-pass inference finds nothing on full drawings.
    """
    model = _get_model()

    if model is not None:
        h, w = image.shape[:2]
        if max(h, w) > TILE_SIZE:
            return detect_tiled(image, conf=conf, iou=iou)
        return _yolo_detect(model, image, conf=conf, iou=iou)

    return _cv_detect(image, min_area=min_area, max_area=max_area, epsilon_factor=epsilon_factor, threshold=threshold)


def detect_tiled(
    image: np.ndarray,
    conf: float = 0.25,
    iou: float = 0.45,
    tile: int = TILE_SIZE,
    stride: int = TILE_STRIDE,
    nms_iou: float = TILE_MERGE_IOU,
) -> list[dict]:
    """
    Sliding-window YOLO detection for full drawing sheets.

    Runs the model on tile x tile windows with the given stride (windows are
    added at the right/bottom edges so the whole sheet is covered), offsets
    each tile's boxes back into sheet coordinates, then merges duplicates from
    overlapping tiles with confidence-sorted greedy NMS at nms_iou.

    Returns boxes sorted top-to-bottom, left-to-right.
    """
    model = _get_model()
    if model is None:
        return []

    h, w = image.shape[:2]
    ys = list(range(0, max(h - tile, 0) + 1, stride)) or [0]
    xs = list(range(0, max(w - tile, 0) + 1, stride)) or [0]
    if ys[-1] + tile < h:
        ys.append(h - tile)
    if xs[-1] + tile < w:
        xs.append(w - tile)

    boxes: list[dict] = []
    for oy in ys:
        for ox in xs:
            window = image[oy : oy + tile, ox : ox + tile]
            for b in _yolo_detect(model, window, conf=conf, iou=iou):
                boxes.append(
                    {
                        "x1": b["x1"] + ox,
                        "y1": b["y1"] + oy,
                        "x2": b["x2"] + ox,
                        "y2": b["y2"] + oy,
                        "conf": b["conf"],
                    }
                )

    boxes.sort(key=lambda b: -b["conf"])
    kept: list[dict] = []
    for b in boxes:
        if all(_box_iou(b, k) < nms_iou for k in kept):
            kept.append(b)

    kept.sort(key=lambda b: (round(b["y1"], 2), round(b["x1"], 2)))
    return kept


def _box_iou(a: dict, b: dict) -> float:
    ix1, iy1 = max(a["x1"], b["x1"]), max(a["y1"], b["y1"])
    ix2, iy2 = min(a["x2"], b["x2"]), min(a["y2"], b["y2"])
    inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
    if inter == 0:
        return 0.0
    area_a = (a["x2"] - a["x1"]) * (a["y2"] - a["y1"])
    area_b = (b["x2"] - b["x1"]) * (b["y2"] - b["y1"])
    return inter / (area_a + area_b - inter)


# ---------------------------------------------------------------------------
# PDF rendering + coordinate conversion
# ---------------------------------------------------------------------------


def render_pdf_page(pdf_path: str, page_number: int, dpi: int = DETECTION_DPI) -> np.ndarray | None:
    """
    Render one PDF page (1-indexed) to an RGB numpy array at the given DPI.

    Opens the document per call and releases it immediately — safe for very
    large drawing sets (183 MB / 53 pages) because only one page's pixmap is
    ever held in memory.
    """
    try:
        import fitz  # PyMuPDF

        doc = fitz.open(pdf_path)
        try:
            idx = max(0, min(page_number - 1, len(doc) - 1))
            page = doc[idx]
            mat = fitz.Matrix(dpi / 72, dpi / 72)
            pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
            return np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
        finally:
            doc.close()
    except Exception as exc:
        print(f"[detection] PDF render error: {exc}")
        return None


def px_to_pt(value: float, dpi: int) -> float:
    """Render-pixel -> PDF-point (canonical ML label space)."""
    return value * 72.0 / dpi


def pt_to_px(value: float, dpi: int) -> float:
    """PDF-point -> render-pixel at the given DPI."""
    return value * dpi / 72.0


def _yolo_detect(
    model,
    image: np.ndarray,
    conf: float,
    iou: float,
) -> list[dict]:
    results = model.predict(
        source=image,
        conf=conf,
        iou=iou,
        imgsz=640,
        device=get_device(),
        verbose=False,
    )
    boxes: list[dict] = []
    for result in results:
        if result.boxes is None:
            continue
        for box in result.boxes:
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            c = float(box.conf[0])
            boxes.append({"x1": float(x1), "y1": float(y1), "x2": float(x2), "y2": float(y2), "conf": round(c, 4)})
    return boxes


def _cv_detect(
    image: np.ndarray,
    min_area: int,
    max_area: int,
    epsilon_factor: float,
    threshold: int,
) -> list[dict]:
    """OpenCV fallback — original contour-based detection, conf fixed at 1.0."""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image.copy()
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    _, binary = cv2.threshold(blurred, threshold, 255, cv2.THRESH_BINARY_INV)
    edges = cv2.Canny(blurred, 50, 150)
    combined = cv2.bitwise_or(binary, edges)

    contours, _ = cv2.findContours(combined, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    results: list[dict] = []
    seen: set[tuple[int, int, int, int]] = set()

    for contour in contours:
        area = cv2.contourArea(contour)
        if area < min_area or area > max_area:
            continue

        peri = cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon_factor * peri, True)

        if len(approx) != 4:
            continue

        x, y, w, h = cv2.boundingRect(approx)
        key = (x, y, x + w, y + h)
        if key in seen:
            continue
        seen.add(key)

        results.append({"x1": float(x), "y1": float(y), "x2": float(x + w), "y2": float(y + h), "conf": 1.0})

    return results


# ---------------------------------------------------------------------------
# Visualisation helper
# ---------------------------------------------------------------------------


def draw_detections(image: np.ndarray, boxes: list[dict]) -> np.ndarray:
    """Draw detection boxes on image. boxes is list[dict] with x1,y1,x2,y2,conf."""
    annotated = image.copy()
    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = int(box["x1"]), int(box["y1"]), int(box["x2"]), int(box["y2"])
        conf = box.get("conf", 1.0)
        cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 200, 0), 2)
        label = f"{i + 1} {conf:.2f}"
        cv2.putText(
            annotated,
            label,
            (x1 + 4, y1 + 16),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 200, 0),
            1,
            cv2.LINE_AA,
        )
    return annotated