app.py

import logging
import os
import shutil
import subprocess
import tempfile
import time
from hashlib import sha1
from collections import deque
from contextlib import contextmanager

import cv2
import numpy as np
import streamlit as st
from PIL import Image, ImageDraw

from vision import Classifier
from utils import box_iou, nms


LOGGER = logging.getLogger(__name__)


PYRONEAR_LOGO_URL = (
    "https://raw.githubusercontent.com/pyronear/pyro-engine/develop/docs/source/_static/img/pyronear-logo-dark.png"
)


DEFAULT_SPLIT_CFG = {
    "n_samples": 16,
    "max_w": 400,
    "crop_y": (0.25, 0.90),
    "dx_threshold_px": 1.5,
    "min_inlier_ratio": 0.20,
    "min_stable_frames": 2,
    "smooth_window": 2,
    "orb_nfeatures": 800,
    "orb_fast_threshold": 12,
    "min_matches": 25,
    "keep_ratio": 0.4,
    "jump_meanabs_threshold": 18.0,
    "progress_every": 0,
}
INFER_BATCH_SIZE = max(1, int(os.getenv("INFER_BATCH_SIZE", "8")))


def _log_timing_summary(label, stats, wall_time=None, max_items=12):
    if not stats:
        LOGGER.info("%s timing | no data", label)
        return

    entries = sorted(
        ((name, float(value)) for name, value in stats.items() if value is not None),
        key=lambda item: item[1],
        reverse=True,
    )
    if wall_time is None:
        wall_time = stats.get("wall")

    step_entries = [(name, sec) for name, sec in entries if name != "wall"]
    parts = []
    if wall_time is not None:
        parts.append(f"wall={float(wall_time):.3f}s")
    for name, sec in step_entries[:max_items]:
        if wall_time and wall_time > 0:
            parts.append(f"{name}={sec:.3f}s ({(100.0 * sec / float(wall_time)):.1f}%)")
        else:
            parts.append(f"{name}={sec:.3f}s")
    remaining = max(0, len(step_entries) - max_items)
    if remaining:
        parts.append(f"+{remaining} more")

    LOGGER.info("%s timing | %s", label, " | ".join(parts))


def _sample_indices(total, n):
    if total <= 0:
        return []
    if total <= n:
        return list(range(total))
    return np.linspace(0, total - 1, n).astype(int).tolist()


def _format_idx_list(indices, max_items=40):
    if not indices:
        return "[]"
    values = [int(i) for i in indices]
    if len(values) <= max_items:
        return str(values)
    head = values[: max_items // 2]
    tail = values[-(max_items // 2) :]
    return f"{head} ... {tail} (len={len(values)})"


def _parse_fraction(value):
    if not value:
        return None
    txt = str(value).strip()
    if not txt or txt == "0/0":
        return None
    if "/" in txt:
        num, den = txt.split("/", 1)
        try:
            den_f = float(den)
            if den_f == 0:
                return None
            return float(num) / den_f
        except Exception:
            return None
    try:
        return float(txt)
    except Exception:
        return None


def _probe_total_frames_ffprobe(video_path):
    ffprobe = shutil.which("ffprobe")
    if ffprobe is None:
        return None

    timing = {}
    wall_t0 = time.perf_counter()
    video_name = os.path.basename(video_path)

    # Try direct frame count first.
    cmd = [
        ffprobe,
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-show_entries",
        "stream=nb_frames",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        video_path,
    ]
    with timer("ffprobe_nb_frames", timing):
        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
    if proc.returncode == 0:
        raw = proc.stdout.strip()
        if raw.isdigit():
            val = int(raw)
            if val > 0:
                timing["wall"] = time.perf_counter() - wall_t0
                _log_timing_summary(f"ffprobe ({video_name})", timing, wall_time=timing["wall"])
                return val

    # Fallback: estimate from duration * avg frame rate.
    cmd = [
        ffprobe,
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-show_entries",
        "stream=avg_frame_rate,duration",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        video_path,
    ]
    with timer("ffprobe_fps_duration", timing):
        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
    if proc.returncode != 0:
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"ffprobe ({video_name})", timing, wall_time=timing["wall"])
        return None

    lines = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
    if len(lines) < 2:
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"ffprobe ({video_name})", timing, wall_time=timing["wall"])
        return None

    fps = _parse_fraction(lines[0])
    duration = _parse_fraction(lines[1])
    if fps is None or duration is None:
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"ffprobe ({video_name})", timing, wall_time=timing["wall"])
        return None

    estimate = int(round(fps * duration))
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary(f"ffprobe ({video_name})", timing, wall_time=timing["wall"])
    return estimate if estimate > 0 else None


