Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 # ======================================================
 # Load YOLO model
 # ======================================================
-model = YOLO("rix_reg.pt")  # change to your model
+model = YOLO("rix_reg.pt")   # change to your model
 
 def get_model_names():
     if hasattr(model, "names") and model.names is not None:
@@ -16,59 +16,71 @@ def get_model_names():
     return {}
 
 # ======================================================
-# Function to count all objects
+# Function to count all objects in the model
 # ======================================================
-def count_objects(results, cumulative=None):
+def count_objects(results):
     names = get_model_names()
-    counter = cumulative if cumulative is not None else {}
+    counter = {}
+
     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")
+
+            # 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)
     return counter
 
+
 # ======================================================
-# Image Detection
+# Tab 1 - Image processing
 # ======================================================
 def detect_image(img):
     results = model.predict(img, imgsz=640)
     annotated = results[0].plot()
+
     dashboard = count_objects(results)
     return annotated, dashboard
 
+
 # ======================================================
-# Video Detection
+# Tab 2 - Video processing
 # ======================================================
 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
+# Tab 3 - Live camera
 # ======================================================
-def init_camera():
-    return {}  # reset cumulative counts
-
-def detect_camera(frame, state):
+def detect_camera(frame):
     results = model.predict(frame, imgsz=640)
     annotated = results[0].plot()
-    cumulative = count_objects(results, state)
-    return annotated, cumulative, cumulative
 
-def stop_camera(state):
-    return np.zeros((480,640,3), dtype=np.uint8), {}, state
+    dashboard = count_objects(results)
+    return annotated, dashboard
+
 
 # ======================================================
-# GRADIO Interface
+# GRADIO interface
 # ======================================================
 with gr.Blocks(title="Rix Detection") as demo:
 
@@ -76,37 +88,44 @@ with gr.Blocks(title="Rix Detection") as demo:
 
     with gr.Tabs():
 
-        # ==================== TAB 1: Image ====================
+        # ==================== TAB 1 ====================
         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])
 
-        # ==================== TAB 2: Video ====================
+            btn1.click(
+                fn=detect_image,
+                inputs=img_input,
+                outputs=[img_out, dashboard1]
+            )
+
+        # ==================== TAB 2 ====================
         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])
-
-        # ==================== TAB 3: Live Camera ====================
-        with gr.Tab("Live Camera"):
-            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 Detection")
-            stop_btn = gr.Button("Stop Detection")
+            btn2 = gr.Button("Detect Video")
 
-            start_btn.click(fn=init_camera, inputs=None, outputs=state)
-            stop_btn.click(fn=stop_camera, inputs=state, outputs=[cam_out, dashboard3, state])
+            btn2.click(
+                fn=detect_video,
+                inputs=video_input,
+                outputs=[video_out, dashboard2]
+            )
 
-            cam_input.stream(fn=detect_camera, inputs=[cam_input, state], outputs=[cam_out, dashboard3, state])
+        # ==================== TAB 3 ====================
+        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")
+
+            cam_input.stream(
+                fn=detect_camera,
+                inputs=cam_input,
+                outputs=[cam_out, dashboard3]
+            )
 
 demo.launch()