Update app.py

app.py

@@ -19,6 +19,7 @@ from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
+import face_recognition
 from collections import defaultdict
 
 # ========================== # Configuration and Setup # ==========================
@@ -28,134 +29,174 @@ os.makedirs('/tmp/Ultralytics', exist_ok=True)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# ========================== # Enhanced BYTETracker Implementation # ==========================
-class EnhancedBYTETracker:
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+# ========================== # Enhanced Tracker Implementation # ==========================
+class SafetyTracker:
     def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
         self.match_thresh = match_thresh
         self.frame_rate = frame_rate
         self.next_id = 1
-        self.tracks = {}  # Store active tracks
+        
+        # Trackers for different purposes
+        self.worker_tracks = {}  # Active worker tracks
         self.violation_history = defaultdict(dict)  # Track violations per worker
-        self.last_positions = {}  # Last known positions of workers
-        self.violation_cooldown = {}  # Cooldown periods for violations
+        self.face_encodings = {}  # Store face encodings for helmet violations
+        self.position_history = defaultdict(list)  # Track positions for non-helmet violations
+        
+        # Cooldown periods (in seconds)
+        self.VIOLATION_COOLDOWNS = {
+            "no_helmet": 30.0,
+            "no_harness": 20.0,
+            "unsafe_posture": 15.0,
+            "unsafe_zone": 10.0,
+            "improper_tool_use": 15.0
+        }
 
-    def update(self, dets, scores, cls, frame_idx):
-        tracks = []
-        current_time = frame_idx / self.frame_rate
+    def update(self, detections, frame):
+        """Update tracks with new detections and check for violations"""
+        current_time = time.time()
+        active_violations = []
+        new_violations = []
         
-        # Update existing tracks with new detections
-        for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
-            if score < self.track_thresh:
-                continue
-                
-            x, y, w, h = det
-            matched = False
-            best_iou = 0
-            best_track_id = None
+        for det in detections:
+            bbox = det['bbox']
+            label = det['violation']
+            confidence = det['confidence']
             
-            # Try to match with existing tracks
-            for track_id, track_info in self.tracks.items():
-                if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
-                    continue
-                    
-                tx, ty, tw, th = track_info['bbox']
-                iou = self._calculate_iou([x, y, w, h], [tx, ty, tw, th])
-                
-                if iou > self.match_thresh and iou > best_iou:
-                    best_iou = iou
-                    best_track_id = track_id
-                    matched = True
+            # For helmet violations, use face recognition
+            if label == "no_helmet":
+                worker_id = self._match_by_face(bbox, frame)
+            else:
+                # For other violations, use position tracking
+                worker_id = self._match_by_position(bbox, label)
             
-            if matched:
-                # Update existing track
-                self.tracks[best_track_id].update({
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl,
-                    'last_seen': current_time
-                })
-                
-                # Update position
-                self.last_positions[best_track_id] = [x, y]
+            if worker_id is None:
+                worker_id = self.next_id
+                self.next_id += 1
                 