def _extract_bgr_with_ffmpeg(video_path, n):
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg is None:
        raise RuntimeError("ffmpeg is not available")

    timing = {}
    wall_t0 = time.perf_counter()
    video_name = os.path.basename(video_path)

    with timer("probe_total_frames", timing):
        total = _probe_total_frames_ffprobe(video_path)
    if total is None or total <= 0:
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"Frame extraction ({video_name})", timing, wall_time=timing["wall"])
        raise RuntimeError("ffprobe could not determine total frame count")

    with timer("sample_indices", timing):
        indices = _sample_indices(total, int(n))
    if not indices:
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"Frame extraction ({video_name})", timing, wall_time=timing["wall"])
        return []

    LOGGER.info(
        "Frame extraction | video=%s total_frames=%d n_samples=%d sampled_indices=%s",
        os.path.basename(video_path),
        total,
        len(indices),
        _format_idx_list(indices),
    )

    select_expr = "+".join(f"eq(n\\,{int(i)})" for i in indices)
    vf = f"select={select_expr}"

    with tempfile.TemporaryDirectory(prefix="ffmpeg_frames_") as tmpdir:
        pattern = os.path.join(tmpdir, "frame_%06d.jpg")
        cmd = [
            ffmpeg,
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            video_path,
            "-vf",
            vf,
            "-vsync",
            "vfr",
            "-q:v",
            "2",
            pattern,
        ]
        with timer("ffmpeg_extract", timing):
            proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
        if proc.returncode != 0:
            timing["wall"] = time.perf_counter() - wall_t0
            _log_timing_summary(f"Frame extraction ({video_name})", timing, wall_time=timing["wall"])
            raise RuntimeError(proc.stderr.strip() or "ffmpeg extraction failed")

        frames = []
        with timer("read_extracted_images", timing):
            for name in sorted(os.listdir(tmpdir)):
                if not name.lower().endswith(".jpg"):
                    continue
                frame = cv2.imread(os.path.join(tmpdir, name), cv2.IMREAD_COLOR)
                if frame is not None:
                    frames.append(frame)
        LOGGER.info(
            "Frame extraction done | video=%s extracted=%d requested=%d",
            os.path.basename(video_path),
            len(frames),
            len(indices),
        )
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary(f"Frame extraction ({video_name})", timing, wall_time=timing["wall"])
        return frames


def _extract_with_ffmpeg(video_path, n):
    timing = {}
    wall_t0 = time.perf_counter()
    with timer("extract_bgr", timing):
        frames = _extract_bgr_with_ffmpeg(video_path, n)
    with timer("bgr_to_pil", timing):
        pil_frames = [Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames]
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("Extract with ffmpeg", timing, wall_time=timing["wall"])
    return pil_frames


def split_video(video_path, n=8):
    if not video_path or not os.path.exists(video_path):
        return []
    timing = {}
    wall_t0 = time.perf_counter()
    with timer("extract_with_ffmpeg", timing):
        frames = _extract_with_ffmpeg(video_path, n)
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("split_video", timing, wall_time=timing["wall"])
    return frames


@contextmanager
def timer(name, stats):
    t0 = time.perf_counter()
    yield
    stats[name] = stats.get(name, 0.0) + (time.perf_counter() - t0)


def _iter_sampled_frames(video_path, n_samples, sampled_frames=None):
    timing = {}
    wall_t0 = time.perf_counter()
    if sampled_frames is None:
        with timer("extract_bgr_with_ffmpeg", timing):
            frames = _extract_bgr_with_ffmpeg(video_path, int(n_samples))
    else:
        with timer("reuse_sampled_frames", timing):
            frames = sampled_frames
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("Iter sampled frames", timing, wall_time=timing["wall"])
    for out_idx, frame in enumerate(frames):
        yield out_idx, frame


