split videos

README.md

@@ -36,5 +36,4 @@ in your browser, upload an MP4, and click "Detect".
 
 ## Notes
 - The first run downloads the wildfire detection model from Hugging Face.
-- The app prefers OpenCV for random frame access and falls back to imageio if
-  OpenCV is not available.
+- OpenCV is required for video decoding and frame sampling.

app.py

@@ -1,8 +1,14 @@
+import logging
 import os
+import shutil
+import subprocess
+import tempfile
+import time
+from collections import deque
+from contextlib import contextmanager
 
 import cv2
 import gradio as gr
-import imageio.v2 as imageio
 import numpy as np
 from PIL import Image, ImageDraw
 
@@ -10,6 +16,26 @@ from vision import Classifier
 from utils import box_iou, nms
 
 
+LOGGER = logging.getLogger(__name__)
+
+
+DEFAULT_SPLIT_CFG = {
+    "n_samples": 30,
+    "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,
+}
+
+
 def _sample_indices(total, n):
     if total <= 0:
         return []
@@ -18,86 +44,488 @@ def _sample_indices(total, n):
     return np.linspace(0, total - 1, n).astype(int).tolist()
 
 
-def _select_from_list(frames, n):
-    if not frames:
-        return []
-    indices = _sample_indices(len(frames), n)
-    return [frames[i] for i in indices]
-
+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 _extract_with_cv2(video_path, n):
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError("Could not open video file.")
 
-    frames = []
-    try:
-        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        if total > 0:
-            for idx in _sample_indices(total, n):
-                cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
-                ok, frame = cap.read()
-                if not ok:
-                    continue
-                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                frames.append(Image.fromarray(frame))
-            return frames
-
-        all_frames = []
-        while True:
-            ok, frame = cap.read()
-            if not ok:
-                break
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            all_frames.append(Image.fromarray(frame))
-        return _select_from_list(all_frames, n)
-    finally:
-        cap.release()
-
-
-def _get_imageio_count(reader):
-    try:
-        return reader.count_frames()
-    except Exception:
-        pass
+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:
-        meta = reader.get_meta_data()
-        count = meta.get("nframes")
-        if count and count != float("inf"):
-            return int(count)
+        return float(txt)
     except Exception:
-        pass
-    return None
+        return None
 
 
-def _extract_with_imageio(video_path, n):
-    reader = imageio.get_reader(video_path)
-    try:
-        total = _get_imageio_count(reader)
-        if total:
-            frames = []
-            for idx in _sample_indices(total, n):
-                frame = reader.get_data(idx)
-                frames.append(Image.fromarray(frame))
-            return frames
+def _probe_total_frames_ffprobe(video_path):
+    ffprobe = shutil.which("ffprobe")
+    if ffprobe is None:
+        return None
 
-        all_frames = [Image.fromarray(frame) for frame in reader]
-        return _select_from_list(all_frames, n)
-    finally:
-        reader.close()
+    # 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,
+    ]
+    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:
+                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,
+    ]
+    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
+    if proc.returncode != 0:
+        return None
+
+    lines = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
+    if len(lines) < 2:
+        return None
+
+    fps = _parse_fraction(lines[0])
+    duration = _parse_fraction(lines[1])
+    if fps is None or duration is None:
+        return None
+
+    estimate = int(round(fps * duration))
+    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")
+
+    total = _probe_total_frames_ffprobe(video_path)
+    if total is None or total <= 0:
+        raise RuntimeError("ffprobe could not determine total frame count")
+
+    indices = _sample_indices(total, int(n))
+    if not indices:
+        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,
+        ]
+        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False)
+        if proc.returncode != 0:
+            raise RuntimeError(proc.stderr.strip() or "ffmpeg extraction failed")
+
+        frames = []
+        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),
+        )
+        return frames
+
+
+def _extract_with_ffmpeg(video_path, n):
+    frames = _extract_bgr_with_ffmpeg(video_path, n)
+    return [Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames]
 
 
 def split_video(video_path, n=8):
     if not video_path or not os.path.exists(video_path):
         return []
-    # Prefer OpenCV for random access; use imageio as a fallback path if needed.
-    try:
-        return _extract_with_cv2(video_path, n)
-    except Exception:
-        return _extract_with_imageio(video_path, n)
+    return _extract_with_ffmpeg(video_path, n)
+
+
+@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):
+    frames = _extract_bgr_with_ffmpeg(video_path, int(n_samples))
+    for out_idx, frame in enumerate(frames):
+        yield out_idx, frame
+
+
+def iter_frames(video_path, n_samples, max_w, crop_y):
+    for out_idx, frame in _iter_sampled_frames(video_path, n_samples):
+        proc = frame
+        if max_w > 0 and proc.shape[1] != max_w:
+            scale = max_w / float(proc.shape[1])
+            proc = cv2.resize(
+                proc,
+                (max_w, int(proc.shape[0] * scale)),
+                interpolation=cv2.INTER_AREA,
+            )
 
+        if crop_y is not None:
+            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, :]
 
