Update app.py

app.py

@@ -1,41 +1,69 @@
 import gradio as gr
 import cv2
 from ultralytics import YOLO
+from deep_sort_realtime.deepsort_tracker import DeepSort
 import tempfile
+import numpy as np
+import json
 
-# تحميل الموديل
+# ======================
+# MODELS
+# ======================
 model = YOLO("best.pt")
+tracker = DeepSort(max_age=30)
 
-# أسماء الكلاسات
 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"
+
+
+# ======================
+# MAIN FUNCTION
+# ======================
 def process_video(video_file, selected_classes):
-    cap = cv2.VideoCapture(video_file)
 
-    if not cap.isOpened():
-        return None, "Error opening video", "", 0, 0
+    cap = cv2.VideoCapture(video_file)
 
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    width = int(cap.get(3))
+    height = int(cap.get(4))
     fps = cap.get(cv2.CAP_PROP_FPS)
 
+    total_frames = 0
+
     temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    output_path = temp_output.name
+    out_path = temp_output.name
 
     out = cv2.VideoWriter(
-        output_path,
-        cv2.VideoWriter_fourcc(*'mp4v'),
+        out_path,
+        cv2.VideoWriter_fourcc(*"mp4v"),
         fps,
         (width, height)
     )
 
-    total_objects = 0
-    active_frames = 0
-    total_frames = 0
-
-    # تحويل أسماء الكلاسات لأرقام
-    selected_ids = [k for k, v in class_names.items() if v in selected_classes]
-
+    # ======================
+    # VIDEO LOOP
+    # ======================
     while True:
         ret, frame = cap.read()
         if not ret:
@@ -43,45 +71,131 @@ def process_video(video_file, selected_classes):
 
         total_frames += 1
 
-        results = model(frame)
+        results = model(frame)[0]
+
+        detections = []
+
+        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]:
+                continue
 
-        boxes = results[0].boxes
-        count_in_frame = 0
+            x1, y1, x2, y2 = box.xyxy[0].tolist()
+            conf = float(box.conf[0])
 
-        if boxes is not None:
-            for box in boxes:
-                cls_id = int(box.cls[0])
+            detections.append(([x1, y1, x2-x1, y2-y1], conf, cls_name))
 
-                if cls_id in selected_ids:
-                    count_in_frame += 1
+        tracks = tracker.update_tracks(detections, frame=frame)
 
-        total_objects += count_in_frame
+        # ======================
+        # TRACK PROCESSING
+        # ======================
+        for t in tracks:
+            if not t.is_confirmed():
+                continue
 
-        # active لو فيه object
-        if count_in_frame > 0:
-            active_frames += 1
+            track_id = t.track_id
+            x1, y1, x2, y2 = t.to_ltrb()
 
-        # رسم النتائج (فلترة)
-        annotated_frame = results[0].plot()
+            cx = (x1 + x2) / 2
+            cy = (y1 + y2) / 2
 
-        out.write(annotated_frame)
+            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)
+
+            cv2.putText(frame,
+                        f"ID:{track_id} {activity}",
+                        (int(x1), int(y1)-10),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.6,
+                        (0,255,0),
+                        2)
+
+        out.write(frame)
 
     cap.release()
     out.release()
 
-    # حساب الزمن
-    active_time = active_frames / fps
-    inactive_time = (total_frames - active_frames) / fps
+    # ======================
+    # BUILD JSON OUTPUT
+    # ======================
+    results_json = []
+
+    for track_id, data in track_stats.items():
+
+        total = data["active"] + data["idle"]
+
+        active_sec = data["active"] / fps
+        idle_sec = data["idle"] / fps
+
+        utilization = (active_sec / (active_sec + idle_sec + 1e-6)) * 100
+
+        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)
+            }
+        }
+
+        results_json.append(equipment_json)
 
-    summary = f"""
-    📊 Total Objects Detected: {total_objects}
-    ⏱️ Active Time: {active_time:.2f} sec
-    💤 Inactive Time: {inactive_time:.2f} sec
-    """
+    json_output = json.dumps(results_json, indent=2)
 
-    return output_path, summary, output_path, active_time, inactive_time
+    return out_path, json_output, out_path
 
 
+# ======================
+# GRADIO UI
+# ======================
 demo = gr.Interface(
     fn=process_video,
     inputs=[
@@ -90,13 +204,11 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Video(label="Processed Video"),
-        gr.Textbox(label="Summary"),
-        gr.File(label="Download Video"),
-        gr.Number(label="Active Time (sec)"),
-        gr.Number(label="Inactive Time (sec)")
+        gr.Textbox(label="JSON Output"),
+        gr.File(label="Download Video")
     ],
-    title="YOLO Smart Video Analyzer 🎯🔥",
-    description="ارفع فيديو + اختار الكلاسات، وهيطلعلك تحليل كامل"
+    title="Smart Excavator Analytics"
+    description="YOLO + Tracking + Activity + Utilization JSON"
 )
 
 demo.launch()