Update app.py

app.py

@@ -1,7 +1,6 @@
 """
-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.
+RF-DETR Truck Counter — counts trucks crossing a fixed horizontal line
+at the center of the video (either direction counts as one crossing).
 """
 
 import os
@@ -19,17 +18,19 @@ from rfdetr import RFDETRNano
 # ---------------------------------------------------------------------------
 # Target classes (COCO indices)
 # ---------------------------------------------------------------------------
-# COCO has no "goat" or "donkey" — closest proxies: sheep ≈ goat, horse ≈ donkey
+# We still DETECT all of these so users see them tracked on screen,
+# but only TRUCKS contribute to the line-crossing count.
 TARGET_CLASSES = {
     0:  "person",
     2:  "car",
     7:  "truck",
     16: "dog",
-    17: "horse",   # also catches donkeys
-    18: "sheep",   # also catches goats
+    17: "horse",
+    18: "sheep",
     19: "cow",
 }
 TARGET_IDS = list(TARGET_CLASSES.keys())
+TRUCK_CLASS_ID = 7
 
 DISPLAY_NAMES = {
     "person": "person",
@@ -75,138 +76,138 @@ except Exception:
 
 print("Model ready.")
 
-BOX_ANNOTATOR = sv.BoxAnnotator(thickness=2)
-LABEL_ANNOTATOR = sv.LabelAnnotator(text_scale=0.45, text_thickness=1, text_padding=3)
-
 
 # ---------------------------------------------------------------------------
-# Drawing helpers
+# Draw a custom truck counter directly on the centerline
 # ---------------------------------------------------------------------------
-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]
-
-    rows = max(1, len(active)) + 1
-    panel_w = 320
-    panel_h = 28 + 22 * rows
+def draw_truck_label_on_line(frame: np.ndarray, line_y: int, total: int) -> np.ndarray:
+    text = f"TRUCKS CROSSED: {total}"
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    scale = max(0.6, frame.shape[1] / 1600)
+    thickness = max(2, int(2 * scale))
+
+    (tw, th), baseline = cv2.getTextSize(text, font, scale, thickness)
+    pad_x, pad_y = 14, 10
+    box_w = tw + 2 * pad_x
+    box_h = th + 2 * pad_y + baseline
+
+    cx = frame.shape[1] // 2
+    x1 = cx - box_w // 2
+    y1 = line_y - box_h // 2
+    x2 = x1 + box_w
+    y2 = y1 + box_h
 
     overlay = frame.copy()
-    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", (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 = 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
+    cv2.rectangle(overlay, (x1, y1), (x2, y2), (245, 66, 161), -1)  # truck-pink
+    frame = cv2.addWeighted(overlay, 0.88, frame, 0.12, 0)
+    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 255), 2)
+
+    text_x = x1 + pad_x
+    text_y = y1 + pad_y + th
+    cv2.putText(frame, text, (text_x, text_y),
+                font, scale, (255, 255, 255), thickness, cv2.LINE_AA)
     return frame
 
 
-def build_counts_html(unique_ids: dict, frame_idx: int, total: int,
-                      elapsed: float) -> str:
-    """Side-panel live counts as HTML cards."""
+# ---------------------------------------------------------------------------
+# HTML side panel
+# ---------------------------------------------------------------------------
+def build_counts_html(truck_total: int, truck_in: int, truck_out: int,
+                      frame_idx: int, total: int, elapsed: float) -> str:
     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>'
+
+    hero = (
+        '<div style="text-align:center;padding:22px 14px;margin-bottom:14px;'
+        'background:linear-gradient(135deg,#f5318b,#c2185b);color:white;'
+        'border-radius:14px;box-shadow:0 4px 12px rgba(245,49,139,0.35);">'
+        '<div style="font-size:32px;line-height:1;">🚛</div>'
+        '<div style="font-size:11px;opacity:0.9;letter-spacing:1.5px;'
+        'margin-top:6px;">TRUCKS CROSSED</div>'
+        f'<div style="font-size:54px;font-weight:800;line-height:1.0;'
+        f'margin-top:4px;">{truck_total}</div>'
+        '<div style="display:flex;justify-content:center;gap:14px;'
+        'margin-top:10px;font-size:12px;opacity:0.95;">'
+        f'<span>↓ {truck_in} down</span>'
+        f'<span>↑ {truck_out} up</span>'
+        '</div></div>'
     )
