Update app.py

app.py

@@ -3,10 +3,7 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 
-# ======================================================
-# Load YOLO model
-# ======================================================
-model = YOLO("rix_reg.pt")   # change to your model
+model = YOLO("rix_reg.pt")
 
 def get_model_names():
     if hasattr(model, "names") and model.names is not None:
@@ -15,117 +12,89 @@ def get_model_names():
         return model.model.names
     return {}
 
-# ======================================================
-# Function to count all objects in the model
-# ======================================================
-def count_objects(results):
+# Function to count objects in a single frame
+def count_objects(results, cumulative=None):
     names = get_model_names()
-    counter = {}
+    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])
-
-            # increment count
             if label not in counter:
                 counter[label] = 1
             else:
                 counter[label] += 1
 
-    counter["Total"] = sum(counter.get(k, 0) for k in counter)
+    counter["Total"] = sum(counter.get(k, 0) for k in counter if k != "Total")
     return counter
 
-
-# ======================================================
-# Tab 1 - Image processing
-# ======================================================
+# Image detection (unchanged)
 def detect_image(img):
     results = model.predict(img, imgsz=640)
     annotated = results[0].plot()
-
     dashboard = count_objects(results)
     return annotated, dashboard
 
-
-# ======================================================
-# Tab 2 - Video processing
-# ======================================================
+# 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"}
-
-    # demo first frame
     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), {}
 
-# ======================================================
-# Tab 3 - Live camera
-# ======================================================
-def detect_camera(frame):
+def detect_camera(frame, state):
     results = model.predict(frame, imgsz=640)
     annotated = results[0].plot()
 
-    dashboard = count_objects(results)
-    return annotated, dashboard
+    cumulative = state.get("cumulative", {})
+    cumulative = count_objects(results, cumulative)
+    state["cumulative"] = cumulative
 
+    return annotated, cumulative, state
 
-# ======================================================
-# GRADIO interface
-# ======================================================
 with gr.Blocks(title="Rix Detection") as demo:
-
     gr.Markdown("## 🛠️ Object Counting Dashboard")
 
     with gr.Tabs():
 
-        # ==================== TAB 1 ====================
+        # 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=detect_image, inputs=img_input, outputs=[img_out, dashboard1])
 
-            btn1.click(
-                fn=detect_image,
-                inputs=img_input,
-                outputs=[img_out, dashboard1]
-            )
-
-        # ==================== TAB 2 ====================
+        # 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=detect_video, inputs=video_input, outputs=[video_out, dashboard2])
 
-            btn2.click(
-                fn=detect_video,
-                inputs=video_input,
-                outputs=[video_out, dashboard2]
-            )
-
-        # ==================== TAB 3 ====================
+        # 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")
-            dashboard3 = gr.JSON(label="Counts")
+            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,
-                outputs=[cam_out, dashboard3]
-            )
+            cam_input.stream(fn=detect_camera, inputs=[cam_input, state], outputs=[cam_out, dashboard3, state])
 
 demo.launch()