Update app.py

app.py

@@ -260,6 +260,10 @@ def process_video(video_data):
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         workers = []  # List to track workers
 
+        # Adding debug logging for violation labels
+        logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
+        logger.info(f"Using confidence threshold: {CONFIG['CONFIDENCE_THRESHOLD']}")
+
         while True:
             ret, frame = video.read()
             if not ret:
@@ -273,21 +277,35 @@ def process_video(video_data):
                 logger.info("Processing time limit reached")
                 break
 
+            # Run detection on this frame
             results = model(frame, device=device)
             current_detections = []
+            
+            # Process detections from the model
             for result in results:
-                for box in result.boxes:
+                boxes = result.boxes
+                logger.info(f"Frame {frame_count}: Found {len(boxes)} potential detections")
+                
+                for box in boxes:
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
-                    # Log detected violations
-                    if label:
-                        logger.info(f"Violation Detected: {label} with confidence: {conf}")
-
-                    if label not in CONFIG["VIOLATION_LABELS"].values() or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                    # Enhanced logging
+                    logger.info(f"Detection: class={cls}, conf={conf:.2f}, label={label}")
+                    
+                    # Skip if not a known violation or below confidence threshold
+                    if label not in CONFIG["VIOLATION_LABELS"].values():
+                        logger.info(f"Skipping unknown class: {cls}")
+                        continue
+                        
+                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                        logger.info(f"Skipping low confidence: {conf:.2f} < {CONFIG['CONFIDENCE_THRESHOLD']}")
                         continue
 
+                    # Process valid detection
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    logger.info(f"Valid detection: {label} with confidence: {conf:.2f}")
+                    
                     current_detections.append({
                         "violation": label,
                         "confidence": round(conf, 2),
@@ -296,10 +314,13 @@ def process_video(video_data):
                         "frame": frame_count
                     })
 
-            # Process detections and workers
+            # Process detections and associate with workers
+            # FIXED: Improved worker tracking logic
             for detection in current_detections:
                 matched_worker = None
                 max_iou = 0
+                
+                # Try to match with existing workers
                 for worker in workers:
                     iou = calculate_iou(detection["bounding_box"], worker["bbox"])
                     if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
@@ -309,6 +330,8 @@ def process_video(video_data):
                 if matched_worker:
                     # Update existing worker
                     if detection["violation"] not in matched_worker["violations"]:
+                        # New violation for this worker
+                        logger.info(f"New violation for worker {matched_worker['id']}: {detection['violation']}")
                         matched_worker["violations"].add(detection["violation"])
                         violations.append({
                             "frame": frame_count,
@@ -318,18 +341,34 @@ def process_video(video_data):
                             "timestamp": detection["timestamp"],
                             "worker_id": matched_worker["id"]
                         })
-                        snapshot_taken[detection["violation"]] = True
+                        
+                        # Save snapshot for this violation type if not already taken
+                        if not snapshot_taken[detection["violation"]]:
+                            snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
+                            snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                            cv2.imwrite(snapshot_path, frame)
+                            snapshot_taken[detection["violation"]] = True
+                            snapshots.append({
+                                "violation": detection["violation"],
+                                "frame": frame_count,
+                                "snapshot_path": snapshot_path,
+                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                            })
+                    
+                    # Update worker position
                     matched_worker["bbox"] = detection["bounding_box"]
                     matched_worker["last_frame"] = frame_count
                 else:
-                    # New worker
+                    # New worker detected
                     worker_id = len(workers) + 1
+                    logger.info(f"New worker {worker_id} with violation: {detection['violation']}")
                     workers.append({
                         "id": worker_id,
                         "violations": {detection["violation"]},
                         "bbox": detection["bounding_box"],
                         "last_frame": frame_count
                     })
+                    
                     violations.append({
                         "frame": frame_count,
                         "violation": detection["violation"],
@@ -338,11 +377,37 @@ def process_video(video_data):
                         "timestamp": detection["timestamp"],
                         "worker_id": worker_id
                     })
+                    
+                    # Save snapshot for this violation type if not already taken
+                    if not snapshot_taken[detection["violation"]]:
+                        snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                        cv2.imwrite(snapshot_path, frame)
+                        snapshot_taken[detection["violation"]] = True
+                        snapshots.append({
+                            "violation": detection["violation"],
+                            "frame": frame_count,
+                            "snapshot_path": snapshot_path,
+                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                        })
 
+            # Clean up workers that haven't been seen for a while
+            active_workers = [w for w in workers if frame_count - w["last_frame"] < CONFIG["FRAME_SKIP"] * 5]
+            if len(active_workers) != len(workers):
+                logger.info(f"Cleaned up {len(workers) - len(active_workers)} inactive workers")
+                workers = active_workers
+                
             frame_count += 1
 
         video.release()
         os.remove(video_path)
+        
+        # Final log of violations detected
+        violation_types = {}
+        for v in violations:
+            violation_types[v["violation"]] = violation_types.get(v["violation"], 0) + 1
+        
+        logger.info(f"Detection complete. Found violations: {violation_types}")
 
         if not violations:
             logger.info("No violations detected")
@@ -441,4 +506,4 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     logger.info("Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()