-    return progress + "".join(cards)
 
+    progress = (
+        '<div style="margin:8px 0 4px 0;">'
+        '<div style="display:flex;justify-content:space-between;font-size:11px;'
+        'color:#6b7280;margin-bottom:4px;">'
+        f'<span>frame {frame_idx} / {total}</span>'
+        f'<span>{pct:.1f}% · {elapsed:.1f}s</span>'
+        '</div>'
+        '<div style="height:6px;background:#e5e7eb;border-radius:3px;overflow:hidden;">'
+        f'<div style="height:100%;width:{pct}%;background:#6366f1;transition:width 0.2s;"></div>'
+        '</div></div>'
+    )
 
-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)
+    return hero + progress
 
 
-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]]
+def build_summary_md(truck_total: int, truck_in: int, truck_out: int) -> str:
+    if truck_total == 0:
+        return ("### ℹ️ No trucks crossed the center line.\n"
+                "Try a lower confidence threshold or a smaller frame stride.")
+    return (f"### 🚛 Total trucks crossed: **{truck_total}**\n\n"
+            f"- ↓ Going down: {truck_in}\n"
+            f"- ↑ Going up: {truck_out}")
 
 
 # ---------------------------------------------------------------------------
 # 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:
         yield (None,
                '<div style="padding:12px;color:#b91c1c;">⚠️ Please upload a video first.</div>',
                "Submit a video to start.",
-               None,
-               [])
+               None)
         return
 
     video_info = sv.VideoInfo.from_video_path(video_path)
+    width, height = video_info.width, video_info.height
+
+    # ---- Line zone fixed at vertical center ----
+    line_y = height // 2
+    line_zone = sv.LineZone(
+        start=sv.Point(0, line_y),
+        end=sv.Point(width, line_y),
+    )
+
+    # ---- Annotators (sized to frame) ----
+    scale = max(0.5, width / 1280)
+    box_ann = sv.BoxAnnotator(thickness=max(2, int(2 * scale)))
+    label_ann = sv.LabelAnnotator(
+        text_scale=0.5 * scale,
+        text_thickness=max(1, int(1 * scale)),
+        text_padding=4,
+    )
+    trace_ann = sv.TraceAnnotator(thickness=max(2, int(2 * scale)), trace_length=40)
+    # We draw the line ourselves so we control the label completely
+    line_ann = sv.LineZoneAnnotator(
+        thickness=max(2, int(3 * scale)),
+        text_thickness=1, text_scale=0.01,    # effectively hide default text
+        display_in_count=False,
+        display_out_count=False,
+    )
+
     frame_gen = sv.get_video_frames_generator(video_path)
     tracker = sv.ByteTrack(frame_rate=int(video_info.fps))
 
-    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
+    last_rgb = None
 
     yield (None,
-           build_counts_html(unique_ids, 0, total, 0.0),
+           build_counts_html(0, 0, 0, 0, total, 0.0),
            "### 🎬 Starting analysis…",
-           None,
-           [])
-
-    last_rgb = None
+           None)
 
     with sv.VideoSink(target_path=out_path, video_info=video_info) as sink:
         for i, frame in enumerate(frame_gen):
@@ -223,16 +224,16 @@ def process_video(video_path, confidence, frame_stride):
                 detections = tracker.update_with_detections(detections)
                 last_detections = detections
 
-                for cid, tid in zip(detections.class_id, detections.tracker_id):
-                    if tid is None:
-                        continue
-                    name = TARGET_CLASSES.get(int(cid))
-                    if name:
-                        unique_ids[name].add(int(tid))
+                # ---- Only trucks feed the line zone ----
+                if len(detections) > 0:
+                    truck_mask = detections.class_id == TRUCK_CLASS_ID
+                    truck_detections = detections[truck_mask]
+                    if len(truck_detections) > 0:
+                        line_zone.trigger(truck_detections)
             else:
                 detections = last_detections
 
