Update app.py

app.py

@@ -19,7 +19,6 @@ from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
-from collections import defaultdict
 
 # ========================== # Configuration and Setup # ==========================
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
@@ -28,174 +27,148 @@ os.makedirs('/tmp/Ultralytics', exist_ok=True)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# Suppress warnings
-warnings.filterwarnings("ignore")
-
-# ========================== # Enhanced Tracker Implementation # ==========================
-class SafetyTracker:
+# ========================== # ByteTrack Implementation # ==========================
+class BYTETracker:
     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
-        
-        # Trackers for different purposes
-        self.worker_tracks = {}  # Active worker tracks
-        self.violation_history = defaultdict(dict)  # Track violations per worker
-        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
-        }
+        self.tracks = {}  # Store active tracks
+        self.worker_history = {}  # Track worker positions over time
+        self.last_positions = {}  # Last known positions of workers
 
-    def update(self, detections, frame):
-        """Update tracks with new detections and check for violations"""
+    def update(self, dets, scores, cls):
+        tracks = []
         current_time = time.time()
-        active_violations = []
-        new_violations = []
         
-        for det in detections:
-            bbox = det['bbox']
-            label = det['violation']
-            confidence = det['confidence']
+        # 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 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)
+            # 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
             
-            if worker_id is None:
-                worker_id = self.next_id
-                self.next_id += 1
+            if matched:
+                # Update existing track
+                self.tracks[best_track_id].update({
+                    '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
-                }
-                new_violations.append(violation)
+                # Update position history
+                if best_track_id not in self.worker_history:
+                    self.worker_history[best_track_id] = []
+                self.worker_history[best_track_id].append([x, y])
+                self.last_positions[best_track_id] = [x, y]
                 
-                # Update violation history
-                self.violation_history[worker_id][label] = current_time
+                tracks.append({
+                    'id': best_track_id,
+                    'bbox': [x, y, w, h],
+                    'score': score,
+                    'cls': cl
+                })
+            else:
+                # Create new track
+                # Check if this detection might be the same worker from a different angle
+                same_worker = False
+                for worker_id, last_pos in self.last_positions.items():
+                    if self._is_same_worker([x, y], last_pos):
+                        self.tracks[worker_id] = {
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl,
+                            'last_seen': current_time
+                        }
+                        tracks.append({
+                            'id': worker_id,
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl
+                        })
+                        same_worker = True
+                        break
                 
-                # 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
-            }
+                if not same_worker:
+                    self.tracks[self.next_id] = {
+                        'bbox': [x, y, w, h],
+                        'score': score,
+                        'cls': cl,
+                        'last_seen': current_time
+                    }
+                    self.worker_history[self.next_id] = [[x, y]]
+                    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
         
         # Clean up old tracks
-        self._cleanup_tracks(current_time)
-        
-        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)]
-        
-        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)
+        current_time = time.time()
+        stale_ids = []
+        for track_id, track_info in self.tracks.items():
+            if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
+                stale_ids.append(track_id)
         
-        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
+        for track_id in stale_ids:
+            del self.tracks[track_id]
+            if track_id in self.worker_history:
+                del self.worker_history[track_id]
+            if track_id in self.last_positions:
+                del self.last_positions[track_id]
             
-        last_detection = self.violation_history[worker_id][label]
-        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
+        return tracks
+    
+    def _calculate_iou(self, box1, box2):
+        """Calculate IOU between two boxes"""
+        x1, y1, w1, h1 = box1
+        x2, y2, w2, h2 = box2
         
-        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)]
+        # 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
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
             
-        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)
-        ]
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
         
-        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)
+        box1_area = w1 * h1
+        box2_area = w2 * h2
+        
+        iou = intersection_area / (box1_area + box2_area - intersection_area)
+        return iou
+    
+    def _is_same_worker(self, pos1, pos2, threshold=100):
+        """Check if two positions likely belong to the same worker"""
+        x1, y1 = pos1
+        x2, y2 = pos2
+        distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
+        return distance < threshold
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -238,11 +211,17 @@ CONFIG = {
         "improper_tool_use": 0.3
     },
     "MIN_VIOLATION_FRAMES": 1,
-    "FRAME_SKIP": 2,
+    "VIOLATION_COOLDOWN": 30.0,  # Increased cooldown period
+    "WORKER_TRACKING_DURATION": 5.0,
+    "MAX_PROCESSING_TIME": 60,
+    "FRAME_SKIP": 2,  # Skip more frames for faster processing
     "BATCH_SIZE": 16,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "SNAPSHOT_QUALITY": 95,