def iter_frames(video_path, n_samples, max_w, crop_y, sampled_frames=None):
    timing = {"resize": 0.0, "crop": 0.0}
    wall_t0 = time.perf_counter()
    frame_count = 0
    try:
        for out_idx, frame in _iter_sampled_frames(video_path, n_samples, sampled_frames=sampled_frames):
            frame_count += 1
            proc = frame
            if max_w > 0 and proc.shape[1] != max_w:
                t_resize = time.perf_counter()
                scale = max_w / float(proc.shape[1])
                proc = cv2.resize(
                    proc,
                    (max_w, int(proc.shape[0] * scale)),
                    interpolation=cv2.INTER_AREA,
                )
                timing["resize"] += time.perf_counter() - t_resize

            if crop_y is not None:
                t_crop = time.perf_counter()
                h = proc.shape[0]
                y0 = int(max(0.0, min(1.0, float(crop_y[0]))) * h)
                y1 = int(max(0.0, min(1.0, float(crop_y[1]))) * h)
                if y1 > y0:
                    proc = proc[y0:y1, :]
                timing["crop"] += time.perf_counter() - t_crop

            yield out_idx, proc
    finally:
        timing["wall"] = time.perf_counter() - wall_t0
        LOGGER.info(
            "iter_frames summary | n_samples=%d yielded=%d max_w=%d crop_y=%s",
            int(n_samples),
            frame_count,
            int(max_w),
            crop_y,
        )
        _log_timing_summary("iter_frames", timing, wall_time=timing["wall"])


def quick_jump_score(prev_gray, gray, small_w=160):
    h, w = prev_gray.shape[:2]
    if w > small_w:
        scale = small_w / float(w)
        prev_s = cv2.resize(prev_gray, (small_w, int(h * scale)), interpolation=cv2.INTER_AREA)
        gray_s = cv2.resize(gray, (small_w, int(h * scale)), interpolation=cv2.INTER_AREA)
    else:
        prev_s = prev_gray
        gray_s = gray

    diff = cv2.absdiff(prev_s, gray_s)
    return float(np.mean(diff))


def estimate_dx_orb_affine(prev_gray, gray, orb, bf, min_matches, keep_ratio, timing_pair):
    with timer("orb_detect_compute", timing_pair):
        kp1, des1 = orb.detectAndCompute(prev_gray, None)
        kp2, des2 = orb.detectAndCompute(gray, None)

    if des1 is None or des2 is None or len(kp1) < 8 or len(kp2) < 8:
        return None

    with timer("bf_match", timing_pair):
        matches = bf.match(des1, des2)

    if len(matches) < min_matches:
        return None

    with timer("match_sort_filter", timing_pair):
        matches = sorted(matches, key=lambda m: m.distance)
        keep_n = max(8, int(len(matches) * keep_ratio))
        matches = matches[:keep_n]

        pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
        pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])

    with timer("ransac_affine", timing_pair):
        M, inliers = cv2.estimateAffinePartial2D(
            pts1,
            pts2,
            method=cv2.RANSAC,
            ransacReprojThreshold=3.0,
            maxIters=1500,
            confidence=0.99,
        )

    if M is None:
        return None

    dx = float(M[0, 2])
    dy = float(M[1, 2])
    inlier_ratio = float(np.mean(inliers)) if inliers is not None else 0.0

    return {
        "dx": dx,
        "dy": dy,
        "score_dx": float(abs(dx)),
        "score_px": float(np.hypot(dx, dy)),
        "inlier_ratio": inlier_ratio,
        "matches": len(matches),
        "M": M,
    }


