Update app.py

app.py

@@ -1,481 +1,739 @@
 import os
 import sys
+import subprocess
 import logging
-import time
+import warnings
 import cv2
-import numpy as np
-import torch
 import gradio as gr
+import torch
+import numpy as np
 from ultralytics import YOLO
+import time
+from simple_salesforce import Salesforce
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.units import inch
+from io import BytesIO
+import base64
+from retrying import retry
+import uuid
+from multiprocessing import Pool, cpu_count
+from functools import partial
+import face_recognition
 from collections import defaultdict
-from typing import List, Dict, Tuple, Optional
-
-# ========================== # Configuration # ==========================
-class Config:
-    # Model settings
-    MODEL_PATH = "yolov8_safety.pt"
-    FALLBACK_MODEL = "yolov8n.pt"
-    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    # Violation settings
-    VIOLATION_LABELS = {
-        0: "no_helmet",
-        1: "no_harness",
-        2: "unsafe_posture",
-        3: "unsafe_zone",
-        4: "improper_tool_use"
-    }
-    
-    # Tracking settings
-    FACE_TRACKING = True  # Enable face tracking for helmet violations
-    TRACK_BUFFER = 30  # Number of frames to keep track of a person
-    MIN_FACE_CONFIDENCE = 0.7  # Minimum confidence for face detection
-    MIN_VIOLATION_CONFIDENCE = 0.5  # Minimum confidence for violation detection
-    
-    # Output settings
-    OUTPUT_DIR = "static/output"
-    SNAPSHOT_QUALITY = 90
-    FRAME_SKIP = 2  # Process every nth frame for efficiency
-    
-    # Violation suppression
-    VIOLATION_COOLDOWN = 5.0  # Seconds before same violation can be reported again
-    POSITION_THRESHOLD = 50  # Pixel distance to consider same location
-    
-    # Display colors
-    COLOR_MAP = {
-        "no_helmet": (0, 0, 255),
-        "no_harness": (0, 165, 255),
-        "unsafe_posture": (0, 255, 0),
-        "unsafe_zone": (255, 0, 0),
-        "improper_tool_use": (255, 255, 0)
-    }
-    
-    # Penalty scores for safety calculation
-    PENALTIES = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 35,
-        "improper_tool_use": 25
-    }
 
