Update app.py

app.py

@@ -10,86 +10,42 @@ model = YOLO("best.pt")
 class_names = model.names
 tracker = DeepSort(max_age=30)
 
-
 def analyze_articulated_motion(frame, prev_frame, bbox):
     x1, y1, x2, y2 = map(int, bbox)
-
     x1, y1 = max(0, x1), max(0, y1)
     x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
 
     h, w = y2 - y1, x2 - x1
-    if h < 20 or w < 20:
-        return False, "none"
-
-    mid_y = y1 + int(h * 0.4)  
-    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)
-    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)
-
-    def get_motion_score(img1, img2):
-        diff = cv2.absdiff(img1, img2)
-        _, thresh = cv2.threshold(diff, 15, 255, cv2.THRESH_BINARY)
-        return np.mean(thresh)
+    if h < 10 or w < 10: return False, "none"
 
-    upper_score = get_motion_score(upper_curr, upper_prev)
-    lower_score = get_motion_score(lower_curr, lower_prev)
-
-    sensitivity = 0.5 
+    mid_y = y1 + int(h * 0.5)
     
-    if upper_score > sensitivity:
-        if lower_score < (sensitivity * 1.5):
+    try:
+        roi_curr = cv2.cvtColor(frame[y1:mid_y, x1:x2], cv2.COLOR_BGR2GRAY)
+        roi_prev = cv2.cvtColor(prev_frame[y1:mid_y, x1:x2], cv2.COLOR_BGR2GRAY)
+        
+        diff = cv2.absdiff(roi_curr, roi_prev)
+        _, thresh = cv2.threshold(diff, 12, 255, cv2.THRESH_BINARY)
+        motion_score = np.mean(thresh)
+        
+        if motion_score > 0.15:
             return True, "arm_only"
-        else:
-            return True, "full_body"
-    elif lower_score > sensitivity:
-        return True, "full_body"
-    
+    except:
+        pass
     return False, "none"
 
-
-def get_activity(history, is_active, motion_source):
-    if not is_active:
-        return "WAITING"
-
-    if len(history) < 5:
-        return "DIGGING"
-
-    dx = history[-1][0] - history[0][0]
-    dy = history[-1][1] - history[0][1]
-
-    if motion_source == "arm_only":
-        if dy > 3: return "DIGGING"
-        if dy < -3: return "DUMPING"
-        return "DIGGING" 
-
-    if abs(dx) > abs(dy) * 1.5:
-        return "SWINGING/LOADING"
-    if dy > 5:
-        return "DIGGING"
-    if dy < -5:
-        return "DUMPING"
-
-    return "WORKING"
-
-def process_video(video_file, selected_classes):
+def process_video(video_file):
     cap = cv2.VideoCapture(video_file)
     fps = cap.get(cv2.CAP_PROP_FPS) or 24
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
+    
     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))
 
     frame_id = 0
     prev_frame = None
-    
-    track_memory = {}   
-    track_stats = {}    
-    frame_results = []  
-
-    selected_ids = [k for k, v in class_names.items() if v in selected_classes]
+    track_stats = {}
+    final_json_data = []
 
     while True:
         ret, frame = cap.read()
@@ -97,15 +53,18 @@ def process_video(video_file, selected_classes):
 
         frame_id += 1
         frame_resized = cv2.resize(frame, (640, 360))
-        results = model(frame_resized, verbose=False)[0]
         
+        results = model(frame_resized, verbose=False)[0]
         detections = []
+        
         for box in results.boxes:
             cls_id = int(box.cls[0])
-            if cls_id in selected_ids:
+            label = class_names[cls_id]
+            
+            if label == "C_E":  
                 x1, y1, x2, y2 = box.xyxy[0].tolist()
                 conf = float(box.conf[0])
-                detections.append(([x1, y1, x2 - x1, y2 - y1], conf, cls_id))
+                detections.append(([x1, y1, x2 - x1, y2 - y1], conf, "excavator"))
 
         tracks = tracker.update_tracks(detections, frame=frame_resized)
 
@@ -114,52 +73,47 @@ def process_video(video_file, selected_classes):
 
             track_id = t.track_id
             bbox = t.to_ltrb()
