Update app.py

app.py

@@ -12,7 +12,6 @@ def get_model_names():
         return model.model.names
     return {}
 
-# Function to count objects in a single frame
 def count_objects(results, cumulative=None):
     names = get_model_names()
     counter = cumulative if cumulative is not None else {}
@@ -21,45 +20,36 @@ def count_objects(results, cumulative=None):
         for cls_id in r.boxes.cls:
             cls_id = int(cls_id)
             label = str(names[cls_id])
-            if label not in counter:
-                counter[label] = 1
-            else:
-                counter[label] += 1
+            counter[label] = counter.get(label, 0) + 1
 
     counter["Total"] = sum(counter.get(k, 0) for k in counter if k != "Total")
     return counter
 
-# Image detection (unchanged)
-def detect_image(img):
-    results = model.predict(img, imgsz=640)
-    annotated = results[0].plot()
-    dashboard = count_objects(results)
-    return annotated, dashboard
+# 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
 
-# Video detection (unchanged)
-def detect_video(video_path):
-    cap = cv2.VideoCapture(video_path)
     ret, frame = cap.read()
     if not ret:
-        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 start_camera(_):
-    return np.zeros((480, 640, 3), dtype=np.uint8), {}
+        return np.zeros((480,640,3), dtype=np.uint8), state.get("cumulative", {}), state
 
-def detect_camera(frame, state):
     results = model.predict(frame, imgsz=640)
     annotated = results[0].plot()
-
-    cumulative = state.get("cumulative", {})
-    cumulative = count_objects(results, cumulative)
+    cumulative = count_objects(results, state.get("cumulative", {}))
     state["cumulative"] = cumulative
 
     return annotated, cumulative, state
@@ -75,7 +65,7 @@ with gr.Blocks(title="Rix Detection") as demo:
             img_out = gr.Image(label="Result Image")
             dashboard1 = gr.JSON(label="Counts")
             btn1 = gr.Button("Detect")
-            btn1.click(fn=detect_image, inputs=img_input, outputs=[img_out, dashboard1])
+            btn1.click(fn=lambda img: detect_image(img), inputs=img_input, outputs=[img_out, dashboard1])
 
         # Tab 2: Video
         with gr.Tab("Video Detection"):
@@ -83,18 +73,26 @@ with gr.Blocks(title="Rix Detection") as demo:
             video_out = gr.Image(label="Demo Frame Result")
             dashboard2 = gr.JSON(label="Counts")
             btn2 = gr.Button("Detect Video")
-            btn2.click(fn=detect_video, inputs=video_input, outputs=[video_out, dashboard2])
+            btn2.click(fn=lambda vid: detect_video(vid), inputs=video_input, outputs=[video_out, dashboard2])
 
         # Tab 3: Live Camera
         with gr.Tab("Live Camera"):
-            cam_input = gr.Image(sources=["webcam"], type="numpy", label="Camera")
-            cam_out = gr.Image(label="Real-time Result")
+            camera_index = gr.Number(value=0, label="Camera Index")
+            cam_out = gr.Image(label="Live Result")
             dashboard3 = gr.JSON(label="Cumulative Counts")
             state = gr.State()
 
-            start_btn = gr.Button("Start Detection")
-            start_btn.click(fn=start_camera, inputs=None, outputs=[cam_out, dashboard3])
-
-            cam_input.stream(fn=detect_camera, inputs=[cam_input, state], outputs=[cam_out, dashboard3, state])
+            start_btn = gr.Button("Start Camera")
+            stop_btn = gr.Button("Stop Camera")
+
+            start_btn.click(fn=init_camera, inputs=camera_index, 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]
+            )
 
 demo.launch()