Update app.py

app.py

@@ -11,23 +11,26 @@ import subprocess
 from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
 from draw_boxes import draw_bounding_boxes
 
-# -------------------------------------------------
-# Model loading (same as original)
-# -------------------------------------------------
+print("=" * 60)
+print("RT-DETR Space – Browser-compatible version starting")
+print("=" * 60)
+
 image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
 model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd").to("cuda")
 
 SUBSAMPLE = 2
+SEGMENT_SECONDS = 1.0          # smaller segments → faster first output
 
 
 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)
-    """
+    """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}")
+
     cmd = [
         "ffmpeg", "-y",
         "-i", input_path,
@@ -36,145 +39,134 @@ def reencode_to_browser_h264(input_path: str) -> str:
         "-crf", "23",
         "-pix_fmt", "yuv420p",
         "-movflags", "+faststart",
-        "-an",                    # no audio
+        "-an",
         output_path
     ]
-    subprocess.run(
-        cmd,
-        check=True,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL
-    )
-    # Clean intermediate file
-    if os.path.exists(input_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
+        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
 def stream_object_detection(video, conf_threshold):
+    print(f"\n[START] New video received. conf_threshold={conf_threshold}")
     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}")
 
-    # 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
+    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 = []
-
-    # First segment
-    name = f"output_{uuid.uuid4()}.mp4"
-    # We deliberately use mp4v for writing reliability,
-    # then re-encode. This is the most robust pattern on HF Spaces.
+    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))
 
-    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) == 2 * desired_fps:
-            inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
-
-            print(f"starting batch of size {len(batch)}")
-            start = time.time()
-            with torch.no_grad():
-                outputs = model(**inputs)
-            end = time.time()
-            print("time taken for inference", end - start)
-
-            start = time.time()
-            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)
-
-            batch = []
-            segment_file.release()
-
-            # === 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, fourcc, desired_fps, (width, height))
+    def flush_segment(current_batch, writer, current_name):
+        if len(current_batch) == 0:
+            writer.release()
+            return None
 
-        iterating, frame = cap.read()
-        n_frames += 1
+        print(f"[PROCESS] batch size = {len(current_batch)}")
+        inputs = image_processor(images=current_batch, return_tensors="pt").to("cuda")
 
-    # Flush remaining frames if any
-    if batch:
-        inputs = image_processor(images=batch, return_tensors="pt").to("cuda")
+        start = 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(batch)),
+            target_sizes=torch.tensor([(height, width)] * len(current_batch)),
             threshold=conf_threshold
         )
-        for array, box in zip(batch, boxes):
+
+        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()
-            segment_file.write(frame_bgr)
+            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)
 
-        segment_file.release()
-        playable_name = reencode_to_browser_h264(name)
-        yield playable_name
+        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")
 
 
-# -------------------------------------------------
-# Gradio UI (unchanged structure)
-# -------------------------------------------------
 with gr.Blocks() as app:
-    gr.HTML(
-        """
+    gr.HTML("""
         <h1 style='text-align: center'>
-        Video Object Detection with RT-DETR (Browser-Compatible Fix)
+        RT-DETR Object Detection<br>
+        <span style='font-size:0.55em;color:#555'>(Short video + Chrome H.264 fixed)</span>
         </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'>Model</a>
-        </h3>
-        """
-    )
+    """)
     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.05, value=0.30
             )
         with gr.Column():
             output_video = gr.Video(