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
-import face_recognition
 from collections import defaultdict
 
 # ========================== # Configuration and Setup # ==========================
@@ -29,98 +28,134 @@ os.makedirs('/tmp/Ultralytics', exist_ok=True)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# ========================== # Face Recognition Setup # ==========================
-class FaceTracker:
-    def __init__(self):
-        self.known_faces = {}
-        self.next_face_id = 1
-        self.tolerance = 0.6
-        self.frame_skip = 5  # Process face recognition every N frames
-        
-    def get_face_encoding(self, frame, box):
-        """Extract face encoding from bounding box"""
-        x, y, w, h = box
-        x1, y1 = int(x - w/2), int(y - h/2)
-        x2, y2 = int(x + w/2), int(y + h/2)
-        
-        # Expand the face area slightly
-        expand = 0.2
-        h_expand = int((y2 - y1) * expand)
-        w_expand = int((x2 - x1) * expand)
+# ========================== # Enhanced BYTETracker Implementation # ==========================
+class EnhancedBYTETracker:
+    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
+        self.violation_history = defaultdict(dict)  # Track violations per worker
+        self.last_positions = {}  # Last known positions of workers
+        self.violation_cooldown = {}  # Cooldown periods for violations
+
+    def update(self, dets, scores, cls, frame_idx):
+        tracks = []
+        current_time = frame_idx / self.frame_rate
         
-        y1 = max(0, y1 - h_expand)
-        y2 = min(frame.shape[0], y2 + h_expand)
-        x1 = max(0, x1 - w_expand)
-        x2 = min(frame.shape[1], x2 + w_expand)
+        # 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
+            
+            # 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 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]
+                
+                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
+                }
+                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
         
-        face_frame = frame[y1:y2, x1:x2]
+        # 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]
         
-        if face_frame.size == 0:
-            return None
+        for track_id in stale_ids:
+            del self.tracks[track_id]
+            if track_id in self.last_positions:
+                del self.last_positions[track_id]
             
-        # Convert to RGB (face_recognition uses RGB)
-        rgb_frame = cv2.cvtColor(face_frame, cv2.COLOR_BGR2RGB)
+        return tracks
+    
+    def _calculate_iou(self, box1, box2):
+        """Calculate IOU between two boxes"""
+        x1, y1, w1, h1 = box1
+        x2, y2, w2, h2 = box2
         
-        # Get face encodings
-        encodings = face_recognition.face_encodings(rgb_frame)
-        return encodings[0] if encodings else None
+        # 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)
         
-    def identify_face(self, frame, box):
-        """Identify or register a new face"""
-        encoding = self.get_face_encoding(frame, box)
-        if encoding is None:
-            return None
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
             
-        # Compare with known faces
-        for face_id, known_encoding in self.known_faces.items():
-            matches = face_recognition.compare_faces([known_encoding], encoding, tolerance=self.tolerance)
-            if matches[0]:
-                return face_id
-                
-        # If no match, register new face
-        face_id = f"face_{self.next_face_id}"
-        self.known_faces[face_id] = encoding
-        self.next_face_id += 1
-        return face_id
-
-# ========================== # Position-Based Tracker # ==========================
-class PositionTracker:
-    def __init__(self, distance_threshold=100, cooldown=30):
-        self.workers = {}
-        self.distance_threshold = distance_threshold
-        self.cooldown = cooldown
-        self.next_id = 1
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
         
-    def track(self, position, violation_type, current_time):
-        """Track worker position and return worker ID"""
-        # Check if this is a known worker
-        for worker_id, worker_data in self.workers.items():
-            last_pos = worker_data['position']
-            last_time = worker_data['last_seen']
-            
-            # Calculate distance and time difference
-            distance = np.sqrt((position[0] - last_pos[0])**2 + (position[1] - last_pos[1])**2)
-            time_diff = current_time - last_time
+        box1_area = w1 * h1
+        box2_area = w2 * h2
+        
+        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 close enough and not too much time has passed
-            if distance < self.distance_threshold and time_diff < self.cooldown:
-                # Check if this violation type was already recorded
-                if violation_type not in worker_data['violations']:
-                    worker_data['position'] = position
-                    worker_data['last_seen'] = current_time
-                    worker_data['violations'].add(violation_type)
-                    return worker_id
-                return None  # Violation already recorded
+            # For no_helmet, we only report once per worker
+            if violation_type == "no_helmet":
+                return False
                 
