Update app.py

app.py

@@ -1,7 +1,8 @@
 """
-RF-DETR Object Counter — Gradio app for Hugging Face Spaces.
-Counts people, bicycles, cars, trucks, and animals in video using
-RF-DETR Medium + ByteTrack (so each object is counted only once).
+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.
 """
 
 import os
@@ -12,47 +13,47 @@ import cv2
 import gradio as gr
 import numpy as np
 import supervision as sv
+import torch
 from rfdetr import RFDETRMedium
-from rfdetr.assets.coco_classes import COCO_CLASSES
 
 # ---------------------------------------------------------------------------
-# Target classes (COCO indices) — exactly what the user asked for
+# 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.
 TARGET_CLASSES = {
     0:  "person",
-    1:  "bicycle",
     2:  "car",
     7:  "truck",
-    # animals
-    14: "bird",
-    15: "cat",
     16: "dog",
-    17: "horse",
-    18: "sheep",
+    17: "horse",   # also catches donkeys
+    18: "sheep",   # also catches goats
     19: "cow",
-    20: "elephant",
-    21: "bear",
-    22: "zebra",
-    23: "giraffe",
 }
 TARGET_IDS = list(TARGET_CLASSES.keys())
 
-# Per-class colour palette (BGR) for the live overlay
+# Friendly UI labels
+DISPLAY_NAMES = {
+    "person": "person",
+    "car":    "car",
+    "truck":  "truck",
+    "dog":    "dog",
+    "horse":  "horse / donkey",
+    "sheep":  "sheep / goat",
+    "cow":    "cow",
+}
+
+# Per-class colours for the live overlay panel (BGR)
 CLASS_COLORS = {
-    "person":   (66, 135, 245),
-    "bicycle":  (245, 173, 66),
-    "car":      (66, 245, 167),
-    "truck":    (245, 66, 161),
-    "bird":     (245, 230, 66),
-    "cat":      (200, 120, 245),
-    "dog":      (120, 245, 200),
-    "horse":    (245, 120, 120),
-    "sheep":    (220, 220, 220),
-    "cow":      (140, 90, 60),
-    "elephant": (160, 160, 200),
-    "bear":     (90, 60, 30),
-    "zebra":    (40, 40, 40),
-    "giraffe":  (220, 180, 90),
+    "person":  (66, 135, 245),
+    "car":     (66, 245, 167),
+    "truck":   (245, 66, 161),
+    "dog":     (120, 245, 200),
+    "horse":   (245, 120, 120),
+    "sheep":   (220, 220, 220),
+    "cow":     (140, 90, 60),
 }
 
 # Example video lives next to app.py
@@ -60,15 +61,26 @@ APP_DIR = os.path.dirname(os.path.abspath(__file__))
 EXAMPLE_VIDEO = os.path.join(APP_DIR, "example.mp4")
 
 # ---------------------------------------------------------------------------
-# Load model once at startup
+# Load model — pinned to CPU
 # ---------------------------------------------------------------------------
-print("Loading RF-DETR Medium…")
-MODEL = RFDETRMedium()
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Loading RF-DETR Medium on {DEVICE}…")
+MODEL = RFDETRMedium(device=DEVICE)
+
+# optimize_for_inference is GPU-only (TensorRT-style ops). Skip on CPU.
+if DEVICE == "cuda":
+    try:
+        MODEL.optimize_for_inference()
+        print("Optimized for GPU inference.")
+    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:
-    MODEL.optimize_for_inference()  # speeds up subsequent predicts
-    print("Model optimized for inference.")
-except Exception as e:
-    print(f"(Optimization skipped: {e})")
+    torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1))
+except Exception:
+    pass
+
 print("Model ready.")
 
 # Annotators
