Update app.py

app.py

@@ -1,55 +1,27 @@
 import gradio as gr
 import cv2
+import tempfile
+import time
 from ultralytics import YOLO
 from deep_sort_realtime.deepsort_tracker import DeepSort
-import tempfile
-import numpy as np
-import json
 
 # ======================
-# MODELS
+# LOAD MODEL
 # ======================
-model = YOLO("best.pt")
-tracker = DeepSort(max_age=30)
-
+model = YOLO("best.pt")   # يفضل yolov8n
 class_names = model.names
 
-# نخزن history لكل ماكينة
-track_memory = {}
-track_stats = {}
-
-# ======================
-# ACTIVITY CLASSIFIER
-# ======================
-def classify_activity(dx, dy, speed):
-
-    if speed < 1.2:
-        return "WAITING"
-
-    if abs(dy) > abs(dx) and dy > 2:
-        return "DIGGING"
-
-    if abs(dx) > 3:
-        return "SWINGING"
-
-    if speed > 7:
-        return "DUMPING"
-
-    return "ACTIVE"
-
+tracker = DeepSort(max_age=15)
 
 # ======================
-# MAIN FUNCTION
+# PROCESS VIDEO
 # ======================
 def process_video(video_file, selected_classes):
 
     cap = cv2.VideoCapture(video_file)
 
-    width = int(cap.get(3))
-    height = int(cap.get(4))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-
-    total_frames = 0
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    w, h = 640, 360
 
     temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
     out_path = temp_output.name
@@ -58,156 +30,127 @@ def process_video(video_file, selected_classes):
         out_path,
         cv2.VideoWriter_fourcc(*"mp4v"),
         fps,
-        (width, height)
+        (w, h)
     )
 
-    # ======================
-    # VIDEO LOOP
-    # ======================
+    frame_id = 0
+
+    # counters
+    total_objects = 0
+    active_frames = 0
+    inactive_frames = 0
+
+    start_time = time.time()
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
 
-        total_frames += 1
+        frame_id += 1
+
+        # ⚡ SKIP FRAMES (speed boost)
+        if frame_id % 2 != 0:
+            continue
+
+        frame = cv2.resize(frame, (w, h))
 
         results = model(frame)[0]
 
         detections = []
 
+        # map selected classes
+        selected_ids = [k for k, v in class_names.items() if v in selected_classes]
+
         for box in results.boxes:
             cls_id = int(box.cls[0])
-            cls_name = class_names[cls_id].lower()
 
-            if cls_name not in [c.lower() for c in selected_classes]:
+            if cls_id not in selected_ids:
                 continue
 
             x1, y1, x2, y2 = box.xyxy[0].tolist()
             conf = float(box.conf[0])
 
-            detections.append(([x1, y1, x2-x1, y2-y1], conf, cls_name))
+            detections.append(([x1, y1, x2-x1, y2-y1], conf, cls_id))
 
         tracks = tracker.update_tracks(detections, frame=frame)
 
-        # ======================
-        # TRACK PROCESSING
-        # ======================
+        active_in_frame = 0
+
         for t in tracks:
             if not t.is_confirmed():
                 continue
 
+            x1, y1, x2, y2 = map(int, t.to_ltrb())
             track_id = t.track_id
-            x1, y1, x2, y2 = t.to_ltrb()
 