def split_video_into_stable_segments_fast(
    video_path,
    n_samples=16,
    max_w=400,
    crop_y=(0.25, 0.90),
    dx_threshold_px=1.5,
    min_inlier_ratio=0.20,
    min_stable_frames=2,
    smooth_window=2,
    orb_nfeatures=800,
    orb_fast_threshold=12,
    min_matches=25,
    keep_ratio=0.4,
    jump_meanabs_threshold=18.0,
    progress_every=200,
    sampled_frames=None,
):
    wall_t0 = time.perf_counter()
    timing_total = {}
    timing_pair = {}

    with timer("setup", timing_total):
        orb = cv2.ORB_create(nfeatures=orb_nfeatures, fastThreshold=orb_fast_threshold)
        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

    metrics = []
    prev_gray = None
    frame_count = 0

    with timer("loop_total", timing_total):
        for _, frame in iter_frames(
            video_path,
            n_samples=n_samples,
            max_w=max_w,
            crop_y=crop_y,
            sampled_frames=sampled_frames,
        ):
            frame_count += 1

            with timer("to_gray", timing_total):
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            if prev_gray is not None:
                with timer("quick_jump", timing_total):
                    q = quick_jump_score(prev_gray, gray)

                if q >= jump_meanabs_threshold:
                    metrics.append(
                        {
                            "dx": np.nan,
                            "dy": np.nan,
                            "score_dx": 1e9,
                            "score_px": 1e9,
                            "inlier_ratio": 0.0,
                            "matches": 0,
                            "M": None,
                            "quick_jump": q,
                        }
                    )
                else:
                    m = estimate_dx_orb_affine(
                        prev_gray,
                        gray,
                        orb=orb,
                        bf=bf,
                        min_matches=min_matches,
                        keep_ratio=keep_ratio,
                        timing_pair=timing_pair,
                    )
                    if m is None:
                        metrics.append(
                            {
                                "dx": np.nan,
                                "dy": np.nan,
                                "score_dx": 1e9,
                                "score_px": 1e9,
                                "inlier_ratio": 0.0,
                                "matches": 0,
                                "M": None,
                                "quick_jump": q,
                            }
                        )
                    else:
                        m["quick_jump"] = q
                        metrics.append(m)

                if progress_every and (len(metrics) % progress_every == 0):
                    print(f"processed pairs: {len(metrics)}")

            prev_gray = gray

    if frame_count < 2:
        timing_total["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary("Segmentation total", timing_total, wall_time=timing_total["wall"])
        if timing_pair:
            _log_timing_summary(
                "Segmentation pair internals",
                timing_pair,
                wall_time=max(timing_total.get("loop_total", 0.0), 1e-9),
            )
        return [], metrics, [], {"total": timing_total, "per_pair": timing_pair}

    with timer("post_smooth", timing_total):
        raw_dx = [m["score_dx"] for m in metrics]
        raw_inlier = [m["inlier_ratio"] for m in metrics]

        smoothed_dx = []
        q = deque(maxlen=max(1, int(smooth_window)))
        for v in raw_dx:
            if not np.isfinite(v):
                q.clear()
                smoothed_dx.append(np.nan)
            else:
                q.append(v)
                smoothed_dx.append(float(np.mean(q)))

    with timer("post_segments", timing_total):
        min_len = max(1, int(min_stable_frames))

        stable_flags = []
        for dx_s, r in zip(smoothed_dx, raw_inlier):
            if not np.isfinite(dx_s):
                stable_flags.append(False)
            else:
                stable_flags.append((dx_s < dx_threshold_px) and (r >= min_inlier_ratio))

        segments = []
        start = None
        for i, is_stable in enumerate(stable_flags):
            if is_stable and start is None:
                start = i
            if (not is_stable) and start is not None:
                end = i
                if (end - start) >= min_len:
                    segments.append((start, end))
                start = None

        if start is not None:
            end = len(stable_flags)
            if (end - start) >= min_len:
                segments.append((start, end))

    LOGGER.info(
        "Segmentation summary | sampled_frames=%d pair_metrics=%d stable_segments=%d",
        frame_count,
        len(metrics),
        len(segments),
    )
    if segments:
        LOGGER.info("Segment ranges (sample indices) | %s", segments)
    timing_total["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("Segmentation total", timing_total, wall_time=timing_total["wall"])
    if timing_pair:
        _log_timing_summary(
            "Segmentation pair internals",
            timing_pair,
            wall_time=max(timing_total.get("loop_total", 0.0), 1e-9),
        )

    return segments, metrics, smoothed_dx, {"total": timing_total, "per_pair": timing_pair}


def _bgr_to_pil(frame):
    return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))