-    "FACE_RECOGNITION_INTERVAL": 5  # Process face recognition every 5 frames
+    "TRACK_BUFFER": 30,
+    "TRACK_THRESH": 0.3,
+    "MATCH_THRESH": 0.7,
+    "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")
@@ -317,13 +296,22 @@ def calculate_safety_score(violations):
         "improper_tool_use": 25
     }
     
-    # Count unique violation types
-    unique_violations = set()
+    # Count unique violation types per worker
+    worker_violations = {}
     for v in violations:
+        worker_id = v.get("worker_id", "Unknown")
         violation_type = v.get("violation", "Unknown")
-        unique_violations.add(violation_type)
+        
+        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
     
-    total_penalty = sum(penalties.get(v, 0) for v in unique_violations)
     score = max(0, 100 - total_penalty)
     return score
 
@@ -354,10 +342,18 @@ 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")
         }
         
@@ -365,27 +361,30 @@ def generate_violation_pdf(violations, score):
             c.drawString(1 * inch, y_position, f"{key}: {value}")
             y_position -= 0.25 * inch
 
-        # Detailed Violations
+        # Detailed Violations by Worker
         y_position -= 0.5 * inch
         c.setFont("Helvetica-Bold", 12)
-        c.drawString(1 * inch, y_position, "Violation Details:")
+        c.drawString(1 * inch, y_position, "Violations by Worker:")
         y_position -= 0.3 * inch
         
         c.setFont("Helvetica", 10)
-        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)
+        for worker_id, worker_vios in worker_violations.items():
+            c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
             y_position -= 0.2 * inch
             
-            if y_position < 1 * inch:
-                c.showPage()
-                c.setFont("Helvetica", 10)
-                y_position = 10 * 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
 
         c.save()
         pdf_file.seek(0)
@@ -499,7 +498,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 enhanced tracking"""
+    """Process video to detect safety violations"""
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
@@ -521,12 +520,19 @@ 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 = SafetyTracker(frame_rate=fps)
+        tracker = BYTETracker(
+            track_thresh=CONFIG["TRACK_THRESH"],
+            track_buffer=CONFIG["TRACK_BUFFER"],
+            match_thresh=CONFIG["MATCH_THRESH"],
+            frame_rate=fps
+        )
+
+        # Track unique violations by worker ID
+        unique_violations = {}  # {worker_id: {violation_type: first_detection_time}}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
-        frame_counter = 0
 
         while processed_frames < total_frames:
             batch_frames = []
@@ -551,7 +557,6 @@ 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
@@ -569,7 +574,7 @@ def process_video(video_data):
                     start_time = time.time()
 
                 boxes = result.boxes
-                detections = []
+                track_inputs = []
                 
                 for box in boxes:
                     cls = int(box.cls)
@@ -583,54 +588,83 @@ def process_video(video_data):
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
-                    detections.append({
+                    track_inputs.append({
                         "bbox": bbox,
-                        "violation": label,
-                        "confidence": conf
+                        "conf": conf,
+                        "cls": cls
                     })
 
-                if not detections:
+                if not track_inputs:
                     continue
                     
-                # 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])
-                    
-                    # 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)
+                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])
+                )
+
+                # 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']
                     
-                    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}"
-                    })
+                    if label is None:
+                        continue
+                        
+                    # Initialize worker if not seen before
+                    if worker_id not in unique_violations:
+                        unique_violations[worker_id] = {}
                     
-                    logger.info(f"Captured snapshot for {violation['violation']} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
+                    # 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] = 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):
@@ -639,15 +673,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
+        # Convert tracked violations to final violation list
         violations = []
-        for worker_id, worker_violations in tracker.violation_history.items():
+        for worker_id, worker_violations in unique_violations.items():
             for label, detection_time in worker_violations.items():
-                violations.append({
+                violation = {
                     "worker_id": worker_id,
                     "violation": label,
                     "timestamp": detection_time
-                })
+                }
+                violations.append(violation)
 
         if not violations:
             logger.info("No violations detected after processing")
@@ -664,15 +699,16 @@ 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) |\n"
-        violation_table += "|-----------|-----------|----------|\n"
+        violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
+        violation_table += "|-----------|-----------|----------|------------|\n"
         
-        for v in sorted(violations, key=lambda x: x.get("timestamp", 0.0)):
+        for v in sorted(violations, key=lambda x: (x.get("worker_id", "Unknown"), 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} |\n"
+            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
 
         # Format snapshots for display
         snapshots_text = ""