Update app.py

app.py

@@ -3,7 +3,10 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 
-model = YOLO("rix_reg.pt")
+# ======================================================
+# Load YOLO model
+# ======================================================
+model = YOLO("rix_reg.pt")  # change to your model
 
 def get_model_names():
     if hasattr(model, "names") and model.names is not None:
@@ -12,87 +15,98 @@ def get_model_names():
         return model.model.names
     return {}
 
+# ======================================================
+# Function to count all objects
+# ======================================================
 def count_objects(results, cumulative=None):
     names = get_model_names()
     counter = cumulative if cumulative is not None else {}
-
     for r in results:
         for cls_id in r.boxes.cls:
             cls_id = int(cls_id)
             label = str(names[cls_id])
             counter[label] = counter.get(label, 0) + 1
-
     counter["Total"] = sum(counter.get(k, 0) for k in counter if k != "Total")
     return counter
 
-# Initialize camera state
-def init_camera(camera_index):
-    cap = cv2.VideoCapture(camera_index)
-    return {"cap": cap, "cumulative": {}}
-
-# Release camera
-def stop_camera(state):
-    cap = state.get("cap")
-    if cap is not None:
-        cap.release()
-    return np.zeros((480,640,3), dtype=np.uint8), {}, state
-
-# Detect frame and update cumulative count
-def detect_camera(state):
-    cap = state.get("cap")
-    if cap is None or not cap.isOpened():
-        return np.zeros((480,640,3), dtype=np.uint8), state.get("cumulative", {}), state
-
+# ======================================================
+# Image Detection
+# ======================================================
+def detect_image(img):
+    results = model.predict(img, imgsz=640)
+    annotated = results[0].plot()
+    dashboard = count_objects(results)
+    return annotated, dashboard
+
+# ======================================================
+# Video Detection
+# ======================================================
+def detect_video(video_path):
+    cap = cv2.VideoCapture(video_path)
     ret, frame = cap.read()
     if not ret:
-        return np.zeros((480,640,3), dtype=np.uint8), state.get("cumulative", {}), state
+        return None, {"Error": "Cannot read video"}
+    results = model.predict(frame, imgsz=640)
+    annotated = results[0].plot()
+    dashboard = count_objects(results)
+    cap.release()
+    return annotated, dashboard
 
+# ======================================================
+# Live Camera Detection with cumulative counting
+# ======================================================
+def init_camera():
+    return {}  # reset cumulative counts
+
+def detect_camera(frame, state):
     results = model.predict(frame, imgsz=640)
     annotated = results[0].plot()
-    cumulative = count_objects(results, state.get("cumulative", {}))
-    state["cumulative"] = cumulative
+    cumulative = count_objects(results, state)
+    return annotated, cumulative, cumulative
 
-    return annotated, cumulative, state
+def stop_camera(state):
+    return np.zeros((480,640,3), dtype=np.uint8), {}, state
 
+# ======================================================
+# GRADIO Interface
+# ======================================================
 with gr.Blocks(title="Rix Detection") as demo:
+
     gr.Markdown("## 🛠️ Object Counting Dashboard")
 
     with gr.Tabs():
 
-        # Tab 1: Image
+        # ==================== TAB 1: Image ====================
         with gr.Tab("Image Detection"):
             img_input = gr.Image(type="numpy", label="Upload Image")
             img_out = gr.Image(label="Result Image")
             dashboard1 = gr.JSON(label="Counts")
             btn1 = gr.Button("Detect")
-            btn1.click(fn=lambda img: detect_image(img), inputs=img_input, outputs=[img_out, dashboard1])
+            btn1.click(fn=detect_image, inputs=img_input, outputs=[img_out, dashboard1])
 
-        # Tab 2: Video
+        # ==================== TAB 2: Video ====================
         with gr.Tab("Video Detection"):
             video_input = gr.Video(label="Upload Video")
             video_out = gr.Image(label="Demo Frame Result")
             dashboard2 = gr.JSON(label="Counts")
             btn2 = gr.Button("Detect Video")
-            btn2.click(fn=lambda vid: detect_video(vid), inputs=video_input, outputs=[video_out, dashboard2])
+            btn2.click(fn=detect_video, inputs=video_input, outputs=[video_out, dashboard2])
 
-        # Tab 3: Live Camera
+        # ==================== TAB 3: Live Camera ====================
         with gr.Tab("Live Camera"):
-            camera_index = gr.Number(value=0, label="Camera Index")
+            gr.Markdown("### Select camera and start detection")
+            cam_index = gr.Number(value=0, label="Camera Index (e.g., 0, 1)")
+            cam_input = gr.Image(source="webcam", type="numpy", label="Camera Feed")
             cam_out = gr.Image(label="Live Result")
             dashboard3 = gr.JSON(label="Cumulative Counts")
             state = gr.State()
 
-            start_btn = gr.Button("Start Camera")
-            stop_btn = gr.Button("Stop Camera")
+            start_btn = gr.Button("Start Detection")
+            stop_btn = gr.Button("Stop Detection")
 
-            start_btn.click(fn=init_camera, inputs=camera_index, outputs=state)
+            start_btn.click(fn=init_camera, inputs=None, outputs=state)
             stop_btn.click(fn=stop_camera, inputs=state, outputs=[cam_out, dashboard3, state])
-            
-            cam_out.update(
-                every=0.1,
-                fn=detect_camera,
-                inputs=state,
-                outputs=[cam_out, dashboard3, state]
-            )
+
+            cam_input.stream(fn=detect_camera, inputs=[cam_input, state], outputs=[cam_out, dashboard3, state])
 
 demo.launch()