def extract_segment_frames(video_path, segments, n_samples, sampled_frames=None):
    timing = {}
    wall_t0 = time.perf_counter()

    if not segments:
        LOGGER.info("Segment frame extraction | no segments found")
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary("Segment frame extraction", timing, wall_time=timing["wall"])
        return []

    with timer("normalize_segments", timing):
        normalized_segments = []
        for start, end in segments:
            s = max(0, int(start))
            e = max(s, int(end))
            normalized_segments.append((s, e))

    with timer("prepare_groups", timing):
        normalized_segments.sort(key=lambda x: x[0])
        grouped_frames = [[] for _ in normalized_segments]
        grouped_indices = [[] for _ in normalized_segments]
    segment_idx = 0

    # Detection runs on original sampled frames (no resize / no crop).
    to_pil_time = 0.0
    with timer("assign_frames_to_segments", timing):
        for frame_idx, frame in _iter_sampled_frames(video_path, n_samples=n_samples, sampled_frames=sampled_frames):
            while segment_idx < len(normalized_segments) and frame_idx > normalized_segments[segment_idx][1]:
                segment_idx += 1

            if segment_idx >= len(normalized_segments):
                break

            seg_start, seg_end = normalized_segments[segment_idx]
            if seg_start <= frame_idx <= seg_end:
                t_pil = time.perf_counter()
                grouped_frames[segment_idx].append(_bgr_to_pil(frame))
                to_pil_time += time.perf_counter() - t_pil
                grouped_indices[segment_idx].append(frame_idx)
    timing["to_pil"] = to_pil_time

    LOGGER.info(
        "Segment frame extraction summary | segments=%d n_samples=%d",
        len(normalized_segments),
        n_samples,
    )
    for seg_i, ((seg_start, seg_end), idx_list, frames) in enumerate(
        zip(normalized_segments, grouped_indices, grouped_frames),
        start=1,
    ):
        LOGGER.info(
            "Segment %d | requested_range=[%d,%d] matched_frames=%d matched_indices=%s",
            seg_i,
            seg_start,
            seg_end,
            len(frames),
            _format_idx_list(idx_list),
        )

    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("Segment frame extraction", timing, wall_time=timing["wall"])
    return [frames for frames in grouped_frames if frames]


def split_video_stable(video_path, split_cfg=None, fallback_n=16):
    if not video_path or not os.path.exists(video_path):
        return []

    timing = {}
    wall_t0 = time.perf_counter()
    cfg = DEFAULT_SPLIT_CFG.copy()
    if split_cfg:
        cfg.update(split_cfg)

    LOGGER.info("Split config | %s", cfg)

    with timer("extract_sampled_frames", timing):
        sampled_frames = _extract_bgr_with_ffmpeg(video_path, int(cfg["n_samples"]))

    with timer("split_video_into_stable_segments_fast", timing):
        segments, _, _, _ = split_video_into_stable_segments_fast(video_path, sampled_frames=sampled_frames, **cfg)
    with timer("extract_segment_frames", timing):
        frame_groups = extract_segment_frames(
            video_path,
            segments,
            n_samples=cfg["n_samples"],
            sampled_frames=sampled_frames,
        )

    if frame_groups:
        LOGGER.info(
            "Split result | stable_splits=%d split_frame_counts=%s",
            len(frame_groups),
            [len(group) for group in frame_groups],
        )
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary("split_video_stable", timing, wall_time=timing["wall"])
        return frame_groups

    LOGGER.info("Split result | no stable segment, using fallback sampling n=%d", fallback_n)
    if int(fallback_n) == int(cfg["n_samples"]):
        with timer("fallback_reuse_sampled_frames", timing):
            fallback_frames = [_bgr_to_pil(frame) for frame in sampled_frames]
    else:
        with timer("fallback_split_video", timing):
            fallback_frames = split_video(video_path, n=fallback_n)
    LOGGER.info("Fallback frame count | %d", len(fallback_frames))
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("split_video_stable", timing, wall_time=timing["wall"])
    return [fallback_frames] if fallback_frames else []


@st.cache_resource(show_spinner=False)
def _load_model():
    model_t0 = time.perf_counter()
    clf = Classifier(format="onnx", conf=0.05)
    LOGGER.info("Model init timing | wall=%.3fs", time.perf_counter() - model_t0)
    return clf


model = _load_model()


def _resolve_video_path(video_input):
    if not video_input:
        return None
    if isinstance(video_input, str):
        return video_input
    if isinstance(video_input, dict):
        for key in ("name", "path", "data", "video"):
            value = video_input.get(key)
            if isinstance(value, str) and os.path.exists(value):
                return value
    if isinstance(video_input, (list, tuple)):
        for value in video_input:
            if isinstance(value, str) and os.path.exists(value):
                return value
    return None