-            cx, cy = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
-
-            if track_id not in track_memory:
-                track_memory[track_id] = []
-                track_stats[track_id] = {"active_frames": 0, "total_frames": 0}
-
-            track_memory[track_id].append((cx, cy))
-            if len(track_memory[track_id]) > 20: track_memory[track_id].pop(0)
-
+            
             is_active = False
             motion_src = "none"
             if prev_frame is not None:
                 is_active, motion_src = analyze_articulated_motion(frame_resized, prev_frame, bbox)
-            
-            current_act = get_activity(track_memory[track_id], is_active, motion_src)
 
-            track_stats[track_id]["total_frames"] += 1
-            if current_act != "WAITING":
-                track_stats[track_id]["active_frames"] += 1
+            if track_id not in track_stats:
+                track_stats[track_id] = {"active_f": 0, "total_f": 0}
+            
+            track_stats[track_id]["total_f"] += 1
+            if is_active: track_stats[track_id]["active_f"] += 1
 
             color = (0, 255, 0) if is_active else (0, 0, 255)
             ix1, iy1, ix2, iy2 = map(int, bbox)
             cv2.rectangle(frame_resized, (ix1, iy1), (ix2, iy2), color, 2)
-            cv2.putText(frame_resized, f"ID:{track_id} {current_act}", (ix1, iy1-10), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-
-            total_sec = track_stats[track_id]["total_frames"] / fps
-            active_sec = track_stats[track_id]["active_frames"] / fps
-            util_pct = (active_sec / total_sec) * 100 if total_sec > 0 else 0
-
-            entry = {
-                "frame_id": frame_id,
-                "equipment_id": f"EX-{track_id}",
-                "timestamp": f"{frame_id/fps:.3f}",
-                "utilization": {
-                    "current_state": "ACTIVE" if is_active else "INACTIVE", 
-                    "current_activity": current_act,
-                    "motion_source": motion_src
-                },
-                "time_analytics": {
-                    "total_tracked_seconds": round(total_sec, 1),
-                    "total_active_seconds": round(active_sec, 1),
-                    "utilization_percent": round(util_pct, 1)
-                }
-            }
-            track_stats[track_id]["last_entry"] = entry
+            cv2.putText(frame_resized, f"EX-{track_id} | {'ACTIVE' if is_active else 'IDLE'}", 
+                        (ix1, iy1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+            if frame_id % int(fps) == 0 or frame_id == 1:
+                total_sec = track_stats[track_id]["total_f"] / fps
+                active_sec = track_stats[track_id]["active_f"] / fps
+                idle_sec = total_sec - active_sec
+                util_pct = (active_sec / total_sec) * 100 if total_sec > 0 else 0
+
+                final_json_data.append({
+                    "frame_id": frame_id,
+                    "equipment_id": f"EX-{track_id}",
+                    "equipment_class": "excavator",
+                    "timestamp": f"00:00:{frame_id/fps:06.3f}",
+                    "utilization": {
+                        "current_state": "ACTIVE" if is_active else "INACTIVE",
+                        "current_activity": "WORKING" if is_active else "WAITING",
+                        "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(util_pct, 1)
+                    }
+                })
 
         out.write(frame_resized)
         prev_frame = frame_resized.copy()
@@ -167,27 +121,21 @@ def process_video(video_file, selected_classes):
     cap.release()
     out.release()
 
-    final_json_data = [stats["last_entry"] for stats in track_stats.values() if "last_entry" in stats]
-    
     json_path = tempfile.NamedTemporaryFile(delete=False, suffix=".json").name
     with open(json_path, "w") as f:
         json.dump(final_json_data, f, indent=2)
 
     return output_video_path, json.dumps(final_json_data, indent=2), json_path
 
-
 demo = gr.Interface(
     fn=process_video,
-    inputs=[
-        gr.Video(label="Upload Construction Video"),
-        gr.CheckboxGroup(choices=["excavator", "dump truck", "loader"], label="Equipment to Track", value=["excavator"])
-    ],
+    inputs=gr.Video(label="Upload Construction Video"),
     outputs=[
-        gr.Video(label="Detection Feed"),
-        gr.Textbox(label="Real-time Analytics (JSON)", lines=15),
-        gr.File(label="Download Full Report")
+        gr.Video(label="Processed Video"),
+        gr.Textbox(label="Kafka-ready JSON Payload", lines=20),
+        gr.File(label="Download Full JSON Report")
     ],
-    title="Equipment Utilization Prototype"
+    title=" Equipment Utilization System",
 )
 
 if __name__ == "__main__":