Update app.py

app.py

@@ -6,41 +6,82 @@ import torch
 import time
 import numpy as np
 import uuid
-
+import os
+import subprocess
 from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
-
 from draw_boxes import draw_bounding_boxes
 
+# -------------------------------------------------
+# Model loading (same as original)
+# -------------------------------------------------
 image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
 model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd").to("cuda")
 
-
 SUBSAMPLE = 2
 
+
+def reencode_to_browser_h264(input_path: str) -> str:
+    """
+    Force a browser-compatible H.264 stream.
+    - libx264
+    - yuv420p (mandatory for Chrome/Firefox/Safari)
+    - +faststart (moov atom at the beginning → progressive playback)
+    """
+    output_path = input_path.replace(".mp4", "_h264.mp4")
+    cmd = [
+        "ffmpeg", "-y",
+        "-i", input_path,
+        "-c:v", "libx264",
+        "-preset", "veryfast",
+        "-crf", "23",
+        "-pix_fmt", "yuv420p",
+        "-movflags", "+faststart",
+        "-an",                    # no audio
+        output_path
+    ]
+    subprocess.run(
+        cmd,
+        check=True,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL
+    )
+    # Clean intermediate file
+    if os.path.exists(input_path):
+        os.remove(input_path)
+    return output_path
+
+
 @spaces.GPU
 def stream_object_detection(video, conf_threshold):
     cap = cv2.VideoCapture(video)
 
-    video_codec = cv2.VideoWriter_fourcc(*"mp4v") # type: ignore
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 30
+    desired_fps = max(1, fps // SUBSAMPLE)
 
-    desired_fps = fps // SUBSAMPLE
-    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
+    # Force even dimensions (H.264 requirement)
+    orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    width  = (orig_w // 2) // 2 * 2
+    height = (orig_h // 2) // 2 * 2
 
     iterating, frame = cap.read()
-
     n_frames = 0
+    batch = []
 
+    # First segment
     name = f"output_{uuid.uuid4()}.mp4"
-    segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height)) # type: ignore
-    batch = []
+    # We deliberately use mp4v for writing reliability,
+    # then re-encode. This is the most robust pattern on HF Spaces.
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    segment_file = cv2.VideoWriter(name, fourcc, desired_fps, (width, height))
 
     while iterating:
-        frame = cv2.resize( frame, (0,0), fx=0.5, fy=0.5)
+        frame = cv2.resize(frame, (width, height))
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
         if n_frames % SUBSAMPLE == 0:
             batch.append(frame)
+
         if len(batch) == 2 * desired_fps:
             inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
 
@@ -55,40 +96,76 @@ def stream_object_detection(video, conf_threshold):
             boxes = image_processor.post_process_object_detection(
                 outputs,
                 target_sizes=torch.tensor([(height, width)] * len(batch)),
-                threshold=conf_threshold)
-            
-            for i, (array, box) in enumerate(zip(batch, boxes)):
-                pil_image = draw_bounding_boxes(Image.fromarray(array), box, model, conf_threshold)
-                frame = np.array(pil_image)
-                # Convert RGB to BGR
-                frame = frame[:, :, ::-1].copy()
-                segment_file.write(frame)
+                threshold=conf_threshold
+            )
+
+            for array, box in zip(batch, boxes):
+                pil_image = draw_bounding_boxes(
+                    Image.fromarray(array), box, model, conf_threshold
+                )
+                frame_bgr = np.array(pil_image)[:, :, ::-1].copy()
+                segment_file.write(frame_bgr)
 
             batch = []
             segment_file.release()
-            yield name
+
+            # === Critical fix: convert to real H.264 before yielding ===
+            playable_name = reencode_to_browser_h264(name)
+            yield playable_name
+
             end = time.time()
             print("time taken for processing boxes", end - start)
+
+            # Prepare next segment
             name = f"output_{uuid.uuid4()}.mp4"
-            segment_file = cv2.VideoWriter(name, video_codec, desired_fps, (width, height)) # type: ignore
+            segment_file = cv2.VideoWriter(name, fourcc, desired_fps, (width, height))
 
         iterating, frame = cap.read()
         n_frames += 1
 
+    # Flush remaining frames if any
+    if batch:
+        inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        boxes = image_processor.post_process_object_detection(
+            outputs,
+            target_sizes=torch.tensor([(height, width)] * len(batch)),
+            threshold=conf_threshold
+        )
+        for array, box in zip(batch, boxes):
+            pil_image = draw_bounding_boxes(
+                Image.fromarray(array), box, model, conf_threshold
+            )
+            frame_bgr = np.array(pil_image)[:, :, ::-1].copy()
+            segment_file.write(frame_bgr)
+
+        segment_file.release()
+        playable_name = reencode_to_browser_h264(name)
+        yield playable_name
+
+    cap.release()
 
+
+# -------------------------------------------------
+# Gradio UI (unchanged structure)
+# -------------------------------------------------
 with gr.Blocks() as app:
     gr.HTML(
         """
-    <h1 style='text-align: center'>
-    Video Object Detection with RT-DETR
-    </h1>
-    """)
+        <h1 style='text-align: center'>
+        Video Object Detection with RT-DETR (Browser-Compatible Fix)
+        </h1>
+        """
+    )
     gr.HTML(
         """
         <h3 style='text-align: center'>
-        <a href='https://arxiv.org/abs/2304.08069' target='_blank'>arXiv</a> | <a href='https://huggingface.co/PekingU/rtdetr_r101vd_coco_o365' target='_blank'>github</a>
+        <a href='https://arxiv.org/abs/2304.08069' target='_blank'>arXiv</a> |
+        <a href='https://huggingface.co/PekingU/rtdetr_r101vd_coco_o365' target='_blank'>Model</a>
         </h3>
-        """)
+        """
+    )
     with gr.Row():
         with gr.Column():
             video = gr.Video(label="Video Source")
@@ -100,7 +177,11 @@ with gr.Blocks() as app:
                 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,
@@ -108,5 +189,5 @@ with gr.Blocks() as app:
         outputs=[output_video],
     )
 
-if __name__ == '__main__':
-    app.launch()
+if __name__ == "__main__":
+    app.launch()