-                tracks.append({
-                    'id': best_track_id,
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl
-                })
-            else:
-                # Create new track
-                self.tracks[self.next_id] = {
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl,
-                    'last_seen': current_time
+            # Check if this is a new violation for this worker
+            if self._is_new_violation(worker_id, label, current_time):
+                # Record the violation
+                violation = {
+                    'worker_id': worker_id,
+                    'violation': label,
+                    'confidence': confidence,
+                    'bbox': bbox,
+                    'timestamp': current_time
                 }
-                self.last_positions[self.next_id] = [x, y]
-                tracks.append({
-                    'id': self.next_id,
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl
-                })
-                self.next_id += 1
+                new_violations.append(violation)
+                
+                # Update violation history
+                self.violation_history[worker_id][label] = current_time
+                
+                # For helmet violations, store face encoding
+                if label == "no_helmet":
+                    self._store_face_encoding(worker_id, bbox, frame)
+            
+            # Keep track of active workers
+            self.worker_tracks[worker_id] = {
+                'bbox': bbox,
+                'last_seen': current_time,
+                'label': label
+            }
         
         # Clean up old tracks
-        stale_ids = [tid for tid, track in self.tracks.items() 
-                    if current_time - track['last_seen'] > self.track_buffer / self.frame_rate]
+        self._cleanup_tracks(current_time)
         
-        for track_id in stale_ids:
-            del self.tracks[track_id]
-            if track_id in self.last_positions:
-                del self.last_positions[track_id]
-            
-        return tracks
-    
-    def _calculate_iou(self, box1, box2):
-        """Calculate IOU between two boxes"""
-        x1, y1, w1, h1 = box1
-        x2, y2, w2, h2 = box2
+        return new_violations
+
+    def _match_by_face(self, bbox, frame):
+        """Match detection by face recognition (for helmet violations)"""
+        x, y, w, h = bbox
+        face_region = frame[max(0, int(y-h/2)):int(y+h/2), max(0, int(x-w/2)):int(x+w/2)]
         
-        # Calculate intersection coordinates
-        x_left = max(x1 - w1/2, x2 - w2/2)
-        y_top = max(y1 - h1/2, y2 - h2/2)
-        x_right = min(x1 + w1/2, x2 + w2/2)
-        y_bottom = min(y1 + h1/2, y2 + h2/2)
+        if face_region.size == 0:
+            return None
+            
+        try:
+            # Get face encodings from current detection
+            face_locations = face_recognition.face_locations(face_region)
+            if not face_locations:
+                return None
+                
+            current_encoding = face_recognition.face_encodings(face_region, face_locations)[0]
+            
+            # Compare with known faces
+            for worker_id, encodings in self.face_encodings.items():
+                matches = face_recognition.compare_faces(encodings, current_encoding, tolerance=0.6)
+                if any(matches):
+                    return worker_id
+                    
+        except Exception as e:
+            logger.warning(f"Face recognition error: {e}")
+            
+        return None
+
+    def _match_by_position(self, bbox, label):
+        """Match detection by position (for non-helmet violations)"""
+        x, y, w, h = bbox
+        current_pos = (x, y)
         
-        if x_right < x_left or y_bottom < y_top:
-            return 0.0
+        for worker_id, positions in self.position_history.items():
+            if label not in self.violation_history[worker_id]:
+                continue
+                
+            # Check if current position is near any previous positions for this worker
+            for pos in positions:
+                distance = np.sqrt((current_pos[0]-pos[0])**2 + (current_pos[1]-pos[1])**2)
+                if distance < 100:  # Within 100 pixels
+                    return worker_id
+                    
+        return None
+
+    def _is_new_violation(self, worker_id, label, current_time):
+        """Check if this is a new violation for this worker"""
+        if label not in self.violation_history[worker_id]:
+            return True
             
-        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+        last_detection = self.violation_history[worker_id][label]
+        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
         
-        box1_area = w1 * h1
-        box2_area = w2 * h2
+        return (current_time - last_detection) > cooldown
+
+    def _store_face_encoding(self, worker_id, bbox, frame):
+        """Store face encoding for a worker"""
+        x, y, w, h = bbox
+        face_region = frame[max(0, int(y-h/2)):int(y+h/2), max(0, int(x-w/2)):int(x+w/2)]
         
-        iou = intersection_area / (box1_area + box2_area - intersection_area)
-        return iou
-    
-    def is_new_violation(self, worker_id, violation_type, bbox, current_time):
-        """Check if this is a new violation that should be reported"""
-        # Check if this worker has had this violation type before
-        if violation_type in self.violation_history[worker_id]:
-            last_time, last_bbox = self.violation_history[worker_id][violation_type]
+        if face_region.size == 0:
+            return
             
-            # For no_helmet, we only report once per worker
-            if violation_type == "no_helmet":
-                return False
-                
-            # For other violations, check if it's in a new location or after cooldown
-            if current_time - last_time < CONFIG["VIOLATION_COOLDOWN"]:
-                # Check if it's in the same area (using IOU)
-                iou = self._calculate_iou(bbox, last_bbox)
-                if iou > 0.3:  # If overlap > 30%, consider it the same violation
-                    return False
-                    
-        # If we get here, it's a new violation
-        self.violation_history[worker_id][violation_type] = (current_time, bbox)
-        return True
+        try:
+            face_locations = face_recognition.face_locations(face_region)
+            if face_locations:
+                encoding = face_recognition.face_encodings(face_region, face_locations)[0]
+                if worker_id not in self.face_encodings:
+                    self.face_encodings[worker_id] = []
+                self.face_encodings[worker_id].append(encoding)
+        except Exception as e:
+            logger.warning(f"Error storing face encoding: {e}")
+
+    def _cleanup_tracks(self, current_time):
+        """Clean up old tracks and face encodings"""
+        # Remove inactive workers
+        inactive_ids = [
+            worker_id for worker_id, track in self.worker_tracks.items()
+            if (current_time - track['last_seen']) > (self.track_buffer / self.frame_rate)
+        ]
+        
+        for worker_id in inactive_ids:
+            self.worker_tracks.pop(worker_id, None)
+            self.position_history.pop(worker_id, None)
+            
+            # Keep face encodings for a longer period (for helmet violations)
+            if (current_time - max(self.violation_history[worker_id].values(), default=0)) > 300:  # 5 minutes
+                self.face_encodings.pop(worker_id, None)
+                self.violation_history.pop(worker_id, None)
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -197,18 +238,12 @@ CONFIG = {
         "unsafe_zone": 0.3,
         "improper_tool_use": 0.3
     },
-    "MIN_VIOLATION_FRAMES": 5,
-    "VIOLATION_COOLDOWN": 30.0,  # 30 seconds cooldown for same violation type in same area
-    "WORKER_TRACKING_DURATION": 5.0,
-    "MAX_PROCESSING_TIME": 60,
+    "MIN_VIOLATION_FRAMES": 1,
     "FRAME_SKIP": 2,
     "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
+    "FACE_RECOGNITION_INTERVAL": 5  # Process face recognition every 5 frames
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -283,22 +318,13 @@ def calculate_safety_score(violations):
         "improper_tool_use": 25
     }
     