def _draw_detections(pil_img, preds, subtitle=None):
    img = pil_img.copy()
    draw = ImageDraw.Draw(img)
    width, height = img.size
    color = (255, 80, 0)
    preds = np.asarray(preds)

    for x1, y1, x2, y2, conf in preds:
        x1 = int(max(0.0, min(1.0, float(x1))) * width)
        y1 = int(max(0.0, min(1.0, float(y1))) * height)
        x2 = int(max(0.0, min(1.0, float(x2))) * width)
        y2 = int(max(0.0, min(1.0, float(y2))) * height)
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        draw.text((x1 + 4, y1 + 4), f"{conf:.2f}", fill=color)

    draw.text((6, 6), f"detections : {len(preds)}", fill=color)
    if subtitle:
        draw.text((6, 26), subtitle, fill=color)
    return img


def _combine_predictions_per_split(frame_preds):
    n_frames = len(frame_preds)
    if n_frames == 0:
        return np.zeros((0, 5), dtype=np.float64)

    boxes = np.zeros((0, 5), dtype=np.float64)
    for bbox in frame_preds:
        if bbox.size > 0:
            boxes = np.vstack([boxes, bbox])

    if boxes.size == 0:
        return np.zeros((0, 5), dtype=np.float64)

    main_bboxes = np.asarray(nms(boxes), dtype=np.float64)
    if main_bboxes.size == 0:
        return np.zeros((0, 5), dtype=np.float64)

    matches_per_main = np.zeros(len(main_bboxes), dtype=int)
    for bbox in frame_preds:
        if bbox.size == 0:
            continue
        ious = box_iou(bbox[:, :4], main_bboxes[:, :4])
        matches_per_main += (ious > 0).any(axis=1).astype(int)

    keep_main = matches_per_main > n_frames // 4
    if np.any(keep_main):
        return main_bboxes[keep_main]
    return np.zeros((0, 5), dtype=np.float64)


def infer(video_file):
    timing = {}
    wall_t0 = time.perf_counter()

    with timer("resolve_video_path", timing):
        video_path = _resolve_video_path(video_file)
    LOGGER.info("Inference start | video=%s", video_path)
    LOGGER.info("Inference config | batch_size=%d", INFER_BATCH_SIZE)
    with timer("split_video_stable", timing):
        split_frames = split_video_stable(video_path)
    if not split_frames:
        LOGGER.info("Inference stop | no frames available")
        timing["wall"] = time.perf_counter() - wall_t0
        _log_timing_summary("Inference", timing, wall_time=timing["wall"])
        return []

    outputs = []
    infer_model = 0.0
    combine_time = 0.0
    iou_time = 0.0
    draw_time = 0.0
    split_loop_time = 0.0
    for split_idx, frames in enumerate(split_frames):
        split_t0 = time.perf_counter()
        LOGGER.info("Inference split %d | frames=%d", split_idx + 1, len(frames))
        t_model = time.perf_counter()
        if hasattr(model, "infer_batch"):
            frame_preds = model.infer_batch(frames, batch_size=INFER_BATCH_SIZE)
        else:
            frame_preds = [model(frame) for frame in frames]
        frame_preds = [np.asarray(bbox, dtype=np.float64).reshape(-1, 5) for bbox in frame_preds]

        split_model = time.perf_counter() - t_model
        infer_model += split_model
        split_iou = 0.0
        split_draw = 0.0

        t_combine = time.perf_counter()
        kept_main = _combine_predictions_per_split(frame_preds)
        dt_combine = time.perf_counter() - t_combine
        combine_time += dt_combine
        LOGGER.info(
            "Inference split %d | combined_detections=%d",
            split_idx + 1,
            len(kept_main),
        )
        if kept_main.size == 0:
            split_elapsed = time.perf_counter() - split_t0
            split_loop_time += split_elapsed
            LOGGER.info(
                (
                    "Inference split %d timing | total=%.3fs | model=%.3fs | combine=%.3fs | "
                    "iou=%.3fs | draw=%.3fs | avg_model_ms=%.1f"
                ),
                split_idx + 1,
                split_elapsed,
                split_model,
                dt_combine,
                split_iou,
                split_draw,
                (1000.0 * split_model / max(len(frames), 1)),
            )
            continue

        for det_idx, main_box in enumerate(kept_main):
            for frame, bbox in zip(frames, frame_preds):
                if bbox.size == 0:
                    continue
                t_iou = time.perf_counter()
                ious = box_iou(bbox[:, :4], main_box[:4].reshape(1, 4))
                dt_iou = time.perf_counter() - t_iou
                split_iou += dt_iou
                iou_time += dt_iou
                if (ious > 0).any():
                    match_idx = int(np.argmax(ious[0]))
                    subtitle = f"segment {split_idx + 1} / detection {det_idx + 1}"
                    t_draw = time.perf_counter()
                    outputs.append(_draw_detections(frame, bbox[match_idx : match_idx + 1], subtitle=subtitle))
                    dt_draw = time.perf_counter() - t_draw
                    split_draw += dt_draw
                    draw_time += dt_draw
                    break

        split_elapsed = time.perf_counter() - split_t0
        split_loop_time += split_elapsed
        LOGGER.info(
            (
                "Inference split %d timing | total=%.3fs | model=%.3fs | combine=%.3fs | "
                "iou=%.3fs | draw=%.3fs | avg_model_ms=%.1f"
            ),
            split_idx + 1,
            split_elapsed,
            split_model,
            dt_combine,
            split_iou,
            split_draw,
            (1000.0 * split_model / max(len(frames), 1)),
        )

    timing["split_loop"] = split_loop_time
    timing["model_infer"] = infer_model
    timing["combine_predictions"] = combine_time
    timing["iou_matching"] = iou_time
    timing["draw_detections"] = draw_time
    timing["wall"] = time.perf_counter() - wall_t0
    _log_timing_summary("Inference", timing, wall_time=timing["wall"])
    LOGGER.info("Inference done | output_images=%d", len(outputs))
    return outputs


