src/detector.py

"""
detector.py — per-frame slot occupancy inference

Public interface used by api/main.py:

    from src.detector import Detector

    det = Detector()                          # loads model + slot map once
    result = det.run(image_path)              # returns occupancy dict + saves frame
    # result = {
    #   "slot_001": {"status": "occupied", "confidence": 0.97},
    #   "slot_002": {"status": "empty",    "confidence": 0.88},
    #   ...
    # }
"""

import json
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torchvision import models, transforms

from src.database import init_db, log_occupancy

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR       = Path(__file__).resolve().parent.parent          # backend/
MODEL_PATH     = BASE_DIR / "models" / "slot_classifier.pth"
SLOT_MAP_PATH  = BASE_DIR / "data"   / "raw" / "slot_map.json"
FRAME_OUT_PATH = BASE_DIR / "data"   / "annotated_frame.jpg"

DEVICE    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMG_SIZE  = 224
BATCH_SIZE = 64

# Colours for bounding boxes drawn on the annotated frame (BGR for OpenCV)
COLOR_OCCUPIED = (0,   0,   255)   # red
COLOR_EMPTY    = (0,   200,  0)    # green
BOX_THICKNESS  = 2

# ---------------------------------------------------------------------------
# Transform  (identical to training and evaluation)
# ---------------------------------------------------------------------------
_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


def _build_model() -> nn.Module:
    model = models.mobilenet_v2(weights=None)
    model.classifier[1] = nn.Linear(model.last_channel, 2)
    return model


def _load_weights(model: nn.Module, path: Path) -> nn.Module:
    checkpoint = torch.load(path, map_location=DEVICE, weights_only=False)
    if isinstance(checkpoint, dict):
        if "model_state" in checkpoint:
            state = checkpoint["model_state"]
        elif "model_state_dict" in checkpoint:
            state = checkpoint["model_state_dict"]
        elif "state_dict" in checkpoint:
            state = checkpoint["state_dict"]
        else:
            state = checkpoint
    else:
        state = checkpoint
    model.load_state_dict(state)
    model.to(DEVICE)
    model.eval()
    return model


# ---------------------------------------------------------------------------
# Detector class
# ---------------------------------------------------------------------------
class Detector:
    """
    Loads the model and slot map once at construction time.
    Call .run(image_path) for each new frame.
    """

    def __init__(self):
        init_db()   # ensure tables exist before any logging

        print(f"[Detector] Loading model from {MODEL_PATH} on {DEVICE} ...")
        self.model = _load_weights(_build_model(), MODEL_PATH)

        print(f"[Detector] Loading slot map from {SLOT_MAP_PATH} ...")
        with open(SLOT_MAP_PATH, "r") as f:
            raw = json.load(f)

        # slot_map.json structure:
        # {
        #   "reference_frame": "...",
        #   "total_slots": 100,
        #   "slots": [
        #     {"slot_id": 1, "x": 139, "y": 165, "w": 23, "h": 40, "cx": ..., "cy": ...},
        #     ...
        #   ]
        # }
        slot_list = raw["slots"]
        self.slots = [
            {
                "slot_id": f"slot_{s['slot_id']:03d}",   # 1 → "slot_001"
                "bbox":    [s["x"], s["y"], s["w"], s["h"]],
            }
            for s in slot_list
        ]

        print(f"[Detector] {len(self.slots)} slots loaded.")

    # -----------------------------------------------------------------------
    def run(self, image_path: str | Path) -> dict:
        """
        Run inference on one parking lot image.

        Parameters
        ----------
        image_path : path to a JPG/PNG image from the PKLot dataset

        Returns
        -------
        dict  {slot_id: {"status": "occupied"|"empty", "confidence": float}}
        """
        image_path = Path(image_path)
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")

        # --- Load image once (PIL for cropping, OpenCV copy for annotation) ---
        pil_img = Image.open(image_path).convert("RGB")
        cv_img  = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

        # --- Crop every slot ------------------------------------------------
        crops    = []
        slot_ids = []

        for slot in self.slots:
            slot_id = slot["slot_id"]
            x, y, w, h = [int(v) for v in slot["bbox"]]

            # Guard against out-of-bounds or degenerate boxes
            img_w, img_h = pil_img.size
            x  = max(0, min(x, img_w - 1))
            y  = max(0, min(y, img_h - 1))
            x2 = min(x + w, img_w)
            y2 = min(y + h, img_h)

            if (x2 - x) < 2 or (y2 - y) < 2:
                continue

            crop = pil_img.crop((x, y, x2, y2))
            crops.append(_transform(crop))
            slot_ids.append((slot_id, x, y, x2, y2))

        if not crops:
            raise RuntimeError("No valid crops produced — check slot_map.json bbox values.")

        # --- Batched inference ----------------------------------------------
        all_preds  = []
        all_probs  = []

        softmax = nn.Softmax(dim=1)

        for start in range(0, len(crops), BATCH_SIZE):
            batch   = torch.stack(crops[start : start + BATCH_SIZE]).to(DEVICE)
            with torch.no_grad():
                logits = self.model(batch)          # (B, 2)
                probs  = softmax(logits)            # (B, 2)
                preds  = logits.argmax(dim=1)       # (B,)

            all_preds.extend(preds.cpu().tolist())
            all_probs.extend(probs.cpu().tolist())

        # --- Build result dict + annotate image ----------------------------
        result   = {}
        db_rows  = []

        for i, (slot_id, x, y, x2, y2) in enumerate(slot_ids):
            pred       = all_preds[i]               # 0 = empty, 1 = occupied
            confidence = float(all_probs[i][pred])
            status     = "occupied" if pred == 1 else "empty"

            result[slot_id] = {"status": status, "confidence": confidence}
            db_rows.append({"slot_id": slot_id, "status": status, "confidence": confidence})

            # Draw bounding box on the OpenCV copy
            color = COLOR_OCCUPIED if pred == 1 else COLOR_EMPTY
            cv2.rectangle(cv_img, (x, y), (x2, y2), color, BOX_THICKNESS)

            # Small label (slot id + confidence) above each box
            label = f"{slot_id} {confidence:.2f}"
            cv2.putText(
                cv_img, label,
                (x, max(y - 4, 10)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.35, color, 1, cv2.LINE_AA,
            )

        # --- Save annotated frame ------------------------------------------
        FRAME_OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(FRAME_OUT_PATH), cv_img)

        # --- Log to SQLite --------------------------------------------------
        log_occupancy(db_rows)

        return result


# ---------------------------------------------------------------------------
# Quick smoke-test:  python -m src.detector  (run from backend/)
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import sys, pprint

    if len(sys.argv) < 2:
        print("Usage: python -m src.detector <path_to_image>")
        sys.exit(1)

    det    = Detector()
    result = det.run(sys.argv[1])

    occupied = sum(1 for v in result.values() if v["status"] == "occupied")
    empty    = len(result) - occupied

    print(f"\nSlots processed : {len(result)}")
    print(f"Occupied        : {occupied}")
    print(f"Empty           : {empty}")
    print(f"\nAnnotated frame : {FRAME_OUT_PATH}")
    print("\nFirst 5 results:")
    pprint.pprint(dict(list(result.items())[:5]))