-model = Classifier(format="onnx")
+        yield out_idx, proc
+
+
+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=30,
+    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,
+):
+    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):
+            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:
+        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)
+
+    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):
+    if not segments:
+        LOGGER.info("Segment frame extraction | no segments found")
+        return []
+
+    normalized_segments = []
+    for start, end in segments:
+        s = max(0, int(start))
+        e = max(s, int(end))
+        normalized_segments.append((s, e))
+
+    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).
+    for frame_idx, frame in _iter_sampled_frames(video_path, n_samples=n_samples):
+        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:
+            grouped_frames[segment_idx].append(_bgr_to_pil(frame))
+            grouped_indices[segment_idx].append(frame_idx)
+
+    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),
+        )
+
+    return [frames for frames in grouped_frames if frames]
+
+
+def split_video_stable(video_path, split_cfg=None, fallback_n=30):
+    if not video_path or not os.path.exists(video_path):
+        return []
+
+    cfg = DEFAULT_SPLIT_CFG.copy()
+    if split_cfg:
+        cfg.update(split_cfg)
+
+    LOGGER.info("Split config | %s", cfg)
+
+    segments, _, _, _ = split_video_into_stable_segments_fast(video_path, **cfg)
+    frame_groups = extract_segment_frames(
+        video_path,
+        segments,
+        n_samples=cfg["n_samples"],
+    )
+
+    if frame_groups:
+        LOGGER.info(
+            "Split result | stable_splits=%d split_frame_counts=%s",
+            len(frame_groups),
+            [len(group) for group in frame_groups],
+        )
+        return frame_groups
+
+    LOGGER.info("Split result | no stable segment, using fallback sampling n=%d", fallback_n)
+    fallback_frames = split_video(video_path, n=fallback_n)
+    LOGGER.info("Fallback frame count | %d", len(fallback_frames))
+    return [fallback_frames] if fallback_frames else []
+
+
+model = Classifier(format="onnx", conf=0.05)
 
 
 def _resolve_video_path(video_input):
@@ -117,7 +545,7 @@ def _resolve_video_path(video_input):
     return None
 
 
-def _draw_detections(pil_img, preds):
+def _draw_detections(pil_img, preds, subtitle=None):
     img = pil_img.copy()
     draw = ImageDraw.Draw(img)
     width, height = img.size
@@ -133,32 +561,28 @@ def _draw_detections(pil_img, preds):
         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 infer(video_file):
-    video_path = _resolve_video_path(video_file)
-    frames = split_video(video_path, n=8)
-    if not frames:
-        return []
-    n_frames = len(frames)
-    boxes = np.zeros((0, 5), dtype=np.float64)
-    frame_preds = []
+def _combine_predictions_per_split(frame_preds):
+    n_frames = len(frame_preds)
+    if n_frames == 0:
+        return np.zeros((0, 5), dtype=np.float64)
 
-    for frame in frames:
-        bbox = np.asarray(model(frame), dtype=np.float64)
-        frame_preds.append(bbox)
+    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 []
+        return np.zeros((0, 5), dtype=np.float64)
 
     main_bboxes = np.asarray(nms(boxes), dtype=np.float64)
     if main_bboxes.size == 0:
-        return []
+        return np.zeros((0, 5), dtype=np.float64)
 
-    # Keep main boxes that appear in enough frames.
     matches_per_main = np.zeros(len(main_bboxes), dtype=int)
     for bbox in frame_preds:
         if bbox.size == 0:
@@ -166,23 +590,49 @@ def infer(video_file):
         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 // 2
-    kept_main = main_bboxes[keep_main] if np.any(keep_main) else np.zeros((0, 5), dtype=np.float64)
+    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)
 
-    if kept_main.size == 0:
+
+def infer(video_file):
+    video_path = _resolve_video_path(video_file)
+    LOGGER.info("Inference start | video=%s", video_path)
+    split_frames = split_video_stable(video_path)
+    if not split_frames:
+        LOGGER.info("Inference stop | no frames available")
         return []
 
     outputs = []
-    for main_box in kept_main:
-        for frame, bbox in zip(frames, frame_preds):
-            if bbox.size == 0:
-                continue
-            ious = box_iou(bbox[:, :4], main_box[:4].reshape(1, 4))
-            if (ious > 0).any():
-                match_idx = int(np.argmax(ious[0]))
-                outputs.append(_draw_detections(frame, bbox[match_idx : match_idx + 1]))
-                break
+    for split_idx, frames in enumerate(split_frames):
+        LOGGER.info("Inference split %d | frames=%d", split_idx + 1, len(frames))
+        frame_preds = []
+        for frame in frames:
+            bbox = np.asarray(model(frame), dtype=np.float64).reshape(-1, 5)
+            frame_preds.append(bbox)
+
+        kept_main = _combine_predictions_per_split(frame_preds)
+        LOGGER.info(
+            "Inference split %d | combined_detections=%d",
+            split_idx + 1,
+            len(kept_main),
+        )
+        if kept_main.size == 0:
+            continue
 
+        for det_idx, main_box in enumerate(kept_main):
+            for frame, bbox in zip(frames, frame_preds):
+                if bbox.size == 0:
+                    continue
+                ious = box_iou(bbox[:, :4], main_box[:4].reshape(1, 4))
+                if (ious > 0).any():
+                    match_idx = int(np.argmax(ious[0]))
+                    subtitle = f"split {split_idx + 1} / det {det_idx + 1}"
+                    outputs.append(_draw_detections(frame, bbox[match_idx : match_idx + 1], subtitle=subtitle))
+                    break
+
+    LOGGER.info("Inference done | output_images=%d", len(outputs))
     return outputs
 
 
@@ -190,7 +640,7 @@ with gr.Blocks() as demo:
     gr.Markdown("## Pyronear Wildfire Detection")
     with gr.Row():
         video_in = gr.Video(label="Upload MP4")
-    gallery_out = gr.Gallery(label="Wildfire detected", columns=2, height=360)
+    gallery_out = gr.Gallery(label="Wildfire detected (per split)", columns=2, height=360)
     run_btn = gr.Button("Detect")
     run_btn.click(fn=infer, inputs=video_in, outputs=gallery_out)
 

requirements.txt

@@ -2,7 +2,5 @@ gradio
 numpy
 Pillow
 opencv-python
-imageio
-imageio-ffmpeg
 onnxruntime
 tqdm