app.py

"""
RT-DETR Object Detection — browser-compatible build (Plan A: single-file output)

Fixes applied vs. the upstream gradio/rt-detr-object-detection Space:
  1. Output is re-encoded to real H.264 (libx264 / yuv420p / +faststart) with ffmpeg.
     OpenCV's VideoWriter CANNOT produce H.264 in the pip opencv wheels used on
     HF Spaces — 'avc1'/'H264'/'X264' all return isOpened() == False.
  2. No streaming, no segments. One complete file is returned when done.
     This removes the entire class of MSE / chunk-boundary / short-video bugs.
  3. Every decoded+subsampled frame is written. The upstream code silently
     discarded the trailing partial batch (12 of 36 frames on a 3 s / 24 fps clip).
  4. Output dimensions forced even (H.264 yuv420p requirement).
  5. Frames are resized to the writer's exact size, so no silent write() no-ops.
  6. Optional HUD overlay showing per-frame detection count — useful when the
     purpose of the clip is to demonstrate the LIMITS of a general-purpose
     COCO-trained detector on industrial parts.
"""

import os
import time
import uuid
import shutil
import subprocess

import cv2
import numpy as np
import torch
import gradio as gr
import spaces
from PIL import Image
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor

from draw_boxes import draw_bounding_boxes

# --------------------------------------------------------------------------
# Startup banner — use this to PROVE your edited code is the code running.
# --------------------------------------------------------------------------
BUILD_TAG = "PLANA-H264-SINGLEFILE-v1"
print("=" * 70, flush=True)
print(f"[BOOT] RT-DETR build tag: {BUILD_TAG}", flush=True)
print(f"[BOOT] ffmpeg on PATH: {shutil.which('ffmpeg')}", flush=True)
print(f"[BOOT] opencv: {cv2.__version__}", flush=True)
print("=" * 70, flush=True)

MODEL_ID = "PekingU/rtdetr_r50vd"
image_processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
model = RTDetrForObjectDetection.from_pretrained(MODEL_ID).to("cuda")

SUBSAMPLE = 2        # keep every Nth source frame
CHUNK = 16           # frames per forward pass (memory bound, not timing bound)
WORK_DIR = os.path.join(os.getcwd(), "rtdetr_work")
os.makedirs(WORK_DIR, exist_ok=True)


def encode_h264(src_path: str) -> str:
    """Re-encode to a stream every browser can decode. Returns the new path."""
    dst_path = src_path.replace(".mp4", "_h264.mp4")
    cmd = [
        "ffmpeg", "-y", "-i", src_path,
        "-c:v", "libx264",
        "-profile:v", "baseline",   # widest possible decoder support
        "-level", "3.1",
        "-preset", "veryfast",
        "-crf", "23",
        "-pix_fmt", "yuv420p",      # mandatory: browsers reject yuv444p/yuvj420p
        "-movflags", "+faststart",  # moov atom first -> progressive playback
        "-an",
        dst_path,
    ]
    print(f"[H264] encoding -> {dst_path}", flush=True)
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0 or not os.path.exists(dst_path) or os.path.getsize(dst_path) == 0:
        print("[H264] FAILED. ffmpeg stderr tail:", flush=True)
        print("\n".join(proc.stderr.strip().splitlines()[-15:]), flush=True)
        raise gr.Error("H.264 re-encode failed — see container logs for ffmpeg output.")
    os.remove(src_path)
    print(f"[H264] OK  {os.path.getsize(dst_path)} bytes", flush=True)
    return dst_path


def _run_batch(frames_rgb, conf_threshold, height, width, writer, show_hud, counter):
    """Detect on a list of RGB frames, draw, write as BGR. Returns detections found."""
    if not frames_rgb:
        return 0

    inputs = image_processor(images=frames_rgb, return_tensors="pt").to("cuda")
    t0 = time.time()
    with torch.no_grad():
        outputs = model(**inputs)
    t_inf = time.time() - t0

    boxes = image_processor.post_process_object_detection(
        outputs,
        target_sizes=torch.tensor([(height, width)] * len(frames_rgb)),
        threshold=conf_threshold,
    )

    found = 0
    for array, box in zip(frames_rgb, boxes):
        n_det = int(box["scores"].shape[0])
        found += n_det
        pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
        frame_bgr = np.array(pil_image)[:, :, ::-1].copy()

        if show_hud:
            counter["i"] += 1
            label = f"frame {counter['i']:04d}   detections: {n_det}"
            colour = (0, 200, 0) if n_det else (0, 0, 255)
            cv2.rectangle(frame_bgr, (0, 0), (frame_bgr.shape[1], 26), (0, 0, 0), -1)
            cv2.putText(frame_bgr, label, (8, 18),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.55, colour, 1, cv2.LINE_AA)

        writer.write(frame_bgr)

    print(f"[BATCH] n={len(frames_rgb)} inference={t_inf:.3f}s detections={found}", flush=True)
    return found