-        # If no match, create new worker
-        worker_id = f"worker_{self.next_id}"
-        self.workers[worker_id] = {
-            'position': position,
-            'last_seen': current_time,
-            'violations': {violation_type}
-        }
-        self.next_id += 1
-        return worker_id
+            # 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
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -139,7 +174,7 @@ CONFIG = {
         "no_harness": (0, 165, 255),       # Orange
         "unsafe_posture": (0, 255, 0),     # Green
         "unsafe_zone": (255, 0, 0),        # Blue
-        "improper_tool_use": (255, 255, 0)  # Cyan
+        "improper_tool_use": (255, 255, 0) # Cyan
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -162,15 +197,16 @@ CONFIG = {
         "unsafe_zone": 0.3,
         "improper_tool_use": 0.3
     },
-    "MIN_VIOLATION_FRAMES": 1,
-    "VIOLATION_COOLDOWN": 30.0,
+    "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,
     "FRAME_SKIP": 2,
     "BATCH_SIZE": 16,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "FACE_RECOGNITION_INTERVAL": 5,  # Process face recognition every N frames
-    "POSITION_TRACKING_THRESHOLD": 100,  # pixels
+    "TRACK_BUFFER": 30,
+    "TRACK_THRESH": 0.3,
+    "MATCH_THRESH": 0.7,
     "SNAPSHOT_QUALITY": 95,
     "MAX_WORKER_DISTANCE": 100
 }
@@ -248,14 +284,20 @@ def calculate_safety_score(violations):
     }
     
     # Count unique violation types per worker
-    worker_violations = defaultdict(set)
+    worker_violations = {}
     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 = sum(penalties.get(v, 0) for violations_set in worker_violations.values() for v in violations_set)
+    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
     
     score = max(0, 100 - total_penalty)
     return score
@@ -288,9 +330,11 @@ def generate_violation_pdf(violations, score):
         y_position -= 0.3 * inch
         
         # Group violations by worker
-        worker_violations = defaultdict(list)
+        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)
@@ -441,7 +485,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"""
+    """Process video to detect safety violations with enhanced tracking"""
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
@@ -463,11 +507,11 @@ 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}")
 
-        # Initialize trackers
-        face_tracker = FaceTracker()
-        position_tracker = PositionTracker(
-            distance_threshold=CONFIG["POSITION_TRACKING_THRESHOLD"],
-            cooldown=CONFIG["VIOLATION_COOLDOWN"]
+        tracker = EnhancedBYTETracker(
+            track_thresh=CONFIG["TRACK_THRESH"],
+            track_buffer=CONFIG["TRACK_BUFFER"],
+            match_thresh=CONFIG["MATCH_THRESH"],
+            frame_rate=fps
         )
 
         violations = []
@@ -475,7 +519,6 @@ def process_video(video_data):
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
-        frame_count = 0
 
         while processed_frames < total_frames:
             batch_frames = []
@@ -500,7 +543,6 @@ def process_video(video_data):
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
                 processed_frames += 1
-                frame_count += 1
 
             if not batch_frames:
                 break
@@ -518,7 +560,7 @@ def process_video(video_data):
                     start_time = time.time()
 
                 boxes = result.boxes
-                detections = []
+                track_inputs = []
                 
                 for box in boxes:
                     cls = int(box.cls)
@@ -532,35 +574,42 @@ def process_video(video_data):
                         continue
 
                     bbox = box.xywh.cpu().numpy()[0]
+                    track_inputs.append({
+                        "bbox": bbox,
+                        "conf": conf,
+                        "cls": cls
+                    })
+
+                if not track_inputs:
+                    continue
                     
-                    # For helmet violations, use face recognition
-                    if label == "no_helmet" and frame_count % CONFIG["FACE_RECOGNITION_INTERVAL"] == 0:
-                        worker_id = face_tracker.identify_face(batch_frames[i], bbox)
-                    else:
-                        # For other violations, use position tracking
-                        position = (bbox[0], bbox[1])
-                        worker_id = position_tracker.track(position, label, current_time)
+                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']
                     
-                    if worker_id is None:
-                        continue  # Skip if this is a duplicate violation
+                    if label is None:
+                        continue
                         
-                    detection = {
-                        "worker_id": worker_id,
-                        "violation": label,
-                        "confidence": round(conf, 2),
-                        "bounding_box": bbox,
-                        "timestamp": current_time
-                    }
-                    detections.append(detection)
-
-                # Process new violations
-                for detection in detections:
-                    # Check if we already have this violation for this worker
-                    existing = next((v for v in violations 
-                                  if v["worker_id"] == detection["worker_id"] 
-                                  and v["violation"] == detection["violation"]), None)
-                    
-                    if not existing:
+                    # 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
@@ -579,7 +628,7 @@ def process_video(video_data):
                         )
                         
                         # Save snapshot with high quality
-                        snapshot_filename = f"violation_{detection['violation']}_worker{detection['worker_id']}_{int(current_time*100)}.jpg"
+                        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(
@@ -589,14 +638,14 @@ def process_video(video_data):
                         )
                         
                         snapshots.append({
-                            "violation": detection["violation"],
-                            "worker_id": detection["worker_id"],
+                            "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 {detection['violation']} violation by worker {detection['worker_id']} at {current_time:.2f}s")
+                        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):