Update app.py

app.py

@@ -1,179 +1,161 @@
-import spaces
-import gradio as gr
-import cv2
-from PIL import Image
-import torch
+"""
+RT-DETR Object Detection — streaming variant (Plan B).
+
+Use this ONLY if you specifically need progressive playback while processing.
+It is inherently more fragile than app.py (single-file). Differences from upstream:
+
+  * Every yielded segment is re-encoded to real H.264 with ffmpeg. Gradio's
+    streaming-outputs guide requires .mp4 or h.264-in-.ts for video chunks.
+  * Segment length is time-based and >= 1 s (Gradio requires this for smooth
+    playback). Upstream hardcoded exactly 2*desired_fps frames, so any video
+    whose frame count was not an exact multiple of that lost its tail.
+  * The trailing partial segment is flushed instead of discarded.
+  * Even dimensions enforced; frames resized to the writer's exact size.
+"""
+
+import os
 import time
-import numpy as np
 import uuid
-import os
+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
 
-print("=" * 60)
-print("RT-DETR Space – Browser-compatible version starting")
-print("=" * 60)
+BUILD_TAG = "PLANB-H264-STREAMING-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("=" * 70, flush=True)
 
-image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
-model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd").to("cuda")
+MODEL_ID = "PekingU/rtdetr_r50vd"
+image_processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
+model = RTDetrForObjectDetection.from_pretrained(MODEL_ID).to("cuda")
 
 SUBSAMPLE = 2
-SEGMENT_SECONDS = 1.0          # smaller segments → faster first output
-
+SEGMENT_SECONDS = 1.5          # must stay >= 1.0
+WORK_DIR = os.path.join(os.getcwd(), "rtdetr_stream")
+os.makedirs(WORK_DIR, exist_ok=True)
 
-def reencode_to_browser_h264(input_path: str) -> str:
-    """Force H.264 that Chrome can actually play."""
-    if not os.path.exists(input_path):
-        print(f"[ERROR] Input file does not exist: {input_path}")
-        return input_path
-
-    output_path = input_path.replace(".mp4", "_h264.mp4")
-    print(f"[RE-ENCODE] {input_path} → {output_path}")
 
+def encode_h264(src_path: str) -> str:
+    dst_path = src_path.replace(".mp4", "_h264.mp4")
     cmd = [
-        "ffmpeg", "-y",
-        "-i", input_path,
-        "-c:v", "libx264",
-        "-preset", "veryfast",
-        "-crf", "23",
-        "-pix_fmt", "yuv420p",
-        "-movflags", "+faststart",
-        "-an",
-        output_path
+        "ffmpeg", "-y", "-i", src_path,
+        "-c:v", "libx264", "-profile:v", "baseline", "-level", "3.1",
+        "-preset", "veryfast", "-crf", "23",
+        "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-an",
+        dst_path,
     ]
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+    if proc.returncode != 0 or not os.path.exists(dst_path):
+        print("[H264] FAILED:", proc.stderr[-800:], flush=True)
+        raise gr.Error("H.264 re-encode failed — see container logs.")
+    os.remove(src_path)
+    print(f"[H264] OK -> {dst_path} ({os.path.getsize(dst_path)} bytes)", flush=True)
+    return dst_path
 
-    try:
-        result = subprocess.run(
-            cmd,
-            check=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True
-        )
-        print(f"[RE-ENCODE] SUCCESS → {output_path}")
-        # remove intermediate
-        os.remove(input_path)
-        return output_path
-    except subprocess.CalledProcessError as e:
-        print(f"[RE-ENCODE] FAILED!")
-        print("ffmpeg stderr:", e.stderr)
-        # return original as last resort
-        return input_path
-    except FileNotFoundError:
-        print("[RE-ENCODE] ffmpeg not found in PATH!")
-        return input_path
-
-
-@spaces.GPU
+
+@spaces.GPU(duration=180)
 def stream_object_detection(video, conf_threshold):
-    print(f"\n[START] New video received. conf_threshold={conf_threshold}")
+    print(f"\n[RUN] {BUILD_TAG} | conf={conf_threshold}", flush=True)
     cap = cv2.VideoCapture(video)
     if not cap.isOpened():