-            # ---- Annotate ----
+            # ---- Annotate everything detected (visual richness) ----
             if len(detections) > 0:
                 tids = (detections.tracker_id
                         if detections.tracker_id is not None
@@ -245,34 +246,40 @@ def process_video(video_path, confidence, frame_stride):
                     tid_str = f"#{tid} " if tid is not None else ""
                     labels.append(f"{tid_str}{display} {conf:.2f}")
 
-                frame = BOX_ANNOTATOR.annotate(frame, detections)
-                frame = LABEL_ANNOTATOR.annotate(frame, detections, labels)
+                frame = trace_ann.annotate(scene=frame, detections=detections)
+                frame = box_ann.annotate(scene=frame, detections=detections)
+                frame = label_ann.annotate(scene=frame, detections=detections,
+                                            labels=labels)
 
-            counts_now = {name: len(ids) for name, ids in unique_ids.items()}
-            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)
+            # ---- Draw line + custom truck counter on the line ----
+            frame = line_ann.annotate(frame=frame, line_counter=line_zone)
+            truck_in = int(line_zone.in_count)
+            truck_out = int(line_zone.out_count)
+            truck_total = truck_in + truck_out
+            frame = draw_truck_label_on_line(frame, line_y, truck_total)
 
             sink.write_frame(frame)
             last_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
-            # ---- Yield to UI (throttled to ~5 updates/sec) ----
+            # ---- Yield to UI (throttled ~5/sec) ----
             now = time.time()
             if now - last_yield > 0.20 or i == total - 1:
                 last_yield = now
+                elapsed = time.time() - start_time
                 yield (last_rgb,
-                       build_counts_html(unique_ids, i + 1, total, elapsed),
+                       build_counts_html(truck_total, truck_in, truck_out,
+                                         i + 1, total, elapsed),
                        "### 🔴 Live analysis in progress…",
-                       None,
-                       [])
+                       None)
 
-    # ---- Final yield: include saved video + summary table ----
     elapsed = time.time() - start_time
+    truck_in = int(line_zone.in_count)
+    truck_out = int(line_zone.out_count)
+    truck_total = truck_in + truck_out
     yield (last_rgb,
-           build_counts_html(unique_ids, total, total, elapsed),
-           build_summary_md(unique_ids),
-           out_path,
-           build_table(unique_ids))
+           build_counts_html(truck_total, truck_in, truck_out, total, total, elapsed),
+           build_summary_md(truck_total, truck_in, truck_out),
+           out_path)
 
 
 # ---------------------------------------------------------------------------
@@ -289,16 +296,16 @@ CUSTOM_CSS = """
 footer {visibility: hidden;}
 """
 
-with gr.Blocks(title="RF-DETR Live Object Counter") as demo:
+with gr.Blocks(title="RF-DETR Truck Counter") as demo:
 
     with gr.Row(elem_id="title-row"):
         gr.Markdown(
             """
-            # 🐄 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).
+            # 🚛 RF-DETR Truck Counter
+            Counts trucks crossing a horizontal line at the **center of the video** —
+            in either direction. Powered by
+            [RF-DETR Nano](https://github.com/roboflow/rf-detr) + ByteTrack +
+            `sv.LineZone`.
             """
         )
 
@@ -342,15 +349,14 @@ with gr.Blocks(title="RF-DETR Live Object Counter") as demo:
                     with gr.Column(scale=3):
                         gr.Markdown("### 🔴 Live View")
                         live_frame = gr.Image(
-                            label=None,
                             show_label=False,
                             elem_id="live-frame",
-                            height=420,
+                            height=440,
                         )
-                    with gr.Column(scale=1, min_width=220):
+                    with gr.Column(scale=1, min_width=240):
                         gr.Markdown("### 📊 Live Counts")
                         live_counts = gr.HTML(
-                            value=build_counts_html(defaultdict(set), 0, 0, 0)
+                            value=build_counts_html(0, 0, 0, 0, 0, 0)
                         )
 
     # ---------- Bottom: final results ----------
@@ -361,20 +367,13 @@ with gr.Blocks(title="RF-DETR Live Object Counter") as demo:
                 video_output = gr.Video(label="Download / replay", height=260)
         with gr.Column(scale=1):
             with gr.Group(elem_classes="card"):
-                gr.Markdown("### 📈 Final totals")
+                gr.Markdown("### 📈 Final summary")
                 summary_output = gr.Markdown("Run an analysis to see results.")
-                table_output = gr.Dataframe(
-                    headers=["Class", "Unique count"],
-                    datatype=["str", "number"],
-                    interactive=False,
-                    wrap=True,
-                )
 
     submit_btn.click(
         fn=process_video,
         inputs=[video_input, confidence, frame_stride],
-        outputs=[live_frame, live_counts, summary_output,
-                 video_output, table_output],
+        outputs=[live_frame, live_counts, summary_output, video_output],
     )