Update app.py

app.py

@@ -8,11 +8,19 @@ from deep_sort_realtime.deepsort_tracker import DeepSort
 
 model = YOLO("best.pt")
 class_names = model.names
-tracker = DeepSort(max_age=15)
+tracker = DeepSort(max_age=30)
 
 
+# ===================================================
+# Articulated Motion — نقسم الـ box لـ zones
+# ===================================================
 def analyze_articulated_motion(frame, prev_frame, bbox):
-  
+    """
+    نقسم الـ bounding box لجزئين:
+    - upper zone = الـ arm / bucket  (فوق)
+    - lower zone = الـ tracks / body (تحت)
+    لو الجزء العلوي بيتحرك → ACTIVE حتى لو الجسم واقف
+    """
     x1, y1, x2, y2 = map(int, bbox)
 
     x1 = max(0, x1)
@@ -21,42 +29,69 @@ def analyze_articulated_motion(frame, prev_frame, bbox):
     y2 = min(frame.shape[0], y2)
 
     h = y2 - y1
-    if h < 10:
+    w = x2 - x1
+
+    if h < 20 or w < 20:
         return False, "none"
 
     mid_y = y1 + h // 2
 
+    # الجزء العلوي (الـ arm)
     upper_curr = cv2.cvtColor(frame[y1:mid_y, x1:x2], cv2.COLOR_BGR2GRAY)
     upper_prev = cv2.cvtColor(prev_frame[y1:mid_y, x1:x2], cv2.COLOR_BGR2GRAY)
 
+    # الجزء السفلي (الـ tracks)
     lower_curr = cv2.cvtColor(frame[mid_y:y2, x1:x2], cv2.COLOR_BGR2GRAY)
     lower_prev = cv2.cvtColor(prev_frame[mid_y:y2, x1:x2], cv2.COLOR_BGR2GRAY)
 
     upper_motion = np.mean(cv2.absdiff(upper_curr, upper_prev))
     lower_motion = np.mean(cv2.absdiff(lower_curr, lower_prev))
 
-    if upper_motion > 8 and lower_motion < 5:
-        return True, "arm_only"       
-    elif lower_motion > 8:
-        return True, "full_body"      
+    # threshold منخفض عشان يكتشف حركة الـ arm الصغيرة
+    if upper_motion > 3:
+        if lower_motion < 3:
+            return True, "arm_only"    # الـ arm بس بيتحرك
+        else:
+            return True, "full_body"   # الجسم كله بيتحرك
+    elif lower_motion > 3:
+        return True, "full_body"
     else:
-        return False, "none"          
+        return False, "none"           # واقف تماماً
+
 
+# ===================================================
+# Activity Classification
+# ===================================================
+def get_activity(history, is_active, motion_source):
+    # لو مفيش motion خالص → WAITING
+    if not is_active and motion_source == "none":
+        return "WAITING"
 
-def get_activity(history):
+    # لو التاريخ قصير بس في motion → DIGGING
     if len(history) < 5:
         return "DIGGING"
 
     dx = history[-1][0] - history[0][0]
     dy = history[-1][1] - history[0][1]
 
-    if abs(dx) + abs(dy) < 10:
+    # لو الـ arm بس بيتحرك (جسم واقف)
+    if motion_source == "arm_only":
+        if dy > 5:
+            return "DIGGING"
+        elif dy < -5:
+            return "DUMPING"
+        else:
+            return "DIGGING"
+
+    # لو الجسم بيتحرك
+    total_move = abs(dx) + abs(dy)
+    if total_move < 8:
         return "WAITING"
-    if dy > 15:
+    if dy > 10:
         return "DIGGING"
     if abs(dx) > abs(dy):
         return "SWINGING/LOADING"
-    if dy < -15:
+    if dy < -10:
         return "DUMPING"
 
     return "WAITING"
@@ -65,9 +100,9 @@ def get_activity(history):
 def process_video(video_file, selected_classes):
 
     cap = cv2.VideoCapture(video_file)
-    fps = cap.get(cv2.CAP_PROP_FPS) or 25
+    fps = cap.get(cv2.CAP_PROP_FPS) or 24
 
-   
+    # VideoWriter لنرجع فيديو بالـ bounding boxes
     output_video_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     out = cv2.VideoWriter(output_video_path, fourcc, fps, (640, 360))
@@ -75,10 +110,10 @@ def process_video(video_file, selected_classes):
     frame_id = 0
     prev_frame = None
 
-    track_memory = {}
-    track_stats = {}
-    track_bbox = {}        
-    track_motion = {}      
+    track_memory   = {}
+    track_stats    = {}
+    track_motion   = {}
+    track_activity = {}
 
     selected_ids = [k for k, v in class_names.items() if v in selected_classes]
 
@@ -89,9 +124,7 @@ def process_video(video_file, selected_classes):
 
         frame_id += 1
 
-        if frame_id % 2 != 0:
-            continue
-
+        # ✅ نشتغل على كل الفريمات (مش كل فريم تاني)
         frame = cv2.resize(frame, (640, 360))
 
         results = model(frame)[0]
@@ -118,42 +151,47 @@ def process_video(video_file, selected_classes):
             cy = (y1 + y2) / 2
 
             if track_id not in track_memory:
-                track_memory[track_id] = []
-                track_stats[track_id] = {"active": 0, "idle": 0}
-                track_motion[track_id] = "none"
+                track_memory[track_id]   = []
+                track_stats[track_id]    = {"active": 0, "idle": 0}
+                track_motion[track_id]   = "none"
+                track_activity[track_id] = "WAITING"
 
             track_memory[track_id].append((cx, cy))
-            if len(track_memory[track_id]) > 10:
+            if len(track_memory[track_id]) > 15:
                 track_memory[track_id].pop(0)
 
-            track_bbox[track_id] = bbox
-
+            # Zone-based articulated motion
             if prev_frame is not None:
-                is_active, motion_source = analyze_articulated_motion(frame, prev_frame, bbox)
+                is_active, motion_source = analyze_articulated_motion(
+                    frame, prev_frame, bbox
+                )
                 track_motion[track_id] = motion_source
             else:
-                is_active = False
+                is_active     = False
                 motion_source = "none"
 
-            activity = get_activity(track_memory[track_id])
+            activity = get_activity(
+                track_memory[track_id], is_active, motion_source
+            )
+            track_activity[track_id] = activity
 
-            if activity == "WAITING" and not is_active:
+            if activity == "WAITING":
                 track_stats[track_id]["idle"] += 1
             else:
                 track_stats[track_id]["active"] += 1
 
-           
+            # رسم على الفريم
             ix1, iy1, ix2, iy2 = map(int, bbox)
             color = (0, 255, 0) if activity != "WAITING" else (0, 0, 255)
 
             cv2.rectangle(frame, (ix1, iy1), (ix2, iy2), color, 2)
-
-            label = f"EX-{track_id} | {activity}"
-            cv2.putText(frame, label, (ix1, iy1 - 8),
+            cv2.putText(frame,
+                        f"EX-{track_id} | {activity}",
+                        (ix1, max(iy1 - 8, 10)),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1)
-
-            cv2.putText(frame, f"motion: {track_motion[track_id]}",
-                        (ix1, iy2 + 14),
+            cv2.putText(frame,
+                        f"motion: {motion_source}",
+                        (ix1, min(iy2 + 14, frame.shape[0] - 4)),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.35, (255, 255, 0), 1)
 
         out.write(frame)
@@ -162,46 +200,54 @@ def process_video(video_file, selected_classes):
     cap.release()
     out.release()
 
-    
+    # ===================================================
+    # Build JSON — كل الـ tracks ✅
+    # ===================================================
     final_output = []
 
     for track_id, data in track_stats.items():
-        track_id = int(track_id)   
+        tid    = int(track_id)
         active = data["active"]
-        idle = data["idle"]
-        total = active + idle
+        idle   = data["idle"]
+        total  = active + idle
 
         active_sec = active / fps
-        idle_sec = idle / fps
-        total_sec = total / fps
+        idle_sec   = idle   / fps
+        total_sec  = total  / fps
 
         utilization = (active_sec / (total_sec + 1e-6)) * 100
 
-        timestamp = f"00:00:{(frame_id / fps):06.3f}"
+        total_video_sec = frame_id / fps
+        mm = int(total_video_sec // 60)
+        ss = total_video_sec % 60
+        timestamp = f"00:{mm:02d}:{ss:06.3f}"
+
+        last_activity = track_activity.get(track_id, "WAITING")
+        motion_src    = track_motion.get(track_id,
+                        track_motion.get(str(track_id), "none"))
 
         result = {
             "frame_id": int(frame_id),
-            "equipment_id": f"EX-{track_id:03d}",
+            "equipment_id": f"EX-{tid:03d}",
             "equipment_class": "excavator",
             "timestamp": timestamp,
             "utilization": {
-                "current_state": "ACTIVE" if active_sec > idle_sec else "IDLE",
-                "current_activity": "DIGGING",
-                "motion_source": track_motion.get(track_id, track_motion.get(str(track_id), "none"))
+                "current_state":    "ACTIVE" if active_sec > idle_sec else "IDLE",
+                "current_activity": last_activity,
+                "motion_source":    motion_src
             },
             "time_analytics": {
-                "total_tracked_seconds": round(total_sec, 1),
-                "total_active_seconds": round(active_sec, 1),
-                "total_idle_seconds": round(idle_sec, 1),
-                "utilization_percent": round(utilization, 1)
+                "total_tracked_seconds": round(total_sec,   1),
+                "total_active_seconds":  round(active_sec,  1),
+                "total_idle_seconds":    round(idle_sec,    1),
+                "utilization_percent":   round(utilization, 1)
             }
         }
-
         final_output.append(result)
 
     json_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json").name
     with open(json_file, "w") as f:
-        json.dump(final_output, f, indent=2)   
+        json.dump(final_output, f, indent=2)
 
     return output_video_path, json.dumps(final_output, indent=2), json_file
 
@@ -216,11 +262,11 @@ demo = gr.Interface(
         gr.CheckboxGroup(choices=list(class_names.values()), label="Select Classes")
     ],
     outputs=[
-        gr.Video(label="📹 Processed Video"),          
+        gr.Video(label="📹 Processed Video"),
         gr.Textbox(label="JSON Output", lines=25),
         gr.File(label="⬇️ Download JSON")
     ],
-    title="Equipment Utilization & Activity Classification Prototype"
+    title="🚜 Excavator Activity Analyzer"
 )
 
 demo.launch()