-        raise gr.Error("Cannot open the uploaded video")
-
-    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30
-    desired_fps = max(1, fps // SUBSAMPLE)
-    print(f"[INFO] original fps={fps}, desired_fps={desired_fps}")
-
-    orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    width  = max(2, (orig_w // 2) // 2 * 2)
-    height = max(2, (orig_h // 2) // 2 * 2)
-    print(f"[INFO] output size = {width}x{height}")
-
-    frames_per_segment = max(1, int(desired_fps * SEGMENT_SECONDS))
-    print(f"[INFO] frames_per_segment = {frames_per_segment}")
-
-    iterating, frame = cap.read()
-    n_frames = 0
-    batch = []
-    name = f"/tmp/output_{uuid.uuid4()}.mp4"          # use /tmp for safety
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    segment_file = cv2.VideoWriter(name, fourcc, desired_fps, (width, height))
-
-    def flush_segment(current_batch, writer, current_name):
-        if len(current_batch) == 0:
-            writer.release()
+        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)
+    width = max(2, ((int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2) // 2) * 2)
+    height = max(2, ((int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2) // 2) * 2)
+    per_segment = max(1, int(round(out_fps * SEGMENT_SECONDS)))
+    print(f"[RUN] out {width}x{height} @ {out_fps:.2f}fps | {per_segment} frames/segment",
+          flush=True)
+
+    def new_writer():
+        p = os.path.join(WORK_DIR, f"seg_{uuid.uuid4().hex}.mp4")
+        return cv2.VideoWriter(p, cv2.VideoWriter_fourcc(*"mp4v"), out_fps, (width, height)), p
+
+    writer, path = new_writer()
+    batch, n_read, seg_no = [], 0, 0
+
+    def flush(frames, wr, p):
+        if not frames:
+            wr.release()
+            if os.path.exists(p):
+                os.remove(p)
             return None
-
-        print(f"[PROCESS] batch size = {len(current_batch)}")
-        inputs = image_processor(images=current_batch, return_tensors="pt").to("cuda")
-
-        start = time.time()
+        inputs = image_processor(images=frames, return_tensors="pt").to("cuda")
+        t0 = time.time()
         with torch.no_grad():
             outputs = model(**inputs)
-        print(f"time taken for inference {time.time() - start:.3f}")
-
-        start = time.time()
         boxes = image_processor.post_process_object_detection(
             outputs,
-            target_sizes=torch.tensor([(height, width)] * len(current_batch)),
-            threshold=conf_threshold
+            target_sizes=torch.tensor([(height, width)] * len(frames)),
+            threshold=conf_threshold,
         )
+        n_det = 0
+        for array, box in zip(frames, boxes):
+            n_det += int(box["scores"].shape[0])
+            pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
+            wr.write(np.array(pil_image)[:, :, ::-1].copy())
+        wr.release()
+        print(f"[SEG] frames={len(frames)} inference={time.time() - t0:.3f}s dets={n_det}",
+              flush=True)
+        return encode_h264(p)
 
-        for array, box in zip(current_batch, boxes):
-            pil_image = draw_bounding_boxes(
-                Image.fromarray(array), box, model, conf_threshold
-            )
-            frame_bgr = np.array(pil_image)[:, :, ::-1].copy()
-            writer.write(frame_bgr)
-
-        writer.release()
-        print(f"time taken for processing boxes {time.time() - start:.3f}")
-
-        playable = reencode_to_browser_h264(current_name)
-        print(f"[YIELD] → {playable}")
-        return playable
-
-    while iterating:
-        frame = cv2.resize(frame, (width, height))
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-        if n_frames % SUBSAMPLE == 0:
-            batch.append(frame)
-
-        if len(batch) >= frames_per_segment:
-            playable = flush_segment(batch, segment_file, name)
-            if playable:
-                yield playable
-            batch = []
-            name = f"/tmp/output_{uuid.uuid4()}.mp4"
-            segment_file = cv2.VideoWriter(name, fourcc, desired_fps, (width, height))
-
-        iterating, frame = cap.read()
-        n_frames += 1
-
-    # Final remaining frames
-    print(f"[FINAL] remaining frames in batch = {len(batch)}")
-    playable = flush_segment(batch, segment_file, name)
-    if playable:
-        yield playable
-
-    cap.release()
-    print("[END] Video processing finished\n")
-
-
-with gr.Blocks() as app:
-    gr.HTML("""
-        <h1 style='text-align: center'>
-        RT-DETR Object Detection<br>
-        <span style='font-size:0.55em;color:#555'>(Short video + Chrome H.264 fixed)</span>
-        </h1>
-    """)
+    try:
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+            if n_read % SUBSAMPLE == 0:
+                frame = cv2.resize(frame, (width, height))
+                batch.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                if len(batch) >= per_segment:
+                    out = flush(batch, writer, path)
+                    if out:
+                        seg_no += 1
+                        print(f"[YIELD] #{seg_no} {out}", flush=True)
+                        yield out
+                    batch = []
+                    writer, path = new_writer()
+            n_read += 1
+
+        # Trailing frames — upstream dropped these.
+        out = flush(batch, writer, path)
+        if out:
+            seg_no += 1
+            print(f"[YIELD] #{seg_no} (tail) {out}", flush=True)
+            yield out
+    finally:
+        cap.release()
+
+    print(f"[RUN] done. read={n_read} segments={seg_no}", flush=True)
+    if seg_no == 0:
+        raise gr.Error("No segments produced — the video decoded to zero frames.")
+
+
+with gr.Blocks(title="RT-DETR Streaming") as app:
+    gr.HTML(
+        "<h1 style='text-align:center;margin-bottom:0'>RT-DETR Object Detection (streaming)</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.05, value=0.30
+                minimum=0.0, maximum=1.0, step=0.01, value=0.30,
             )
         with gr.Column():
-            output_video = gr.Video(
-                label="Processed Video",
-                streaming=True,
-                autoplay=True
-            )
+            output_video = gr.Video(label="Processed Video", streaming=True, autoplay=True)
 
     video.upload(
         fn=stream_object_detection,