Update app.py

app.py

@@ -38,6 +38,7 @@ class BYTETracker:
         self.tracks = {}  # Store active tracks
         self.worker_history = {}  # Track worker positions over time
         self.last_positions = {}  # Last known positions of workers
+        self.violation_history = {}  # Track violations per worker: {worker_id: set(violation_types)}
 
     def update(self, dets, scores, cls):
         tracks = []
@@ -117,6 +118,7 @@ class BYTETracker:
                     }
                     self.worker_history[self.next_id] = [[x, y]]
                     self.last_positions[self.next_id] = [x, y]
+                    self.violation_history[self.next_id] = set()  # Initialize violation set for new worker
                     tracks.append({
                         'id': self.next_id,
                         'bbox': [x, y, w, h],
@@ -138,9 +140,21 @@ class BYTETracker:
                 del self.worker_history[track_id]
             if track_id in self.last_positions:
                 del self.last_positions[track_id]
+            if track_id in self.violation_history:
+                del self.violation_history[track_id]
             
         return tracks
     
+    def has_violation(self, worker_id, violation_type):
+        """Check if this worker already has this violation type recorded"""
+        return worker_id in self.violation_history and violation_type in self.violation_history[worker_id]
+    
+    def record_violation(self, worker_id, violation_type):
+        """Record that this worker has this violation type"""
+        if worker_id not in self.violation_history:
+            self.violation_history[worker_id] = set()
+        self.violation_history[worker_id].add(violation_type)
+    
     def _calculate_iou(self, box1, box2):
         """Calculate IOU between two boxes"""
         x1, y1, w1, h1 = box1
@@ -211,17 +225,17 @@ CONFIG = {
         "improper_tool_use": 0.3
     },
     "MIN_VIOLATION_FRAMES": 1,
-    "VIOLATION_COOLDOWN": 30.0,
+    "VIOLATION_COOLDOWN": 30.0,  # Increased cooldown period
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,
+    "FRAME_SKIP": 2,  # Skip more frames for faster processing
     "BATCH_SIZE": 16,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 30,
     "TRACK_THRESH": 0.3,
     "MATCH_THRESH": 0.7,
-    "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 100
+    "SNAPSHOT_QUALITY": 95,  # Higher quality for better visibility
+    "MAX_WORKER_DISTANCE": 100  # Maximum pixel distance to consider same worker
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -498,7 +512,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         return None, ""
 
 def process_video(video_data):
-    """Process video to detect safety violations with de-duplication"""
+    """Process video to detect safety violations"""
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
@@ -528,7 +542,7 @@ def process_video(video_data):
         )
 
         # Track unique violations by worker ID
-        unique_violations = {}  # {worker_id: {violation_type: {"timestamp": float, "confidence": float, "bbox": list}}}
+        unique_violations = {}  # {worker_id: {violation_type: first_detection_time}}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
@@ -613,62 +627,61 @@ def process_video(video_data):
                     if label is None:
                         continue
                         
+                    # Skip if this worker already has this violation recorded
+                    if tracker.has_violation(worker_id, label):
+                        continue
+                        
                     # Initialize worker if not seen before
                     if worker_id not in unique_violations:
                         unique_violations[worker_id] = {}
                     
-                    # Check if this violation type has been recorded for this worker
-                    if label not in unique_violations[worker_id]:
-                        # This is a new violation type for this worker
-                        unique_violations[worker_id][label] = {
-                            "timestamp": current_time,
-                            "confidence": round(conf, 2),
-                            "bbox": bbox
-                        }
-                        
-                        # Create detection object
-                        detection = {
-                            "worker_id": worker_id,
-                            "violation": label,
-                            "confidence": round(conf, 2),
-                            "bounding_box": bbox,
-                            "timestamp": current_time
-                        }
-                        
-                        # Take snapshot for the new violation
-                        snapshot_frame = batch_frames[i].copy()
-                        snapshot_frame = draw_detections(snapshot_frame, [detection])
-                        
-                        # Add timestamp to snapshot
-                        cv2.putText(
-                            snapshot_frame,
-                            f"Time: {current_time:.2f}s",
-                            (10, 30),
-                            cv2.FONT_HERSHEY_SIMPLEX,
-                            0.7,
-                            (255, 255, 255),
-                            2
-                        )
-                        
-                        # Save snapshot with high quality
-                        snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        
-                        cv2.imwrite(
-                            snapshot_path,
-                            snapshot_frame,
-                            [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
-                        )
-                        
-                        snapshots.append({
-                            "violation": label,
-                            "worker_id": worker_id,
-                            "timestamp": current_time,
-                            "snapshot_path": snapshot_path,
-                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                        })
-                        
-                        logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
+                    # Record this violation for this worker
+                    tracker.record_violation(worker_id, label)
+                    unique_violations[worker_id][label] = current_time
+                    
+                    # Create detection object
+                    detection = {
+                        "worker_id": worker_id,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    }
+                    
+                    # Take snapshot for the new violation
+                    snapshot_frame = batch_frames[i].copy()
+                    snapshot_frame = draw_detections(snapshot_frame, [detection])
+                    
+                    # Add timestamp to snapshot
+                    cv2.putText(
+                        snapshot_frame,
+                        f"Time: {current_time:.2f}s",
+                        (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.7,
+                        (255, 255, 255),
+                        2
+                    )
+                    
+                    # Save snapshot with high quality
+                    snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
+                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                    
+                    cv2.imwrite(
+                        snapshot_path,
+                        snapshot_frame,
+                        [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
+                    )
+                    
+                    snapshots.append({
+                        "violation": label,
+                        "worker_id": worker_id,
+                        "timestamp": current_time,
+                        "snapshot_path": snapshot_path,
+                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    })
+                    
+                    logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
 
         cap.release()
         if os.path.exists(video_path):
@@ -680,13 +693,11 @@ def process_video(video_data):
         # Convert tracked violations to final violation list
         violations = []
         for worker_id, worker_violations in unique_violations.items():
-            for label, info in worker_violations.items():
+            for label, detection_time in worker_violations.items():
                 violation = {
                     "worker_id": worker_id,
                     "violation": label,
-                    "timestamp": info["timestamp"],
-                    "confidence": info["confidence"],
-                    "bounding_box": info["bbox"]
+                    "timestamp": detection_time
                 }
                 violations.append(violation)