""" 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( "

Video Object Detection with RT-DETR

" f"

build {BUILD_TAG}

" ) 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()