-# Initialize logging
+# ========================== # Configuration and Setup # ==========================
+os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
+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 # ==========================
-class FaceTracker:
-    def __init__(self):
-        try:
-            # Try to load face detection model
-            self.face_cascade = cv2.CascadeClassifier(
-                cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-            self.enabled = True
-        except:
-            logger.warning("Face detection model not available - falling back to position tracking")
-            self.enabled = False
-    
-    def detect_faces(self, frame: np.ndarray) -> List[Tuple[int, int, int, int]]:
-        """Detect faces in a frame and return bounding boxes"""
-        if not self.enabled:
-            return []
-            
-        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        faces = self.face_cascade.detectMultiScale(
-            gray,
-            scaleFactor=1.1,
-            minNeighbors=5,
-            minSize=(30, 30)
-        return faces
-
-# ========================== # Violation Tracker # ==========================
-class ViolationTracker:
-    def __init__(self):
-        self.tracked_workers = {}  # {track_id: {last_seen, violations, face_bbox, position}}
-        self.violation_history = defaultdict(list)  # {violation_type: [detection_times]}
+# 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.face_tracker = FaceTracker() if Config.FACE_TRACKING else None
-    
-    def _calculate_iou(self, box1: Tuple, box2: Tuple) -> float:
-        """Calculate Intersection over Union for two bounding boxes"""
-        x1, y1, w1, h1 = box1
-        x2, y2, w2, h2 = box2
-        
-        # Calculate coordinates
-        x_left = max(x1, x2)
-        y_top = max(y1, y2)
-        x_right = min(x1 + w1, x2 + w2)
-        y_bottom = min(y1 + h1, y2 + h2)
-        
-        if x_right < x_left or y_bottom < y_top:
-            return 0.0
-            
-        intersection_area = (x_right - x_left) * (y_bottom - y_top)
-        box1_area = w1 * h1
-        box2_area = w2 * h2
         
-        return intersection_area / (box1_area + box2_area - intersection_area)
-    
-    def _is_same_position(self, pos1: Tuple, pos2: Tuple) -> bool:
-        """Check if two positions are close enough to be considered the same"""
-        x1, y1 = pos1
-        x2, y2 = pos2
-        distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
-        return distance < Config.POSITION_THRESHOLD
-    
-    def _get_face_id(self, frame: np.ndarray, bbox: Tuple) -> Optional[int]:
-        """Try to identify a person by their face"""
-        if not self.face_tracker or not self.face_tracker.enabled:
-            return None
-            
-        x, y, w, h = bbox
-        face_region = frame[max(0, y):y+h, max(0, x):x+w]
-        
-        faces = self.face_tracker.detect_faces(face_region)
-        if len(faces) == 0:
-            return None
-            
-        # Get the largest face in the region
-        largest_face = max(faces, key=lambda f: f[2]*f[3])
-        fx, fy, fw, fh = largest_face
-        
-        # Check if we've seen this face before
-        for track_id, info in self.tracked_workers.items():
-            if 'face_bbox' not in info:
-                continue
-                
-            # Compare with stored face
-            iou = self._calculate_iou((fx, fy, fw, fh), info['face_bbox'])
-            if iou > 0.5:  # Threshold for face matching
-                return track_id
-                
-        return None
-    
-    def update(self, frame: np.ndarray, detections: List[Dict], timestamp: float) -> List[Dict]:
-        """Update tracker with new detections and return unique violations"""
+        # 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
+        }
+
+    def update(self, detections, frame):
+        """Update tracks with new detections and check for violations"""
         current_time = time.time()
-        unique_violations = []
+        active_violations = []
+        new_violations = []
         
-        # First pass - try to match with existing tracks
         for det in detections:
             bbox = det['bbox']
-            violation_type = det['violation']
+            label = det['violation']
             confidence = det['confidence']
-            x, y, w, h = bbox
-            center = (x + w/2, y + h/2)
-            
-            # Try face recognition for helmet violations
-            track_id = None
-            if violation_type == "no_helmet" and self.face_tracker:
-                track_id = self._get_face_id(frame, bbox)
             
-            # If no face match, try position matching
-            if track_id is None:
-                for tid, info in self.tracked_workers.items():
-                    if current_time - info['last_seen'] > Config.TRACK_BUFFER / 30.0:
-                        continue
-                        
-                    if 'position' in info and self._is_same_position(center, info['position']):
-                        track_id = tid
-                        break
+            # 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 still no match, create new track
-            if track_id is None:
-                track_id = self.next_id
+            if worker_id is None:
+                worker_id = self.next_id
                 self.next_id += 1
-                self.tracked_workers[track_id] = {
-                    'last_seen': current_time,
-                    'violations': set(),
-                    'position': center
+                
+            # 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)
                 
-                # Store face if available
-                if violation_type == "no_helmet" and self.face_tracker:
-                    faces = self.face_tracker.detect_faces(frame[y:y+h, x:x+w])
-                    if len(faces) > 0:
-                        self.tracked_workers[track_id]['face_bbox'] = max(
-                            faces, key=lambda f: f[2]*f[3])
+                # 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)
             
-            # Update track info
-            self.tracked_workers[track_id].update({
+            # Keep track of active workers
+            self.worker_tracks[worker_id] = {
+                'bbox': bbox,
                 'last_seen': current_time,
-                'position': center,
-                'bbox': bbox
-            })
+                'label': label
+            }
+        
+        # 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
             
-            # Check if this is a new violation for this track
-            if violation_type not in self.tracked_workers[track_id]['violations']:
-                # Check if this violation type has been seen recently
-                recent_violations = [t for t in self.violation_history[violation_type] 
-                                    if current_time - t < Config.VIOLATION_COOLDOWN]
+        try:
+            # Get face encodings from current detection
+            face_locations = face_recognition.face_locations(face_region)
+            if not face_locations:
+                return None
                 
-                if not recent_violations:
-                    self.tracked_workers[track_id]['violations'].add(violation_type)
-                    self.violation_history[violation_type].append(current_time)
+            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
                     
-                    unique_violations.append({
-                        'track_id': track_id,
-                        'violation': violation_type,
-                        'bbox': bbox,
-                        'confidence': confidence,
-                        'timestamp': timestamp
-                    })
+        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)
         
-        # Clean up old tracks
-        stale_ids = [tid for tid, info in self.tracked_workers.items() 
-                     if current_time - info['last_seen'] > Config.TRACK_BUFFER / 30.0]
-        for tid in stale_ids:
-            del self.tracked_workers[tid]
+        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
             
-        return unique_violations
+        last_detection = self.violation_history[worker_id][label]
+        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
+        
+        return (current_time - last_detection) > cooldown
 
-# ========================== # Video Processor # ==========================
-class VideoProcessor:
-    def __init__(self):
-        self.model = self._load_model()
-        self.tracker = ViolationTracker()
+    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)]
         