@@ -77,12 +89,12 @@ LABEL_ANNOTATOR = sv.LabelAnnotator(text_scale=0.45, text_thickness=1, text_padd
 
 
 def draw_counter_panel(frame: np.ndarray, counts: dict) -> np.ndarray:
-    """Translucent counter panel in the top-left corner."""
+    """Translucent live-count panel in the top-left corner of the frame."""
     active = [(name, n) for name, n in counts.items() if n > 0]
     if not active:
         active = [("No targets yet", 0)]
 
-    panel_w = 230
+    panel_w = 280
     panel_h = 40 + 22 * len(active)
     overlay = frame.copy()
     cv2.rectangle(overlay, (12, 12), (12 + panel_w, 12 + panel_h), (20, 20, 20), -1)
@@ -95,13 +107,15 @@ def draw_counter_panel(frame: np.ndarray, counts: dict) -> np.ndarray:
     for name, n in active:
         color = CLASS_COLORS.get(name, (200, 200, 200))
         cv2.circle(frame, (28, y - 5), 5, color, -1)
-        cv2.putText(frame, f"{name}: {n}", (44, y),
+        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 process_video(video_path, confidence, frame_stride,
+                  progress=gr.Progress(track_tqdm=True)):
     if video_path is None:
         return None, "⚠️ Please upload a video first.", []
 
@@ -109,15 +123,19 @@ def process_video(video_path, confidence, frame_stride, progress=gr.Progress(tra
     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)            # class_name -> {tracker_id, ...}
     last_detections = sv.Detections.empty()
 
     out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
 
     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")):
-            # Detect every Nth frame; reuse previous detections in-between to keep video smooth
+        for i, frame in enumerate(progress.tqdm(
+                frame_gen,
+                total=video_info.total_frames,
+                desc="Analyzing video")):
+
+            # Detect every Nth frame; reuse previous detections in-between
+            # so the output video stays smooth even with high stride.
             if i % frame_stride == 0:
                 rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 detections = MODEL.predict(rgb, threshold=confidence)
@@ -130,7 +148,7 @@ def process_video(video_path, confidence, frame_stride, progress=gr.Progress(tra
                 detections = tracker.update_with_detections(detections)
                 last_detections = detections
 
-                # Register unique IDs per class
+                # Register unique tracker IDs per class
                 for cid, tid in zip(detections.class_id, detections.tracker_id):
                     if tid is None:
                         continue
@@ -140,17 +158,18 @@ def process_video(video_path, confidence, frame_stride, progress=gr.Progress(tra
             else:
                 detections = last_detections
 
-            # Annotate
+            # Annotate frame
             if len(detections) > 0:
-                labels = [
-                    f"#{tid} {TARGET_CLASSES.get(int(cid), 'obj')} {conf:.2f}"
-                    for cid, tid, conf in zip(
-                        detections.class_id,
-                        detections.tracker_id if detections.tracker_id is not None
-                            else [None] * len(detections),
-                        detections.confidence,
-                    )
-                ]
+                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):
+                    name = TARGET_CLASSES.get(int(cid), "obj")
+                    display = DISPLAY_NAMES.get(name, name)
+                    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)
 
@@ -158,21 +177,26 @@ def process_video(video_path, confidence, frame_stride, progress=gr.Progress(tra
             frame = draw_counter_panel(frame, counts_now)
             sink.write_frame(frame)
 
-    # Build summary outputs
+    # ---------- Build summary outputs ----------
     total = sum(len(ids) for ids in unique_ids.values())
     if total == 0:
-        summary_md = "### ℹ️ No target objects detected.\nTry lowering the confidence threshold."
+        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:
-                lines.append(f"- **{name.capitalize()}** — {n}")
+                display = DISPLAY_NAMES.get(name, name).capitalize()
+                lines.append(f"- **{display}** — {n}")
         summary_md = "\n".join(lines)
 
-    table = [[name.capitalize(), len(unique_ids.get(name, set()))]
-             for name in TARGET_CLASSES.values()
-             if len(unique_ids.get(name, set())) > 0]
+    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]]
 
@@ -192,15 +216,14 @@ CUSTOM_CSS = """
 footer {visibility: hidden;}
 """
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
-               css=CUSTOM_CSS, title="RF-DETR Object Counter") as demo:
+with gr.Blocks(title="RF-DETR Object Counter") as demo:
 
     with gr.Row(elem_id="title-row"):
         gr.Markdown(
             """
-            # 🚦 RF-DETR Object Counter
-            Count **people, bicycles, cars, trucks, and animals** in any video.
-            Powered by [RF-DETR Medium](https://github.com/roboflow/rf-detr) (Roboflow, ICLR 2026) and ByteTrack —
+            # 🐄 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.
             """
         )
@@ -218,14 +241,14 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate")
 
                 with gr.Accordion("⚙️ Advanced settings", open=False):
                     confidence = gr.Slider(
-                        minimum=0.1, maximum=0.9, value=0.5, step=0.05,
+                        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=10, value=2, step=1,
-                        label="Frame stride",
-                        info="Process every Nth frame. Higher = faster, slightly less accurate.",
+                        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.",
                     )
 
                 submit_btn = gr.Button("🔍 Count Objects", variant="primary", size="lg")
@@ -253,12 +276,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate")
 
     gr.Markdown(
         """
-        ---
-        **Detected categories:** person · bicycle · car · truck · bird · cat · dog · horse ·
-        sheep · cow · elephant · bear · zebra · giraffe
-        
-        **Tip:** the first run loads the model (≈45–90 s for Medium). Subsequent runs are much faster.
-        Use *Frame stride* if processing is slow on CPU.
+        **Detected categories:** person · car · truck · dog · horse / donkey · sheep / goat · cow
         """
     )
 
@@ -270,4 +288,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate")
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=8).launch()
+    demo.queue(max_size=4).launch(
+        theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
+        css=CUSTOM_CSS,
+    )