@spaces.GPU(duration=180)
def detect(video, conf_threshold, show_hud):
    if not video:
        raise gr.Error("Please upload a video first.")

    print(f"\n[RUN] {BUILD_TAG} | conf={conf_threshold} | hud={show_hud}", flush=True)
    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        raise gr.Error("Could not open the uploaded video.")

    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    out_fps = max(1.0, src_fps / SUBSAMPLE)
    src_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    src_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    width = max(2, ((src_w // 2) // 2) * 2)      # half size, forced even
    height = max(2, ((src_h // 2) // 2) * 2)
    print(f"[RUN] src {src_w}x{src_h} @ {src_fps:.3f} -> out {width}x{height} @ {out_fps:.3f}",
          flush=True)

    raw_path = os.path.join(WORK_DIR, f"out_{uuid.uuid4().hex}.mp4")
    writer = cv2.VideoWriter(raw_path, cv2.VideoWriter_fourcc(*"mp4v"),
                             out_fps, (width, height))
    if not writer.isOpened():
        cap.release()
        raise gr.Error("OpenCV VideoWriter could not be opened.")

    batch, n_read, n_written, total_det = [], 0, 0, 0
    counter = {"i": 0}

    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            if n_read % SUBSAMPLE == 0:
                frame = cv2.resize(frame, (width, height))          # exact writer size
                batch.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                n_written += 1
                if len(batch) >= CHUNK:
                    total_det += _run_batch(batch, conf_threshold, height, width,
                                            writer, show_hud, counter)
                    batch = []
            n_read += 1

        # Trailing partial batch — the upstream Space threw these frames away.
        if batch:
            print(f"[FLUSH] trailing frames = {len(batch)}", flush=True)
            total_det += _run_batch(batch, conf_threshold, height, width,
                                    writer, show_hud, counter)
    finally:
        writer.release()
        cap.release()

    if n_written == 0:
        raise gr.Error("No frames were decoded from this video.")

    print(f"[RUN] read={n_read} written={n_written} "
          f"duration={n_written / out_fps:.2f}s total_detections={total_det}", flush=True)

    final_path = encode_h264(raw_path)

    summary = (
        f"Frames processed: {n_written} (from {n_read} source frames)\n"
        f"Output: {width}x{height} @ {out_fps:.2f} fps, {n_written / out_fps:.2f} s\n"
        f"Total detections at threshold {conf_threshold:.2f}: {total_det}\n"
        f"Codec: H.264 (libx264, yuv420p, faststart)"
    )
    if total_det == 0:
        summary += (
            "\n\nNote: zero detections. RT-DETR is trained on the 80 COCO classes; "
            "industrial parts are out-of-distribution, so an empty result is the "
            "expected outcome for a general-purpose detector."
        )
    return final_path, summary


with gr.Blocks(title="RT-DETR Object Detection") as app:
    gr.HTML(
        "<h1 style='text-align:center;margin-bottom:0'>Video Object Detection with RT-DETR</h1>"
        f"<p style='text-align:center;color:#888;font-size:0.85em'>build {BUILD_TAG}</p>"
    )
    with gr.Row():
        with gr.Column():
            video = gr.Video(label="Video Source")
            conf_threshold = gr.Slider(
                label="Confidence Threshold",
                minimum=0.0, maximum=1.0, step=0.01, value=0.30,
            )
            show_hud = gr.Checkbox(
                label="Overlay frame index + detection count", value=True
            )
            run_btn = gr.Button("Run detection", variant="primary")
        with gr.Column():
            output_video = gr.Video(label="Processed Video", autoplay=True)
            report = gr.Textbox(label="Report", lines=7, show_copy_button=True)

    run_btn.click(
        fn=detect,
        inputs=[video, conf_threshold, show_hud],
        outputs=[output_video, report],
    )
    video.upload(
        fn=detect,
        inputs=[video, conf_threshold, show_hud],
        outputs=[output_video, report],
    )

if __name__ == "__main__":
    app.launch()