-    def _load_model(self):
-        """Load YOLOv8 model with fallback"""
+        if face_region.size == 0:
+            return
+            
         try:
-            if os.path.exists(Config.MODEL_PATH):
-                model = YOLO(Config.MODEL_PATH).to(Config.DEVICE)
-                logger.info(f"Loaded custom model from {Config.MODEL_PATH}")
-            else:
-                model = YOLO(Config.FALLBACK_MODEL).to(Config.DEVICE)
-                logger.warning("Using fallback YOLOv8n model")
-            return model
+            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.error(f"Failed to load model: {e}")
-            raise
+            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)
             
-    def _draw_detection(self, frame: np.ndarray, detection: Dict) -> np.ndarray:
-        """Draw a single detection on the frame"""
-        x, y, w, h = detection['bbox']
-        label = detection['violation']
-        confidence = detection['confidence']
-        track_id = detection.get('track_id', 0)
-        
-        color = Config.COLOR_MAP.get(label, (0, 0, 255))
-        
-        # Draw bounding box
-        cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
+            # 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 = {
+    "MODEL_PATH": "yolov8_safety.pt",
+    "FALLBACK_MODEL": "yolov8n.pt",
+    "OUTPUT_DIR": "static/output",
+    "VIOLATION_LABELS": {
+        0: "no_helmet",
+        1: "no_harness",
+        2: "unsafe_posture",
+        3: "unsafe_zone",
+        4: "improper_tool_use"
+    },
+    "CLASS_COLORS": {
+        "no_helmet": (0, 0, 255),          # Red
+        "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
+    },
+    "DISPLAY_NAMES": {
+        "no_helmet": "No Helmet Violation",
+        "no_harness": "No Harness Violation",
+        "unsafe_posture": "Unsafe Posture",
+        "unsafe_zone": "Unsafe Zone Entry",
+        "improper_tool_use": "Improper Tool Use"
+    },
+    "SF_CREDENTIALS": {
+        "username": "prashanth1ai@safety.com",
+        "password": "SaiPrash461",
+        "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
+        "domain": "login"
+    },
+    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
+    "CONFIDENCE_THRESHOLDS": {
+        "no_helmet": 0.5,
+        "no_harness": 0.3,
+        "unsafe_posture": 0.3,
+        "unsafe_zone": 0.3,
+        "improper_tool_use": 0.3
+    },
+    "MIN_VIOLATION_FRAMES": 1,
+    "FRAME_SKIP": 2,
+    "BATCH_SIZE": 16,
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),
+    "SNAPSHOT_QUALITY": 95,
+    "FACE_RECOGNITION_INTERVAL": 5  # Process face recognition every 5 frames
+}
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logger.info(f"Using device: {device}")
+
+def load_model():
+    try:
+        if os.path.isfile(CONFIG["MODEL_PATH"]):
+            model_path = CONFIG["MODEL_PATH"]
+            logger.info(f"Model loaded: {model_path}")
+        else:
+            model_path = CONFIG["FALLBACK_MODEL"]
+            logger.warning("Using fallback model. Train yolov8_safety.pt for best results.")
+            if not os.path.isfile(model_path):
+                logger.info(f"Downloading fallback model: {model_path}")
+                torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
+                
+        model = YOLO(model_path).to(device)
+        logger.info(f"Model classes: {model.names}")
+        return model
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+model = load_model()
+
+# ========================== # Helper Functions # ==========================
+def preprocess_frame(frame):
+    """Apply basic preprocessing to enhance detection"""
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
+    return frame
+
+def draw_detections(frame, detections):
+    """Draw bounding boxes and labels on detection frame with improved visibility"""
+    result_frame = frame.copy()
+    
+    for det in detections:
+        label = det.get("violation", "Unknown")
+        confidence = det.get("confidence", 0.0)
+        x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
+        worker_id = det.get("worker_id", "Unknown")
+
+        x1 = int(x - w/2)
+        y1 = int(y - h/2)
+        x2 = int(x + w/2)
+        y2 = int(y + h/2)
         
