models/object_detector.py

"""
Lightweight object/face detection helper.

Samples frames from a video (OpenCV) and runs the available recognizers
(face -> fallback to body). Returns per-frame detections and writes
thumbnails to a temporary folder.
"""
from pathlib import Path
import tempfile
from typing import List, Dict, Optional, Callable

from utils.logger import get_logger

logger = get_logger("models.object_detector")


class ObjectDetector:
    """Detection-only helper that samples frames and returns detections.

    It will try to use `FaceRecognizer` first (if available) and fall back to
    `BodyRecognizer` (YOLO) if face code is not present.
    """

    def __init__(self, use_insightface: bool = True, use_opencv_face_fallback: bool = True):
        # Import lazily so the module can still be imported when optional deps
        # are missing.
        if use_insightface:
            try:
                from models.face_recognizer import FaceRecognizer
                try:
                    # Try to load the face model eagerly so detect_* calls work
                    self.face = FaceRecognizer(load_model=True)
                except Exception as e:
                    logger.warning(f"Could not load FaceRecognizer: {e}")
                    self.face = None
            except Exception:
                self.face = None
        else:
            self.face = None

        try:
            from models.body_recognizer import BodyRecognizer
            self.body = BodyRecognizer(load_model=False)
        except Exception:
            self.body = None

        # OpenCV Haar fallback for face detection (no extra model download).
        self.cv_face_cascade = None
        self.use_opencv_face_fallback = use_opencv_face_fallback
        try:
            import cv2

            cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
            cascade = cv2.CascadeClassifier(cascade_path)
            if not cascade.empty():
                self.cv_face_cascade = cascade
        except Exception:
            self.cv_face_cascade = None

    def _detect_faces_opencv(self, frame, min_size: int = 24) -> List[object]:
        """Detect faces with OpenCV Haar cascade as a lightweight fallback."""
        import cv2

        if self.cv_face_cascade is None:
            return []

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.cv_face_cascade.detectMultiScale(
            gray,
            scaleFactor=1.1,
            minNeighbors=4,
            minSize=(min_size, min_size),
        )

        class _Face:
            pass

        detections: List[object] = []
        for (x, y, w, h) in faces:
            d = _Face()
            d.bbox = (int(x), int(y), int(x + w), int(y + h))
            d.confidence = 0.8  # Haar cascade does not expose a calibrated score
            d.landmarks = None  # Marks this as a face-like detection for labeling
            detections.append(d)

        return detections

    def detect_faces_in_video(
        self,
        video_path: str,
        sample_rate: float = 1.0,
        min_confidence: float = 0.5,
        max_frames: Optional[int] = None,
        output_dir: Optional[str] = None,
        include_full_frame_fallback: bool = False,
        detection_type: str = "face",  # 'face', 'body' or 'both'
        progress_callback: Optional[Callable[[int, int, int, int], None]] = None,
        group_faces: bool = True,
        identity_similarity_threshold: float = 0.45,
        min_face_area: int = 2500,
        min_sharpness: float = 40.0,
        min_quality_score: float = 0.08,
        cross_identity_merge_threshold: float = 0.35,
    ) -> Dict:
        """Sample frames and detect faces/persons.

        Returns a dict with keys: `output_dir`, `fps`, `frame_count`, `detections`.
        Each detection is a dict: `{timestamp, frame_index, detection_index, bbox, confidence, thumbnail}`
        """
        import cv2
        import numpy as np

        video_path = str(video_path)
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Could not open video: {video_path}")

        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)

        # Calculate sampling interval in frames
        frame_interval = max(1, int(round(max(1.0, fps) / max(0.1, sample_rate))))
        total_sampled = max(1, (frame_count + frame_interval - 1) // frame_interval) if frame_count > 0 else 1

        if output_dir:
            out_dir = Path(output_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
        else:
            out_dir = Path(tempfile.mkdtemp(prefix="shortsmith_det_"))

        thumb_dir = out_dir / "thumbnails"
        thumb_dir.mkdir(parents=True, exist_ok=True)

        # We'll collect per-sampled-frame records. Each record has a
        # timestamp, frame_index and a list of detections (may be empty).
        frames: List[Dict] = []
        identities_state: Dict[int, Dict] = {}
        next_identity_id = 0

        def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
            a_n = np.linalg.norm(a)
            b_n = np.linalg.norm(b)
            if a_n == 0 or b_n == 0:
                return -1.0
            return float(np.dot(a, b) / (a_n * b_n))

        def _match_identity(embedding: Optional[np.ndarray]) -> Optional[int]:
            nonlocal next_identity_id
            if embedding is None:
                return None

            best_id = None
            best_sim = -1.0
            for ident_id, st in identities_state.items():
                centroid = st["embedding_sum"] / max(1, st["embedding_count"])
                sim = _cosine_similarity(embedding, centroid)
                if sim > best_sim:
                    best_sim = sim
                    best_id = ident_id

            if best_id is not None and best_sim >= identity_similarity_threshold:
                st = identities_state[best_id]
                st["embedding_sum"] = st["embedding_sum"] + embedding
                st["embedding_count"] += 1
                return best_id

            ident_id = next_identity_id
            next_identity_id += 1
            identities_state[ident_id] = {
                "embedding_sum": embedding.copy(),
                "embedding_count": 1,
                "detections": 0,
                "first_timestamp": None,
                "last_timestamp": None,
                "best_thumbnail": None,
                "best_quality": -1.0,
                "best_confidence": 0.0,
                "best_bbox": None,
                "best_frame_index": None,
                "best_timestamp": None,
                "occurrences": [],
                "last_occurrence_frame": None,
            }
            return ident_id

        frame_idx = 0
        sampled = 0
        if progress_callback is not None:
            try:
                progress_callback(0, total_sampled, 0, frame_count)
            except Exception:
                pass

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if frame_idx % frame_interval == 0:
                timestamp = frame_idx / fps

                # Choose detector(s) based on detection_type
                dets = []
                if detection_type in ("face", "both"):
                    if self.face is not None:
                        try:
                            dets = self.face.detect_faces(frame, max_faces=10, min_confidence=min_confidence)
                        except Exception:
                            dets = []
                    # Optional fallback to OpenCV Haar face detection
                    if not dets and self.use_opencv_face_fallback:
                        dets = self._detect_faces_opencv(frame)
                    # Face-only mode still needs a useful error if no face backend exists
                    if detection_type == "face" and self.face is None:
                        raise RuntimeError(
                            "InsightFace detector not available. "
                            "Install insightface dependencies and ensure model weights are available."
                        )

                # If requested and still empty, run body detector
                if (not dets) and detection_type in ("body", "both"):
                    if self.body is None:
                        if detection_type == "body":
                            raise RuntimeError("Body recognizer not available (ultralytics missing)")
                    else:
                        try:
                            dets = self.body.detect_persons(frame, min_confidence=min_confidence)
                        except Exception:
                            dets = []

                # If no detections found, optionally skip. If caller requests
                # a full-frame fallback (legacy behavior), generate a single
                # detection that covers the entire frame.
                if not dets and include_full_frame_fallback:
                    h, w = frame.shape[:2]
                    class _Full:
                        pass

                    d = _Full()
                    d.bbox = (0, 0, w, h)
                    d.confidence = 1.0
                    dets = [d]

                # If still no detections, skip this sampled frame entirely
                if not dets:
                    sampled += 1
                    if progress_callback is not None:
                        try:
                            progress_callback(sampled, total_sampled, frame_idx, frame_count)
                        except Exception:
                            pass
                    if max_frames and sampled >= max_frames:
                        break
                    frame_idx += 1
                    continue

                det_list: List[Dict] = []
                for i, d in enumerate(dets):
                    x1, y1, x2, y2 = d.bbox
                    # Clamp bbox to image
                    h, w = frame.shape[:2]
                    x1 = max(0, min(int(x1), w - 1))
                    x2 = max(0, min(int(x2), w))
                    y1 = max(0, min(int(y1), h - 1))
                    y2 = max(0, min(int(y2), h))

                    crop = frame[y1:y2, x1:x2]
                    if crop.size == 0:
                        continue
                    confidence = float(getattr(d, "confidence", 1.0))
                    area = float(max(1, (x2 - x1) * (y2 - y1)))
                    gray_crop = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
                    sharpness = float(cv2.Laplacian(gray_crop, cv2.CV_64F).var())
                    area_norm = min(1.0, area / 40000.0)
                    sharpness_norm = min(1.0, sharpness / 300.0)
                    quality_score = float(confidence * area_norm * sharpness_norm)

                    # Filter low-quality faces so gallery/exports contain only usable crops.
                    if area < float(min_face_area):
                        continue
                    if sharpness < float(min_sharpness):
                        continue
                    if quality_score < float(min_quality_score):
                        continue

                    identity_id = None
                    embedding = getattr(d, "embedding", None)
                    if group_faces and detection_type in ("face", "both"):
                        if embedding is not None:
                            try:
                                embedding = np.asarray(embedding, dtype=np.float32)
                                identity_id = _match_identity(embedding)
                            except Exception:
                                identity_id = None

                    should_write_thumb = True
                    if group_faces and identity_id is not None:
                        st = identities_state[identity_id]
                        should_write_thumb = quality_score > st["best_quality"]
                    thumb_str = None
                    if should_write_thumb:
                        thumb_path = thumb_dir / f"det_{sampled:06d}_{i}.jpg"
                        try:
                            cv2.imwrite(str(thumb_path), crop)
                            thumb_str = str(thumb_path)
                        except Exception:
                            thumb_str = None

                    # Decide label: face detector -> 'face', body detector -> 'person'
                    label = 'unknown'
                    # FaceDetection objects come from FaceRecognizer and have 'embedding' attr
                    if hasattr(d, 'embedding') or hasattr(d, 'landmarks'):
                        label = 'face'
                    else:
                        label = 'person'

                    if group_faces and identity_id is not None:
                        st = identities_state[identity_id]
                        st["detections"] += 1
                        if st["first_timestamp"] is None:
                            st["first_timestamp"] = float(timestamp)
                        st["last_timestamp"] = float(timestamp)
                        if st["last_occurrence_frame"] != int(frame_idx):
                            st["occurrences"].append({
                                "timestamp": float(timestamp),
                                "frame_index": int(frame_idx),
                                "bbox": (int(x1), int(y1), int(x2), int(y2)),
                                "confidence": confidence,
                                "sharpness": sharpness,
                                "quality_score": quality_score,
                            })
                            st["last_occurrence_frame"] = int(frame_idx)
                        if should_write_thumb and thumb_str is not None and quality_score > st["best_quality"]:
                            st["best_quality"] = quality_score
                            st["best_thumbnail"] = thumb_str
                            st["best_confidence"] = confidence
                            st["best_bbox"] = (int(x1), int(y1), int(x2), int(y2))
                            st["best_frame_index"] = int(frame_idx)
                            st["best_timestamp"] = float(timestamp)

                    det_list.append({
                        "detection_index": int(i),
                        "bbox": (int(x1), int(y1), int(x2), int(y2)),
                        "confidence": confidence,
                        "sharpness": sharpness,
                        "quality_score": quality_score,
                        "label": label,
                        "identity_id": int(identity_id) if identity_id is not None else None,
                        "thumbnail": thumb_str,
                    })

                # If all detections got filtered out by quality thresholds, skip frame.
                if not det_list:
                    sampled += 1
                    if progress_callback is not None:
                        try:
                            progress_callback(sampled, total_sampled, frame_idx, frame_count)
                        except Exception:
                            pass
                    if max_frames and sampled >= max_frames:
                        break
                    frame_idx += 1
                    continue

                # Record this sampled frame (may have empty detections if fallback disabled)
                frames.append({
                    "timestamp": float(timestamp),
                    "frame_index": int(frame_idx),
                    "detections": det_list,
                })

                sampled += 1
                if progress_callback is not None:
                    try:
                        progress_callback(sampled, total_sampled, frame_idx, frame_count)
                    except Exception:
                        pass
                if max_frames and sampled >= max_frames:
                    break

            frame_idx += 1

        cap.release()
        if progress_callback is not None:
            try:
                progress_callback(total_sampled, total_sampled, frame_count, frame_count)
            except Exception:
                pass

        def _merge_states(
            states: Dict[int, Dict],
            merge_threshold: float,
        ) -> tuple[Dict[int, Dict], Dict[int, int]]:
            ids = sorted(states.keys())
            if not ids:
                return {}, {}

            parent: Dict[int, int] = {i: i for i in ids}

            def find(x: int) -> int:
                while parent[x] != x:
                    parent[x] = parent[parent[x]]
                    x = parent[x]
                return x

            def union(a: int, b: int) -> None:
                ra, rb = find(a), find(b)
                if ra == rb:
                    return
                if ra < rb:
                    parent[rb] = ra
                else:
                    parent[ra] = rb

            centroids: Dict[int, np.ndarray] = {}
            for ident_id in ids:
                st = states[ident_id]
                if st["embedding_count"] > 0:
                    centroids[ident_id] = st["embedding_sum"] / max(1, st["embedding_count"])

            for i, a in enumerate(ids):
                emb_a = centroids.get(a)
                if emb_a is None:
                    continue
                for b in ids[i + 1:]:
                    emb_b = centroids.get(b)
                    if emb_b is None:
                        continue
                    sim = _cosine_similarity(emb_a, emb_b)
                    if sim >= merge_threshold:
                        union(a, b)

            groups: Dict[int, List[int]] = {}
            for ident_id in ids:
                root = find(ident_id)
                groups.setdefault(root, []).append(ident_id)

            merged_states: Dict[int, Dict] = {}
            old_to_new: Dict[int, int] = {}
            next_new_id = 0
            for root in sorted(groups.keys()):
                members = sorted(groups[root])
                first = states[members[0]]
                merged = {
                    "embedding_sum": first["embedding_sum"].copy(),
                    "embedding_count": int(first["embedding_count"]),
                    "detections": int(first["detections"]),
                    "first_timestamp": first["first_timestamp"],
                    "last_timestamp": first["last_timestamp"],
                    "best_thumbnail": first["best_thumbnail"],
                    "best_quality": float(first["best_quality"]),
                    "best_confidence": float(first["best_confidence"]),
                    "best_bbox": first["best_bbox"],
                    "best_frame_index": first["best_frame_index"],
                    "best_timestamp": first["best_timestamp"],
                    "occurrences": list(first["occurrences"]),
                    "last_occurrence_frame": first["last_occurrence_frame"],
                    "merged_from_ids": members.copy(),
                }

                for member in members[1:]:
                    st = states[member]
                    merged["embedding_sum"] = merged["embedding_sum"] + st["embedding_sum"]
                    merged["embedding_count"] += int(st["embedding_count"])
                    merged["detections"] += int(st["detections"])
                    if merged["first_timestamp"] is None or (
                        st["first_timestamp"] is not None and st["first_timestamp"] < merged["first_timestamp"]
                    ):
                        merged["first_timestamp"] = st["first_timestamp"]
                    if merged["last_timestamp"] is None or (
                        st["last_timestamp"] is not None and st["last_timestamp"] > merged["last_timestamp"]
                    ):
                        merged["last_timestamp"] = st["last_timestamp"]
                    merged["occurrences"].extend(st["occurrences"])
                    if st["best_quality"] > merged["best_quality"]:
                        merged["best_quality"] = float(st["best_quality"])
                        merged["best_thumbnail"] = st["best_thumbnail"]
                        merged["best_confidence"] = float(st["best_confidence"])
                        merged["best_bbox"] = st["best_bbox"]
                        merged["best_frame_index"] = st["best_frame_index"]
                        merged["best_timestamp"] = st["best_timestamp"]

                # De-duplicate occurrences by frame index.
                seen = set()
                uniq_occ = []
                for occ in sorted(merged["occurrences"], key=lambda o: float(o.get("timestamp", 0.0))):
                    fi = int(occ.get("frame_index", -1))
                    if fi in seen:
                        continue
                    seen.add(fi)
                    uniq_occ.append(occ)
                merged["occurrences"] = uniq_occ
                merged["last_occurrence_frame"] = int(uniq_occ[-1]["frame_index"]) if uniq_occ else None

                merged_states[next_new_id] = merged
                for member in members:
                    old_to_new[member] = next_new_id
                next_new_id += 1

            return merged_states, old_to_new

        merged_states, old_to_new = _merge_states(
            identities_state,
            merge_threshold=float(cross_identity_merge_threshold),
        )
        states_for_output = merged_states if group_faces else identities_state

        # Remap per-frame identity ids to merged ids so downstream UI uses unified face ids.
        if group_faces and old_to_new:
            for frame in frames:
                for det in frame.get("detections", []):
                    old_id = det.get("identity_id")
                    if old_id is not None and int(old_id) in old_to_new:
                        det["identity_id"] = int(old_to_new[int(old_id)])

        identities: List[Dict] = []
        for ident_id, st in sorted(states_for_output.items(), key=lambda kv: kv[0]):
            embedding_mean = None
            if st["embedding_count"] > 0:
                centroid = st["embedding_sum"] / max(1, st["embedding_count"])
                embedding_mean = centroid.tolist()
            identities.append({
                "identity_id": int(ident_id),
                "detections": int(st["detections"]),
                "first_timestamp": st["first_timestamp"],
                "last_timestamp": st["last_timestamp"],
                "representative_thumbnail": st["best_thumbnail"],
                "representative_confidence": float(st["best_confidence"]),
                "representative_quality": float(st["best_quality"]),
                "representative_bbox": st["best_bbox"],
                "representative_frame_index": st["best_frame_index"],
                "representative_timestamp": st["best_timestamp"],
                "occurrences": st["occurrences"],
                "embedding_mean": embedding_mean,
                "merged_from_ids": st.get("merged_from_ids", [int(ident_id)]),
            })

        return {
            "output_dir": str(out_dir),
            "fps": float(fps),
            "frame_count": frame_count,
            "frames": frames,
            "identities": identities,
        }


__all__ = ["ObjectDetector"]