-            cx = (x1 + x2) / 2
-            cy = (y1 + y2) / 2
-
-            timestamp_sec = total_frames / fps
-            timestamp = f"00:00:{timestamp_sec:06.3f}"
-
-            # init
-            if track_id not in track_memory:
-                track_memory[track_id] = (cx, cy)
-                track_stats[track_id] = {
-                    "active": 0,
-                    "idle": 0,
-                    "activity_log": []
-                }
-
-            px, py = track_memory[track_id]
-
-            dx = cx - px
-            dy = cy - py
-            speed = np.sqrt(dx**2 + dy**2)
-
-            activity = classify_activity(dx, dy, speed)
-
-            # ACTIVE vs IDLE
-            if activity == "WAITING":
-                state = "IDLE"
-                track_stats[track_id]["idle"] += 1
-            else:
-                state = "ACTIVE"
-                track_stats[track_id]["active"] += 1
-
-            track_stats[track_id]["activity_log"].append(activity)
-
-            track_memory[track_id] = (cx, cy)
-
-            # ======================
-            # DRAW
-            # ======================
-            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,255,0), 2)
+            active_in_frame += 1
+            total_objects += 1
 
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
             cv2.putText(frame,
-                        f"ID:{track_id} {activity}",
-                        (int(x1), int(y1)-10),
+                        f"ID {track_id}",
+                        (x1, y1 - 10),
                         cv2.FONT_HERSHEY_SIMPLEX,
-                        0.6,
+                        0.5,
                         (0,255,0),
                         2)
 
+        # Active / Inactive logic
+        if active_in_frame > 0:
+            active_frames += 1
+        else:
+            inactive_frames += 1
+
+        cv2.putText(frame,
+                    f"Objects: {active_in_frame}",
+                    (10, 30),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.8,
+                    (255, 0, 0),
+                    2)
+
         out.write(frame)
 
     cap.release()
     out.release()
 
-    # ======================
-    # BUILD JSON OUTPUT
-    # ======================
-    results_json = []
-
-    for track_id, data in track_stats.items():
+    total_time = time.time() - start_time
 
-        total = data["active"] + data["idle"]
+    summary = f"""
+    🚜 Processing Done!
 
-        active_sec = data["active"] / fps
-        idle_sec = data["idle"] / fps
+    🔢 Total tracked detections: {total_objects}
+    🟢 Active frames: {active_frames}
+    🔴 Inactive frames: {inactive_frames}
+    ⏱ Processing time: {total_time:.2f} sec
+    """
 
-        utilization = (active_sec / (active_sec + idle_sec + 1e-6)) * 100
+    return out_path, summary
 
-        main_activity = max(set(data["activity_log"]), key=data["activity_log"].count)
 
-        equipment_json = {
-            "frame_id": total_frames,
-            "equipment_id": f"EX-{track_id:03d}",
-            "equipment_class": "excavator",
-            "timestamp": f"00:00:{total_frames / fps:06.3f}",
-            "utilization": {
-                "current_state": "ACTIVE" if active_sec > idle_sec else "IDLE",
-                "current_activity": main_activity,
-                "motion_source": "arm_and_body_motion"
-            },
-            "time_analytics": {
-                "total_tracked_seconds": round(total / fps, 2),
-                "total_active_seconds": round(active_sec, 2),
-                "total_idle_seconds": round(idle_sec, 2),
-                "utilization_percent": round(utilization, 2)
-            }
-        }
+# ======================
+# GRADIO UI
+# ======================
+with gr.Blocks() as demo:
 
-        results_json.append(equipment_json)
+    gr.Markdown("# 🚜 Excavator Monitoring AI (YOLO + Tracking)")
 
-    json_output = json.dumps(results_json, indent=2)
+    video_input = gr.Video()
+    class_selector = gr.CheckboxGroup(
+        choices=list(class_names.values()),
+        label="Select Classes"
+    )
 
-    return out_path, json_output, out_path
+    btn = gr.Button("Process Video 🚀")
 
+    output_video = gr.Video()
+    output_text = gr.Textbox()
 
-# ======================
-# GRADIO UI
-# ======================
-demo = gr.Interface(
-    fn=process_video,
-    inputs=[
-        gr.Video(label="Upload Video"),
-        gr.CheckboxGroup(choices=list(class_names.values()), label="Select Classes")
-    ],
-    outputs=[
-        gr.Video(label="Processed Video"),
-        gr.Textbox(label="JSON Output"),
-        gr.File(label="Download Video")
-    ],
-    title="Smart Excavator Analytics"
-)
+    btn.click(
+        fn=process_video,
+        inputs=[video_input, class_selector],
+        outputs=[output_video, output_text]
+    )
 
 demo.launch()