-    # Count unique violation types per worker
-    worker_violations = {}
+    # Count unique violation types
+    unique_violations = set()
     for v in violations:
-        worker_id = v.get("worker_id", "Unknown")
         violation_type = v.get("violation", "Unknown")
-        
-        if worker_id not in worker_violations:
-            worker_violations[worker_id] = set()
-        worker_violations[worker_id].add(violation_type)
-    
-    # Calculate total penalty
-    total_penalty = 0
-    for worker_violations_set in worker_violations.values():
-        worker_penalty = sum(penalties.get(v, 0) for v in worker_violations_set)
-        total_penalty += worker_penalty
+        unique_violations.add(violation_type)
     
+    total_penalty = sum(penalties.get(v, 0) for v in unique_violations)
     score = max(0, 100 - total_penalty)
     return score
 
@@ -329,18 +355,10 @@ def generate_violation_pdf(violations, score):
         c.drawString(1 * inch, y_position, "Summary:")
         y_position -= 0.3 * inch
         
-        # Group violations by worker
-        worker_violations = {}
-        for v in violations:
-            worker_id = v.get("worker_id", "Unknown")
-            if worker_id not in worker_violations:
-                worker_violations[worker_id] = []
-            worker_violations[worker_id].append(v)
-        
         c.setFont("Helvetica", 10)
         summary_data = {
-            "Total Workers with Violations": len(worker_violations),
             "Total Violations Found": len(violations),
+            "Unique Violation Types": len(set(v['violation'] for v in violations)),
             "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
         
@@ -348,30 +366,27 @@ def generate_violation_pdf(violations, score):
             c.drawString(1 * inch, y_position, f"{key}: {value}")
             y_position -= 0.25 * inch
 
-        # Detailed Violations by Worker
+        # Detailed Violations
         y_position -= 0.5 * inch
         c.setFont("Helvetica-Bold", 12)
-        c.drawString(1 * inch, y_position, "Violations by Worker:")
+        c.drawString(1 * inch, y_position, "Violation Details:")
         y_position -= 0.3 * inch
         
         c.setFont("Helvetica", 10)
-        for worker_id, worker_vios in worker_violations.items():
-            c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
+        for v in violations:
+            display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
+            worker_id = v.get("worker_id", "Unknown")
+            time_str = f"{v.get('timestamp', 0.0):.2f}s"
+            conf_str = f"{v.get('confidence', 0.0):.2f}"
+            
+            violation_text = f"- {display_name} by Worker {worker_id} at {time_str} (Confidence: {conf_str})"
+            c.drawString(1.2 * inch, y_position, violation_text)
             y_position -= 0.2 * inch
             
-            for v in worker_vios:
-                display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-                time_str = f"{v.get('timestamp', 0.0):.2f}s"
-                conf_str = f"{v.get('confidence', 0.0):.2f}"
-                
-                violation_text = f"  - {display_name} at {time_str} (Confidence: {conf_str})"
-                c.drawString(1.2 * inch, y_position, violation_text)
-                y_position -= 0.2 * inch
-                
-                if y_position < 1 * inch:
-                    c.showPage()
-                    c.setFont("Helvetica", 10)
-                    y_position = 10 * inch
+            if y_position < 1 * inch:
+                c.showPage()
+                c.setFont("Helvetica", 10)
+                y_position = 10 * inch
 
         c.save()
         pdf_file.seek(0)
@@ -507,18 +522,12 @@ def process_video(video_data):
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        tracker = EnhancedBYTETracker(
-            track_thresh=CONFIG["TRACK_THRESH"],
-            track_buffer=CONFIG["TRACK_BUFFER"],
-            match_thresh=CONFIG["MATCH_THRESH"],
-            frame_rate=fps
-        )
-
-        violations = []
+        tracker = SafetyTracker(frame_rate=fps)
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
+        frame_counter = 0
 
         while processed_frames < total_frames:
             batch_frames = []
@@ -543,6 +552,7 @@ def process_video(video_data):
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
                 processed_frames += 1
+                frame_counter += 1
 
             if not batch_frames:
                 break
@@ -560,7 +570,7 @@ def process_video(video_data):
                     start_time = time.time()
 
                 boxes = result.boxes
-                track_inputs = []
+                detections = []
                 
                 for box in boxes:
                     cls = int(box.cls)
@@ -574,78 +584,54 @@ def process_video(video_data):
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
-                    track_inputs.append({
+                    detections.append({
                         "bbox": bbox,
-                        "conf": conf,
-                        "cls": cls
+                        "violation": label,
+                        "confidence": conf
                     })
 
-                if not track_inputs:
+                if not detections:
                     continue
                     
-                tracked_objects = tracker.update(
-                    np.array([t["bbox"] for t in track_inputs]),
-                    np.array([t["conf"] for t in track_inputs]),
-                    np.array([t["cls"] for t in track_inputs]),
-                    frame_idx
-                )
-
-                # Process tracked objects for violations
-                for obj in tracked_objects:
-                    worker_id = obj['id']
-                    label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
-                    conf = obj['score']
-                    bbox = obj['bbox']
+                # Update tracker with new detections
+                new_violations = tracker.update(detections, batch_frames[i])
+                
+                # Process new violations
+                for violation in new_violations:
+                    # Take snapshot for the new violation
+                    snapshot_frame = batch_frames[i].copy()
+                    snapshot_frame = draw_detections(snapshot_frame, [violation])
                     
-                    if label is None:
-                        continue
-                        
-                    # Check if this is a new violation that should be reported
-                    if tracker.is_new_violation(worker_id, label, bbox, current_time):
-                        # This is a new violation that should be reported
-                        detection = {
-                            "worker_id": worker_id,
-                            "violation": label,
-                            "confidence": round(conf, 2),
-                            "bounding_box": bbox,
-                            "timestamp": current_time
-                        }
-                        violations.append(detection)
-                        
-                        # 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")
+                    # Add timestamp to snapshot
+                    cv2.putText(
+                        snapshot_frame,
+                        f"Time: {violation['timestamp']:.2f}s",
+                        (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.7,
+                        (255, 255, 255),
+                        2
+                    )
+                    
+                    # Save snapshot with high quality
+                    snapshot_filename = f"violation_{violation['violation']}_worker{violation['worker_id']}_{int(violation['timestamp']*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": violation['violation'],
+                        "worker_id": violation['worker_id'],
+                        "timestamp": violation['timestamp'],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    })
+                    
+                    logger.info(f"Captured snapshot for {violation['violation']} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
 
         cap.release()
         if os.path.exists(video_path):
@@ -654,6 +640,16 @@ def process_video(video_data):
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
+        # Get all unique violations from tracker
+        violations = []
+        for worker_id, worker_violations in tracker.violation_history.items():
+            for label, detection_time in worker_violations.items():
+                violations.append({
+                    "worker_id": worker_id,
+                    "violation": label,
+                    "timestamp": detection_time
+                })
+
         if not violations:
             logger.info("No violations detected after processing")
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
@@ -669,16 +665,15 @@ def process_video(video_data):
         report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
         # Format violations table for display
-        violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
-        violation_table += "|-----------|-----------|----------|------------|\n"
+        violation_table = "| Violation | Worker ID | Time (s) |\n"
+        violation_table += "|-----------|-----------|----------|\n"
         
-        for v in sorted(violations, key=lambda x: (x.get("worker_id", "Unknown"), x.get("timestamp", 0.0))):
+        for v in sorted(violations, key=lambda x: x.get("timestamp", 0.0)):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
             worker_id = v.get("worker_id", "Unknown")
             timestamp = v.get("timestamp", 0.0)
-            confidence = v.get("confidence", 0.0)
             
-            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
+            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} |\n"
 
         # Format snapshots for display
         snapshots_text = ""