Update app.py

app.py

@@ -1,12 +1,12 @@
 """
-RF-DETR Object Counter — CPU-optimized Gradio app for Hugging Face Spaces.
-Counts people, cars, trucks, and farm animals (cow, sheep/goat, horse/donkey,
-dog) in video using RF-DETR Medium + ByteTrack so each object is counted
-only once across the whole video.
+RF-DETR Object Counter — live-streaming Gradio app for Hugging Face Spaces.
+Annotated frames stream into the UI in real time while counts update as
+the model processes the video.
 """
 
 import os
 import tempfile
+import time
 from collections import defaultdict
 
 import cv2
@@ -14,15 +14,12 @@ import gradio as gr
 import numpy as np
 import supervision as sv
 import torch
-from rfdetr import RFDETRMedium
+from rfdetr import RFDETRNano
 
 # ---------------------------------------------------------------------------
 # Target classes (COCO indices)
 # ---------------------------------------------------------------------------
-# Note: COCO does NOT contain "goat" or "donkey". We approximate:
-#   goat   ~ sheep (closest 4-legged ruminant in COCO)
-#   donkey ~ horse (closest equid in COCO)
-# Counts for these will be roughly right; labels will say sheep/horse.
+# COCO has no "goat" or "donkey" — closest proxies: sheep ≈ goat, horse ≈ donkey
 TARGET_CLASSES = {
     0:  "person",
     2:  "car",
@@ -34,7 +31,6 @@ TARGET_CLASSES = {
 }
 TARGET_IDS = list(TARGET_CLASSES.keys())
 
-# Friendly UI labels
 DISPLAY_NAMES = {
     "person": "person",
     "car":    "car",
@@ -45,8 +41,7 @@ DISPLAY_NAMES = {
     "cow":    "cow",
 }
 
-# Per-class colours for the live overlay panel (BGR)
-CLASS_COLORS = {
+CLASS_COLORS = {  # BGR
     "person":  (66, 135, 245),
     "car":     (66, 245, 167),
     "truck":   (245, 66, 161),
@@ -56,18 +51,16 @@ CLASS_COLORS = {
     "cow":     (140, 90, 60),
 }
 
-# Example video lives next to app.py
 APP_DIR = os.path.dirname(os.path.abspath(__file__))
 EXAMPLE_VIDEO = os.path.join(APP_DIR, "example.mp4")
 
 # ---------------------------------------------------------------------------
-# Load model — pinned to CPU
+# Load model — CPU-pinned
 # ---------------------------------------------------------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Loading RF-DETR Medium on {DEVICE}…")
-MODEL = RFDETRMedium(device=DEVICE)
+print(f"Loading RF-DETR Nano on {DEVICE}…")
+MODEL = RFDETRNano(device=DEVICE)
 
-# optimize_for_inference is GPU-only (TensorRT-style ops). Skip on CPU.
 if DEVICE == "cuda":
     try:
         MODEL.optimize_for_inference()
@@ -75,7 +68,6 @@ if DEVICE == "cuda":
     except Exception as e:
         print(f"GPU optimization skipped: {e}")
 
-# Use a few threads for torch CPU inference; tune to your Space's vCPU count
 try:
     torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1))
 except Exception:
@@ -83,64 +75,147 @@ except Exception:
 
 print("Model ready.")
 
-# Annotators
 BOX_ANNOTATOR = sv.BoxAnnotator(thickness=2)
 LABEL_ANNOTATOR = sv.LabelAnnotator(text_scale=0.45, text_thickness=1, text_padding=3)
 
 
-def draw_counter_panel(frame: np.ndarray, counts: dict) -> np.ndarray:
-    """Translucent live-count panel in the top-left corner of the frame."""
+# ---------------------------------------------------------------------------
+# Drawing helpers
+# ---------------------------------------------------------------------------
+def draw_counter_panel(frame: np.ndarray, counts: dict,
+                       frame_idx: int, total_frames: int,
+                       fps_proc: float) -> np.ndarray:
+    """Translucent live-info panel in the top-left corner."""
     active = [(name, n) for name, n in counts.items() if n > 0]
-    if not active:
-        active = [("No targets yet", 0)]
 
-    panel_w = 280
-    panel_h = 40 + 22 * len(active)
+    rows = max(1, len(active)) + 1
+    panel_w = 320
+    panel_h = 28 + 22 * rows
+
     overlay = frame.copy()
-    cv2.rectangle(overlay, (12, 12), (12 + panel_w, 12 + panel_h), (20, 20, 20), -1)
+    cv2.rectangle(overlay, (12, 12), (12 + panel_w, 12 + panel_h),
+                  (20, 20, 20), -1)
     frame = cv2.addWeighted(overlay, 0.65, frame, 0.35, 0)
 
-    cv2.putText(frame, "LIVE COUNTS", (24, 38),
-                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2, cv2.LINE_AA)
+    cv2.putText(frame, "● LIVE", (24, 36),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (60, 220, 60), 2, cv2.LINE_AA)
+    cv2.putText(frame, f"frame {frame_idx}/{total_frames}  ·  {fps_proc:.1f} fps",
+                (90, 36),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (200, 200, 200), 1, cv2.LINE_AA)
 
-    y = 62
-    for name, n in active:
-        color = CLASS_COLORS.get(name, (200, 200, 200))
-        cv2.circle(frame, (28, y - 5), 5, color, -1)
-        display = DISPLAY_NAMES.get(name, name)
-        cv2.putText(frame, f"{display}: {n}", (44, y),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (240, 240, 240), 1, cv2.LINE_AA)
-        y += 22
+    y = 60
+    if not active:
+        cv2.putText(frame, "scanning…", (28, y),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (180, 180, 180), 1, cv2.LINE_AA)
+    else:
+        for name, n in active:
+            color = CLASS_COLORS.get(name, (200, 200, 200))
+            cv2.circle(frame, (28, y - 5), 5, color, -1)
+            display = DISPLAY_NAMES.get(name, name)
+            cv2.putText(frame, f"{display}: {n}", (44, y),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+                        (240, 240, 240), 1, cv2.LINE_AA)
+            y += 22
     return frame
 
 
-def process_video(video_path, confidence, frame_stride,
-                  progress=gr.Progress(track_tqdm=True)):
+def build_counts_html(unique_ids: dict, frame_idx: int, total: int,
+                      elapsed: float) -> str:
+    """Side-panel live counts as HTML cards."""
+    pct = (frame_idx / total * 100) if total else 0
+    cards = []
+    for name in TARGET_CLASSES.values():
+        n = len(unique_ids.get(name, set()))
+        display = DISPLAY_NAMES.get(name, name)
+        r, g, b = CLASS_COLORS.get(name, (200, 200, 200))[::-1]   # BGR->RGB
+        opacity = "1.0" if n > 0 else "0.35"
+        cards.append(
+            f'<div style="display:flex;justify-content:space-between;'
+            f'align-items:center;padding:8px 12px;margin:4px 0;'
+            f'border-radius:8px;background:rgba({r},{g},{b},0.10);'
+            f'border-left:4px solid rgb({r},{g},{b});opacity:{opacity};">'
+            f'<span style="font-weight:500;color:#111;">{display}</span>'
+            f'<span style="font-size:18px;font-weight:700;color:rgb({r},{g},{b});">{n}</span>'
+            f'</div>'
+        )
+    progress = (
+        f'<div style="margin:8px 0 14px 0;">'
+        f'<div style="display:flex;justify-content:space-between;font-size:12px;'
+        f'color:#6b7280;margin-bottom:4px;"><span>frame {frame_idx} / {total}</span>'
+        f'<span>{pct:.1f}% · {elapsed:.1f}s</span></div>'
+        f'<div style="height:6px;background:#e5e7eb;border-radius:3px;overflow:hidden;">'
+        f'<div style="height:100%;width:{pct}%;background:#6366f1;'
+        f'transition:width 0.2s;"></div></div></div>'
+    )
+    return progress + "".join(cards)
+
+
+def build_summary_md(unique_ids: dict) -> str:
+    total = sum(len(ids) for ids in unique_ids.values())
+    if total == 0:
+        return ("### ℹ️ No target objects detected.\n"
+                "Try lowering the confidence threshold or the frame stride.")
+    lines = [f"### ✅ Total unique objects detected: **{total}**", ""]
+    for name in TARGET_CLASSES.values():
+        n = len(unique_ids.get(name, set()))
+        if n > 0:
+            lines.append(f"- **{DISPLAY_NAMES.get(name, name).capitalize()}** — {n}")
+    return "\n".join(lines)
+
+
+def build_table(unique_ids: dict):
+    rows = []
+    for name in TARGET_CLASSES.values():
+        n = len(unique_ids.get(name, set()))
+        if n > 0:
+            rows.append([DISPLAY_NAMES.get(name, name).capitalize(), n])
+    return rows if rows else [["—", 0]]
+
+
+# ---------------------------------------------------------------------------
+# Main streaming generator
+# ---------------------------------------------------------------------------
+def process_video(video_path, confidence, frame_stride):
+    """
+    Generator: streams the annotated frame + live counts every iteration,
+    then yields the saved video and final table on completion.
+    """
     if video_path is None:
-        return None, "⚠️ Please upload a video first.", []
+        yield (None,
+               '<div style="padding:12px;color:#b91c1c;">⚠️ Please upload a video first.</div>',
+               "Submit a video to start.",
+               None,
+               [])
+        return
 
     video_info = sv.VideoInfo.from_video_path(video_path)
     frame_gen = sv.get_video_frames_generator(video_path)
     tracker = sv.ByteTrack(frame_rate=int(video_info.fps))
 
-    unique_ids = defaultdict(set)            # class_name -> {tracker_id, ...}
+    unique_ids = defaultdict(set)
     last_detections = sv.Detections.empty()
 
     out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    total = video_info.total_frames or 0
+    start_time = time.time()
+    last_yield = 0.0
+
+    yield (None,
+           build_counts_html(unique_ids, 0, total, 0.0),
+           "### 🎬 Starting analysis…",
+           None,
+           [])
+
+    last_rgb = None
 
     with sv.VideoSink(target_path=out_path, video_info=video_info) as sink:
-        for i, frame in enumerate(progress.tqdm(
-                frame_gen,
-                total=video_info.total_frames,
-                desc="Analyzing video")):
+        for i, frame in enumerate(frame_gen):
 
-            # Detect every Nth frame; reuse previous detections in-between
-            # so the output video stays smooth even with high stride.
+            # ---- Detection (every Nth frame) ----
             if i % frame_stride == 0:
                 rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 detections = MODEL.predict(rgb, threshold=confidence)
 
-                # Keep only the classes we care about
                 if len(detections) > 0:
                     mask = np.isin(detections.class_id, TARGET_IDS)
                     detections = detections[mask]
@@ -148,7 +223,6 @@ def process_video(video_path, confidence, frame_stride,
                 detections = tracker.update_with_detections(detections)
                 last_detections = detections
 
-                # Register unique tracker IDs per class
                 for cid, tid in zip(detections.class_id, detections.tracker_id):
                     if tid is None:
                         continue
@@ -158,13 +232,14 @@ def process_video(video_path, confidence, frame_stride,
             else:
                 detections = last_detections
 
-            # Annotate frame
+            # ---- Annotate ----
             if len(detections) > 0:
                 tids = (detections.tracker_id
                         if detections.tracker_id is not None
                         else [None] * len(detections))
                 labels = []
-                for cid, tid, conf in zip(detections.class_id, tids, detections.confidence):
+                for cid, tid, conf in zip(detections.class_id, tids,
+                                          detections.confidence):
                     name = TARGET_CLASSES.get(int(cid), "obj")
                     display = DISPLAY_NAMES.get(name, name)
                     tid_str = f"#{tid} " if tid is not None else ""
@@ -174,61 +249,61 @@ def process_video(video_path, confidence, frame_stride,
                 frame = LABEL_ANNOTATOR.annotate(frame, detections, labels)
 
             counts_now = {name: len(ids) for name, ids in unique_ids.items()}
-            frame = draw_counter_panel(frame, counts_now)
-            sink.write_frame(frame)
-
-    # ---------- Build summary outputs ----------
-    total = sum(len(ids) for ids in unique_ids.values())
-    if total == 0:
-        summary_md = ("### ℹ️ No target objects detected.\n"
-                      "Try lowering the confidence threshold or the frame stride.")
-    else:
-        lines = [f"### ✅ Total unique objects detected: **{total}**", ""]
-        for name in TARGET_CLASSES.values():
-            n = len(unique_ids.get(name, set()))
-            if n > 0:
-                display = DISPLAY_NAMES.get(name, name).capitalize()
-                lines.append(f"- **{display}** — {n}")
-        summary_md = "\n".join(lines)
-
-    table = []
-    for name in TARGET_CLASSES.values():
-        n = len(unique_ids.get(name, set()))
-        if n > 0:
-            display = DISPLAY_NAMES.get(name, name).capitalize()
-            table.append([display, n])
-    if not table:
-        table = [["—", 0]]
+            elapsed = time.time() - start_time
+            fps_proc = (i + 1) / elapsed if elapsed > 0 else 0
+            frame = draw_counter_panel(frame, counts_now, i + 1, total, fps_proc)
 
-    return out_path, summary_md, table
+            sink.write_frame(frame)
+            last_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+            # ---- Yield to UI (throttled to ~5 updates/sec) ----
+            now = time.time()
+            if now - last_yield > 0.20 or i == total - 1:
+                last_yield = now
+                yield (last_rgb,
+                       build_counts_html(unique_ids, i + 1, total, elapsed),
+                       "### 🔴 Live analysis in progress…",
+                       None,
+                       [])
+
+    # ---- Final yield: include saved video + summary table ----
+    elapsed = time.time() - start_time
+    yield (last_rgb,
+           build_counts_html(unique_ids, total, total, elapsed),
+           build_summary_md(unique_ids),
+           out_path,
+           build_table(unique_ids))
 
 
 # ---------------------------------------------------------------------------
 # UI
 # ---------------------------------------------------------------------------
 CUSTOM_CSS = """
-.gradio-container {max-width: 1200px !important; margin: auto;}
+.gradio-container {max-width: 1240px !important; margin: auto;}
 #title-row {text-align: center; padding: 8px 0 0 0;}
 #title-row h1 {font-weight: 700; letter-spacing: -0.5px; margin-bottom: 4px;}
 #title-row p {color: #6b7280; margin-top: 0;}
 .card {border: 1px solid #e5e7eb; border-radius: 14px; padding: 16px;
        background: #ffffff;}
+#live-frame img {border-radius: 10px;}
 footer {visibility: hidden;}
 """
 
-with gr.Blocks(title="RF-DETR Object Counter") as demo:
+with gr.Blocks(title="RF-DETR Live Object Counter") as demo:
 
     with gr.Row(elem_id="title-row"):
         gr.Markdown(
             """
-            # 🐄 RF-DETR Object Counter
-            Count **people, cars, trucks, and farm animals** in any video.
-            Powered by [RF-DETR Medium](https://github.com/roboflow/rf-detr) + ByteTrack —
-            each object is counted **only once** as it moves across frames.
+            # 🐄 RF-DETR Live Object Counter
+            Watch detections appear **frame by frame** as the model processes your video —
+            counts update in real time. Powered by
+            [RF-DETR Nano](https://github.com/roboflow/rf-detr) + ByteTrack
+            (each object counted only once).
             """
         )
 
     with gr.Row():
+        # ---------- Left: input ----------
         with gr.Column(scale=1):
             with gr.Group(elem_classes="card"):
                 gr.Markdown("### 📥 Input")
@@ -236,54 +311,70 @@ with gr.Blocks(title="RF-DETR Object Counter") as demo:
                     label="Upload a video",
                     sources=["upload"],
                     format="mp4",
-                    height=320,
+                    height=260,
                 )
 
                 with gr.Accordion("⚙️ Advanced settings", open=False):
                     confidence = gr.Slider(
                         minimum=0.1, maximum=0.9, value=0.45, step=0.05,
                         label="Confidence threshold",
-                        info="Higher = fewer but more certain detections.",
                     )
                     frame_stride = gr.Slider(
-                        minimum=1, maximum=15, value=5, step=1,
-                        label="Frame stride (CPU speed control)",
-                        info="Process every Nth frame. On CPU, 5–8 is a good balance.",
+                        minimum=1, maximum=15, value=3, step=1,
+                        label="Frame stride (CPU speed)",
+                        info="Detect every Nth frame. Higher = faster.",
                     )
 
-                submit_btn = gr.Button("🔍 Count Objects", variant="primary", size="lg")
+                submit_btn = gr.Button("▶️ Start Live Analysis",
+                                       variant="primary", size="lg")
 
                 gr.Markdown("#### 🎬 Example video")
                 gr.Examples(
                     examples=[[EXAMPLE_VIDEO]],
                     inputs=video_input,
-                    label=None,
                     examples_per_page=4,
                 )
 
+        # ---------- Right: live view ----------
+        with gr.Column(scale=2):
+            with gr.Group(elem_classes="card"):
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        gr.Markdown("### 🔴 Live View")
+                        live_frame = gr.Image(
+                            label=None,
+                            show_label=False,
+                            elem_id="live-frame",
+                            height=420,
+                        )
+                    with gr.Column(scale=1, min_width=220):
+                        gr.Markdown("### 📊 Live Counts")
+                        live_counts = gr.HTML(
+                            value=build_counts_html(defaultdict(set), 0, 0, 0)
+                        )
+
+    # ---------- Bottom: final results ----------
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group(elem_classes="card"):
+                gr.Markdown("### 📤 Final annotated video")
+                video_output = gr.Video(label="Download / replay", height=260)
         with gr.Column(scale=1):
             with gr.Group(elem_classes="card"):
-                gr.Markdown("### 📤 Annotated output")
-                video_output = gr.Video(label="Annotated video", height=320)
-                summary_output = gr.Markdown("Submit a video to see the results here.")
+                gr.Markdown("### 📈 Final totals")
+                summary_output = gr.Markdown("Run an analysis to see results.")
                 table_output = gr.Dataframe(
                     headers=["Class", "Unique count"],
                     datatype=["str", "number"],
-                    label="Per-class totals",
                     interactive=False,
                     wrap=True,
                 )
 
-    gr.Markdown(
-        """
-        **Detected categories:** person · car · truck · dog · horse / donkey · sheep / goat · cow
-        """
-    )
-
     submit_btn.click(
         fn=process_video,
         inputs=[video_input, confidence, frame_stride],
-        outputs=[video_output, summary_output, table_output],
+        outputs=[live_frame, live_counts, summary_output,
+                 video_output, table_output],
     )