-        # Draw label background
-        label_text = f"{label} (ID: {track_id})"
-        (text_width, text_height), _ = cv2.getTextSize(
-            label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         
-        cv2.rectangle(frame, (x, y - text_height - 10), 
-                     (x + text_width + 10, y), color, -1)
+        # Draw thicker rectangle with border
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
         
-        # Draw label text
-        cv2.putText(frame, label_text, (x + 5, y - 5), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+        # Add black background behind text
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
+        text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
+        cv2.rectangle(result_frame, (x1, y1-text_size[1]-10), (x1+text_size[0]+10, y1), (0, 0, 0), -1)
+        cv2.putText(result_frame, display_text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
         
-        # Draw confidence
-        conf_text = f"{confidence:.2f}"
-        cv2.putText(frame, conf_text, (x + 5, y + h + 15), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+        # Add confidence score
+        conf_text = f"Conf: {confidence:.2f}"
+        cv2.putText(result_frame, conf_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
         
-        return frame
+    return result_frame
+
+def calculate_safety_score(violations):
+    """Calculate safety score based on detected violations"""
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
     
-    def _process_frame(self, frame: np.ndarray, frame_idx: int, fps: float) -> Tuple[List[Dict], np.ndarray]:
-        """Process a single video frame"""
-        timestamp = frame_idx / fps
-        
-        # Run YOLO detection
-        results = self.model(frame, verbose=False)
-        detections = []
-        
-        for result in results:
-            for box in result.boxes:
-                cls_id = int(box.cls)
-                conf = float(box.conf)
-                label = Config.VIOLATION_LABELS.get(cls_id)
-                
-                if label is None or conf < Config.MIN_VIOLATION_CONFIDENCE:
-                    continue
-                
-                # Convert bounding box to (x, y, w, h) format
-                x1, y1, x2, y2 = map(int, box.xyxy[0])
-                detections.append({
-                    'bbox': (x1, y1, x2 - x1, y2 - y1),
-                    'violation': label,
-                    'confidence': conf
-                })
+    # Count unique violation types
+    unique_violations = set()
+    for v in violations:
+        violation_type = v.get("violation", "Unknown")
+        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
+
+def generate_violation_pdf(violations, score):
+    """Generate a PDF report for the detected violations"""
+    try:
+        pdf_filename = f"violations_{int(time.time())}.pdf"
+        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
+        pdf_file = BytesIO()
+        c = canvas.Canvas(pdf_file, pagesize=letter)
+        
+        # Title
+        c.setFont("Helvetica-Bold", 16)
+        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
+        
+        # Basic Information
+        c.setFont("Helvetica", 12)
+        c.drawString(1 * inch, 9.5 * inch, f"Date: {time.strftime('%Y-%m-%d')}")
+        c.drawString(1 * inch, 9.2 * inch, f"Time: {time.strftime('%H:%M:%S')}")
+        
+        # Safety Score
+        c.setFont("Helvetica-Bold", 14)
+        c.drawString(1 * inch, 8.7 * inch, f"Safety Compliance Score: {score}%")
+
+        # Violation Summary
+        y_position = 8.2 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Summary:")
+        y_position -= 0.3 * inch
+        
+        c.setFont("Helvetica", 10)
+        summary_data = {
+            "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")
+        }
+        
+        for key, value in summary_data.items():
+            c.drawString(1 * inch, y_position, f"{key}: {value}")
+            y_position -= 0.25 * inch
+
+        # Detailed Violations
+        y_position -= 0.5 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Violation Details:")
+        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)
+            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)
+
+        # Save PDF file
+        with open(pdf_path, "wb") as f:
+            f.write(pdf_file.getvalue())
+            
+        public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
+        logger.info(f"PDF generated: {public_url}")
+        return pdf_path, public_url, pdf_file
+    except Exception as e:
+        logger.error(f"Error generating PDF: {e}")
+        return "", "", None
+
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
+def connect_to_salesforce():
+    """Connect to Salesforce with retry logic"""
+    try:
+        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        logger.info("Connected to Salesforce")
+        sf.describe()
+        return sf
+    except Exception as e:
+        logger.error(f"Salesforce connection failed: {e}")
+        raise
+
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    """Upload PDF report to Salesforce"""
+    try:
+        if not pdf_file:
+            logger.error("No PDF file provided for upload")
+            return ""
+            
+        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        content_version_data = {
+            "Title": f"Safety_Violation_Report_{int(time.time())}",
+            "PathOnClient": f"safety_violation_{int(time.time())}.pdf",
+            "VersionData": encoded_pdf,
+            "FirstPublishLocationId": report_id
+        }
+        content_version = sf.ContentVersion.create(content_version_data)
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        
+        if not result['records']:
+            logger.error("Failed to retrieve ContentVersion")
+            return ""
+            
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+        logger.info(f"PDF uploaded to Salesforce: {file_url}")
+        return file_url
+    except Exception as e:
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
+
+def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+    """Push violation report to Salesforce"""
+    try:
+        sf = connect_to_salesforce()
+        
+        # Format violations for Salesforce
+        violations_text = ""
+        for v in violations:
+            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)
+            
+            violations_text += f"Worker {worker_id}: {display_name} at {timestamp:.2f}s (Conf: {confidence:.2f})\n"
         
-        # Update tracker and get unique violations
-        unique_violations = self.tracker.update(frame, detections, timestamp)
+        if not violations_text:
+            violations_text = "No violations detected."
+            
+        pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
+
+        record_data = {
+            "Compliance_Score__c": score,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
+            "Status__c": "Pending",
+            "PDF_Report_URL__c": pdf_url
+        }
         
-        # Draw detections on frame
-        output_frame = frame.copy()
-        for violation in unique_violations:
-            output_frame = self._draw_detection(output_frame, violation)
+        logger.info(f"Creating Salesforce record with data: {record_data}")
         
-        return unique_violations, output_frame
-    
-    def process_video(self, video_path: str) -> Tuple[List[Dict], List[Tuple[float, np.ndarray]]]:
-        """Process entire video file"""
+        try:
+            record = sf.Safety_Video_Report__c.create(record_data)
+            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
+        except Exception as e:
+            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
+            record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
+            logger.warning(f"Fell back to Account record: {record['id']}")
+            
+        record_id = record["id"]
+
+        if pdf_file:
+            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
+            if uploaded_url:
+                try:
+                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    logger.info(f"Updated record {record_id} with PDF URL: {uploaded_url}")
+                except Exception as e:
+                    logger.error(f"Failed to update Safety_Video_Report__c: {e}")
+                    sf.Account.update(record_id, {"Description": uploaded_url})
+                    logger.info(f"Updated Account record {record_id} with PDF URL")
+                pdf_url = uploaded_url
+
+        return record_id, pdf_url
+    except Exception as e:
+        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
+        return None, ""
+
+def process_video(video_data):
+    """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']}")
+
+        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
+        with open(video_path, "wb") as f:
+            f.write(video_data)
+        logger.info(f"Video saved: {video_path}")
+
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            raise ValueError(f"Could not open video: {video_path}")
-        
-        fps = cap.get(cv2.CAP_PROP_FPS)
+            os.remove(video_path)
+            raise ValueError("Could not open video file")
+
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        violations = []
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30
+        duration = total_frames / fps
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        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)
         snapshots = []
-        
-        logger.info(f"Processing video: {video_path} ({total_frames} frames, {fps:.1f} FPS)")
-        
-        frame_idx = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
+        start_time = time.time()
+        frame_skip = CONFIG["FRAME_SKIP"]
+        processed_frames = 0
+        frame_counter = 0
+
+        while processed_frames < total_frames:
+            batch_frames = []
+            batch_indices = []
+            
+            for _ in range(CONFIG["BATCH_SIZE"]):
+                frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+                if frame_idx >= total_frames:
+                    break
                 
-            # Skip frames for efficiency
-            if frame_idx % Config.FRAME_SKIP != 0:
-                frame_idx += 1
-                continue
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                frame = preprocess_frame(frame)
                 
-            # Process frame
-            frame_violations, output_frame = self._process_frame(frame, frame_idx, fps)
+                # Skip frames if needed
+                for _ in range(frame_skip - 1):
+                    if not cap.grab():
+                        break
+                
+                batch_frames.append(frame)
+                batch_indices.append(frame_idx)
+                processed_frames += 1
+                frame_counter += 1
+
+            if not batch_frames:
+                break
+
+            # Process batch with YOLO model
+            results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
-            # Save snapshots for new violations
-            for violation in frame_violations:
-                violations.append(violation)
+            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
+                current_time = frame_idx / fps
                 
-                # Save snapshot
-                timestamp = frame_idx / fps
-                snapshot_path = os.path.join(
-                    Config.OUTPUT_DIR,
-                    f"violation_{violation['violation']}_id{violation['track_id']}_{int(timestamp*100)}.jpg")
+                # Update progress every second
+                if time.time() - start_time > 1.0:
+                    progress = (processed_frames / total_frames) * 100
+                    yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames})", "", "", "", ""
+                    start_time = time.time()
+
+                boxes = result.boxes
+                detections = []
                 
-                cv2.imwrite(snapshot_path, output_frame, 
-                           [cv2.IMWRITE_JPEG_QUALITY, Config.SNAPSHOT_QUALITY])
+                for box in boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                    
+                    if label is None:
+                        continue
+                        
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        continue
+
+                    bbox = box.xywh.cpu().numpy()[0]
+                    detections.append({
+                        "bbox": bbox,
+                        "violation": label,
+                        "confidence": conf
+                    })
+
+                if not detections:
+                    continue
+                    
+                # Update tracker with new detections
+                new_violations = tracker.update(detections, batch_frames[i])
                 
-                snapshots.append((timestamp, snapshot_path))
-                logger.info(f"Detected {violation['violation']} at {timestamp:.2f}s (ID: {violation['track_id']})")
-            
-            frame_idx += 1
-            
-            # Yield progress
-            if frame_idx % 10 == 0:
-                progress = frame_idx / total_frames
-                yield progress, violations, snapshots
-        
+                # 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)
+                    
+                    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()
-        yield 1.0, violations, snapshots  # Final yield
-    
-    def calculate_safety_score(self, violations: List[Dict]) -> int:
-        """Calculate safety compliance score (0-100)"""
-        penalty = 0
-        violation_types = set()
-        
-        for violation in violations:
-            violation_types.add(violation['violation'])
-        
-        for v_type in violation_types:
-            penalty += Config.PENALTIES.get(v_type, 0)
-        
-        return max(0, 100 - penalty)
+        if os.path.exists(video_path):
+            os.remove(video_path)
+            
+        processing_time = time.time() - start_time
+        logger.info(f"Processing complete in {processing_time:.2f}s")
 
-# ========================== # Gradio Interface # ==========================
-def setup_output_dir():
-    """Ensure output directory exists"""
-    os.makedirs(Config.OUTPUT_DIR, exist_ok=True)
-    logger.info(f"Output directory: {Config.OUTPUT_DIR}")
-
-def format_violation_table(violations: List[Dict]) -> str:
-    """Format violations as markdown table"""
-    if not violations:
-        return "No violations detected"
-    
-    table = "| Violation | Worker ID | Time (s) | Confidence |\n"
-    table += "|-----------|-----------|----------|------------|\n"
-    
-    for v in sorted(violations, key=lambda x: x['timestamp']):
-        table += f"| {v['violation']} | {v['track_id']} | {v['timestamp']:.2f} | {v['confidence']:.2f} |\n"
-    
-    return table
+        # 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
+                })
 
-def format_snapshots(snapshots: List[Tuple[float, str]]) -> str:
-    """Format snapshots as markdown"""
-    if not snapshots:
-        return "No violations captured"
-    
-    markdown = ""
-    for timestamp, path in snapshots:
-        filename = os.path.basename(path)
-        markdown += f"### Violation at {timestamp:.2f}s\n\n"
-        markdown += f"![Violation](file/{path})\n\n"
-    
-    return markdown
+        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"
+            return
 
-def process_video_wrapper(video_file: str):
-    """Wrapper for Gradio interface"""
-    setup_output_dir()
-    
-    # Create temporary video path
-    temp_video_path = os.path.join(Config.OUTPUT_DIR, f"temp_{int(time.time())}.mp4")
-    with open(temp_video_path, "wb") as f:
-        f.write(open(video_file, "rb").read())
-    
-    processor = VideoProcessor()
-    
-    try:
-        for progress, violations, snapshots in processor.process_video(temp_video_path):
-            yield (
-                f"Processing... {progress*100:.1f}% complete",
-                "",
-                "",
-                ""
-            )
+        # Calculate safety score
+        score = calculate_safety_score(violations)
         
-        # Calculate final results
-        score = processor.calculate_safety_score(violations)
-        violation_table = format_violation_table(violations)
-        snapshots_md = format_snapshots(snapshots)
+        # Generate PDF report
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
         
+        # Push report to Salesforce
+        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"
+        
+        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)
+            
+            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} |\n"
+
+        # Format snapshots for display
+        snapshots_text = ""
+        for s in snapshots:
+            display_name = CONFIG["DISPLAY_NAMES"].get(s["violation"], "Unknown")
+            worker_id = s.get("worker_id", "Unknown")
+            timestamp = s.get("timestamp", 0.0)
+            
+            snapshots_text += f"### {display_name} - Worker {worker_id} at {timestamp:.2f}s\n\n"
+            snapshots_text += f"![Violation]({s['snapshot_url']})\n\n"
+
+        if not snapshots_text:
+            snapshots_text = "No snapshots captured."
+
         yield (
-            "Processing complete",
-            f"Safety Score: {score}%",
             violation_table,
-            snapshots_md
+            f"Safety Score: {score}%",
+            snapshots_text,
+            f"Salesforce Record ID: {report_id or 'N/A'}",
+            final_pdf_url or "N/A"
         )
+
+    except Exception as e:
+        logger.error(f"Error processing video: {e}", exc_info=True)
+        if 'video_path' in locals() and os.path.exists(video_path):
+            os.remove(video_path)
+        yield f"Error processing video: {e}", "", "", "", ""
+
+def gradio_interface(video_file):
+    """Gradio interface for the video processing"""
+    if not video_file:
+        return "No file uploaded.", "", "No file uploaded.", "", ""
         
-    finally:
-        if os.path.exists(temp_video_path):
-            os.remove(temp_video_path)
-
-# ========================== # Main Application # ==========================
-def create_interface():
-    """Create Gradio interface"""
-    with gr.Blocks(title="Safety Compliance Analyzer") as interface:
-        gr.Markdown("# 🚧 Safety Compliance Video Analyzer")
-        gr.Markdown("Upload site videos to detect safety violations (No Helmet, No Harness, etc.)")
-        
-        with gr.Row():
-            with gr.Column():
-                video_input = gr.Video(label="Upload Site Video")
-                submit_btn = gr.Button("Analyze Video", variant="primary")
+    try:
+        with open(video_file, "rb") as f:
+            video_data = f.read()
+
+        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
+            yield status, score, snapshots_text, record_id, details_url
             
-            with gr.Column():
-                progress_out = gr.Textbox(label="Status")
-                score_out = gr.Textbox(label="Safety Score")
-                violations_out = gr.Markdown(label="Detected Violations")
-                snapshots_out = gr.Markdown(label="Violation Snapshots")
-        
-        submit_btn.click(
-            fn=process_video_wrapper,
-            inputs=video_input,
-            outputs=[progress_out, score_out, violations_out, snapshots_out]
-        )
-    
-    return interface
+    except Exception as e:
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+
+# ========================== # Gradio Interface # ==========================
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Video(label="Upload Site Video"),
+    outputs=[
+        gr.Markdown(label="Detected Safety Violations"),
+        gr.Textbox(label="Compliance Score"),
+        gr.Markdown(label="Snapshots"),
+        gr.Textbox(label="Salesforce Record ID"),
+        gr.Textbox(label="Violation Details URL")
+    ],
+    title="Worksite Safety Violation Analyzer",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Each unique violation is detected only once per worker.",
+    allow_flagging="never"
+)
 
 if __name__ == "__main__":
-    setup_output_dir()
-    interface = create_interface()
+    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()