def _upload_signature(uploaded_file):
    buffer = uploaded_file.getbuffer()
    size = uploaded_file.size if uploaded_file.size is not None else len(buffer)
    digest = sha1(buffer).hexdigest()
    return (uploaded_file.name or "uploaded.mp4", int(size), digest)


def _write_uploaded_video(uploaded_file):
    ext = os.path.splitext(uploaded_file.name or "")[1] or ".mp4"
    with tempfile.NamedTemporaryFile(prefix="upload_", suffix=ext, delete=False) as tmp:
        tmp.write(uploaded_file.getbuffer())
        return tmp.name


def _render_outputs(outputs):
    if not outputs:
        st.warning("Aucune detection d'incendie trouvee dans cette video.")
        return

    st.subheader("Incendies detectes")
    columns = st.columns(2)
    for idx, image in enumerate(outputs):
        columns[idx % 2].image(image, caption=f"Detection {idx + 1}", use_container_width=True)


def main():
    st.set_page_config(page_title="Detection d'incendies Pyronear", layout="wide")
    st.image(PYRONEAR_LOGO_URL, width=220)
    st.title("Detection d'incendies Pyronear")
    st.write("Televersez un MP4 pour lancer la detection automatiquement.")

    uploaded = st.file_uploader("Televerser un MP4", type=["mp4"])
    if uploaded is None:
        st.info("En attente du televersement d'une video.")
        return

    signature = _upload_signature(uploaded)
    previous_signature = st.session_state.get("upload_signature")
    if signature != previous_signature:
        temp_path = None
        st.session_state["upload_signature"] = signature
        with st.spinner("Detection d'incendies en cours..."):
            try:
                temp_path = _write_uploaded_video(uploaded)
                st.session_state["output_images"] = infer(temp_path)
                st.session_state["inference_error"] = None
            except Exception as exc:
                LOGGER.exception("Inference failed")
                st.session_state["output_images"] = []
                st.session_state["inference_error"] = str(exc)
            finally:
                if temp_path and os.path.exists(temp_path):
                    os.remove(temp_path)

    if st.session_state.get("inference_error"):
        st.error(f"Echec de la detection : {st.session_state['inference_error']}")
        return

    _render_outputs(st.session_state.get("output_images", []))


if __name__ == "__main__":
    main()