Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import subprocess
 import logging
 import warnings
 import cv2
@@ -15,6 +16,10 @@ 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
 
 # ========================== # Configuration and Setup # ==========================
@@ -23,98 +28,99 @@ os.makedirs('/tmp/Ultralytics', exist_ok=True)
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
-warnings.filterwarnings("ignore")
-
-# ========================== # Position-Based Tracker (No Face Recognition) # ==========================
-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
+
+# ========================== # 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
         
-        self.worker_tracks = {}  # Active worker tracks
-        self.violation_history = defaultdict(dict)  # Track violations per worker
-        self.position_history = defaultdict(list)  # Track positions for all violations
+    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)
         
-        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):
-        current_time = time.time()
-        new_violations = []
+        # Expand the face area slightly
+        expand = 0.2
+        h_expand = int((y2 - y1) * expand)
+        w_expand = int((x2 - x1) * expand)
         
-        for det in detections:
-            bbox = det['bbox']
-            label = det['violation']
-            confidence = det['confidence']
-            
-            worker_id = self._match_by_position(bbox, label)
-            
-            if worker_id is None:
-                worker_id = self.next_id
-                self.next_id += 1
-                
-            if self._is_new_violation(worker_id, label, current_time):
-                violation = {
-                    'worker_id': worker_id,
-                    'violation': label,
-                    'confidence': confidence,
-                    'bbox': bbox,
-                    'timestamp': current_time
-                }
-                new_violations.append(violation)
-                self.violation_history[worker_id][label] = current_time
+        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)
+        
+        face_frame = frame[y1:y2, x1:x2]
+        
+        if face_frame.size == 0:
+            return None
             
-            self.worker_tracks[worker_id] = {
-                'bbox': bbox,
-                'last_seen': current_time,
-                'label': label
-            }
-            self.position_history[worker_id].append((bbox[0], bbox[1]))
+        # Convert to RGB (face_recognition uses RGB)
+        rgb_frame = cv2.cvtColor(face_frame, cv2.COLOR_BGR2RGB)
         
-        self._cleanup_tracks(current_time)
-        return new_violations
-
-    def _match_by_position(self, bbox, label):
-        x, y, w, h = bbox
-        current_pos = (x, y)
+        # Get face encodings
+        encodings = face_recognition.face_encodings(rgb_frame)
+        return encodings[0] if encodings else None
         
-        for worker_id, positions in self.position_history.items():
-            if not positions:
-                continue
-                
-            last_pos = positions[-1]
-            distance = np.sqrt((current_pos[0]-last_pos[0])**2 + (current_pos[1]-last_pos[1])**2)
-            if distance < 100:  # Within 100 pixels
-                return worker_id
-        return None
-
-    def _is_new_violation(self, worker_id, label, current_time):
-        if label not in self.violation_history[worker_id]:
-            return True
+    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
             
-        last_detection = self.violation_history[worker_id][label]
-        cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
-        return (current_time - last_detection) > cooldown
-
-    def _cleanup_tracks(self, current_time):
-        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)
-        ]
+        # 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
         
-        for worker_id in inactive_ids:
-            self.worker_tracks.pop(worker_id, None)
-            self.position_history.pop(worker_id, None)
-            if (current_time - max(self.violation_history[worker_id].values(), default=0)) > 300:
-                self.violation_history.pop(worker_id, None)
+    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
+            
+            # 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
+                
+        # 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
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -129,11 +135,11 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "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)
+        "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",
@@ -156,9 +162,17 @@ CONFIG = {
         "unsafe_zone": 0.3,
         "improper_tool_use": 0.3
     },
+    "MIN_VIOLATION_FRAMES": 1,
+    "VIOLATION_COOLDOWN": 30.0,
+    "WORKER_TRACKING_DURATION": 5.0,
+    "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 2,
-    "BATCH_SIZE": 8,
-    "SNAPSHOT_QUALITY": 90
+    "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
+    "SNAPSHOT_QUALITY": 95,
+    "MAX_WORKER_DISTANCE": 100
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -167,301 +181,515 @@ logger.info(f"Using device: {device}")
 def load_model():
     try:
         if os.path.isfile(CONFIG["MODEL_PATH"]):
-            model = YOLO(CONFIG["MODEL_PATH"]).to(device)
-            logger.info(f"Loaded custom model: {CONFIG['MODEL_PATH']}")
+            model_path = CONFIG["MODEL_PATH"]
+            logger.info(f"Model loaded: {model_path}")
         else:
-            model = YOLO(CONFIG["FALLBACK_MODEL"]).to(device)
-            logger.warning("Using fallback YOLOv8n model")
+            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"Model loading failed: {e}")
+        logger.error(f"Failed to load model: {e}")
         raise
 
 model = load_model()
 
-# ========================== # Core Functions # ==========================
+# ========================== # 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["violation"]
-        confidence = det["confidence"]
-        x, y, w, h = det["bbox"]
-        worker_id = det["worker_id"]
-
-        x1, y1 = int(x - w/2), int(y - h/2)
-        x2, y2 = int(x + w/2), int(y + h/2)
-        color = CONFIG["CLASS_COLORS"][label]
+        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)
+        
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         
+        # Draw thicker rectangle with border
         cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
-        text = f"{CONFIG['DISPLAY_NAMES'][label]} (Worker {worker_id})"
-        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
-        cv2.rectangle(result_frame, (x1, y1-th-10), (x1+tw+10, y1), (0,0,0), -1)
-        cv2.putText(result_frame, text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
-        cv2.putText(result_frame, f"Conf: {confidence:.2f}", (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
+        
+        # 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)
+        
+        # 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 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
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
     }
-    unique_violations = {v["violation"] for v in violations}
-    return max(0, 100 - sum(penalties.get(v, 0) for v in unique_violations))
+    
+    # Count unique violation types per worker
+    worker_violations = defaultdict(set)
+    for v in violations:
+        worker_id = v.get("worker_id", "Unknown")
+        violation_type = v.get("violation", "Unknown")
+        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)
+    
+    score = max(0, 100 - total_penalty)
+    return score
 
 def generate_violation_pdf(violations, score):
+    """Generate a PDF report for the detected violations"""
     try:
-        pdf_buffer = BytesIO()
-        c = canvas.Canvas(pdf_buffer, pagesize=letter)
+        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)
         
-        # Header
+        # Title
         c.setFont("Helvetica-Bold", 16)
-        c.drawString(1*inch, 10*inch, "Worksite Safety Violation Report")
+        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 %H:%M:%S')}")
-        c.drawString(1*inch, 9*inch, f"Safety Score: {score}%")
+        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')}")
         
-        # Violations List
-        y = 8.5*inch
+        # Safety Score
         c.setFont("Helvetica-Bold", 14)
-        c.drawString(1*inch, y, "Detected Violations:")
-        y -= 0.3*inch
-        c.setFont("Helvetica", 10)
+        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
         
+        # Group violations by worker
+        worker_violations = defaultdict(list)
         for v in violations:
-            text = (f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']} "
-                   f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})")
-            if y < 1*inch:
-                c.showPage()
-                y = 10*inch
-            c.drawString(1.2*inch, y, text)
-            y -= 0.2*inch
+            worker_id = v.get("worker_id", "Unknown")
+            worker_violations[worker_id].append(v)
         
-        c.save()
-        pdf_buffer.seek(0)
+        c.setFont("Helvetica", 10)
+        summary_data = {
+            "Total Workers with Violations": len(worker_violations),
+            "Total Violations Found": len(violations),
+            "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
+        }
         
-        # Save to file
-        pdf_filename = f"violation_report_{int(time.time())}.pdf"
-        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
+        for key, value in summary_data.items():
+            c.drawString(1 * inch, y_position, f"{key}: {value}")
+            y_position -= 0.25 * inch
+
+        # Detailed Violations by Worker
+        y_position -= 0.5 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Violations by Worker:")
+        y_position -= 0.3 * inch
+        
+        c.setFont("Helvetica", 10)
+        for worker_id, worker_vios in worker_violations.items():
+            c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
+            y_position -= 0.2 * inch
+            
+            for v in worker_vios:
+                display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
+                time_str = f"{v.get('timestamp', 0.0):.2f}s"
+                conf_str = f"{v.get('confidence', 0.0):.2f}"
+                
+                violation_text = f"  - {display_name} at {time_str} (Confidence: {conf_str})"
+                c.drawString(1.2 * inch, y_position, violation_text)
+                y_position -= 0.2 * inch
+                
+                if y_position < 1 * inch:
+                    c.showPage()
+                    c.setFont("Helvetica", 10)
+                    y_position = 10 * inch
+
+        c.save()
+        pdf_file.seek(0)
+
+        # Save PDF file
         with open(pdf_path, "wb") as f:
-            f.write(pdf_buffer.getvalue())
+            f.write(pdf_file.getvalue())
             
-        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}", pdf_buffer
+        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"PDF generation failed: {e}")
-        return None, None, None
+        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("Salesforce connection established")
+        logger.info("Connected to Salesforce")
+        sf.describe()
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
         raise
 
-def upload_to_salesforce(sf, pdf_file, record_id):
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    """Upload PDF report to Salesforce"""
     try:
-        encoded = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
-        file_data = {
-            "Title": f"Safety_Report_{int(time.time())}",
-            "PathOnClient": "safety_report.pdf",
-            "VersionData": encoded,
-            "FirstPublishLocationId": record_id
+        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
         }
-        result = sf.ContentVersion.create(file_data)
-        return f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{result['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"Salesforce upload failed: {e}")
-        return None
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
 
-def create_salesforce_record(violations, score, pdf_url=None):
+def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+    """Push violation report to Salesforce"""
     try:
         sf = connect_to_salesforce()
-        violations_text = "\n".join(
-            f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} "
-            f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})"
-            for v in violations
-        )
         
+        # 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"
+        
+        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 or ""
+            "PDF_Report_URL__c": pdf_url
         }
         
+        logger.info(f"Creating Salesforce record with data: {record_data}")
+        
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
-            return record["id"], None
-        except:
+            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())}"})
-            return record["id"], "Used Account as fallback"
+            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}")
-        return None, str(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"""
     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("Failed to open video")
+            os.remove(video_path)
+            raise ValueError("Could not open video file")
 
-        fps = cap.get(cv2.CAP_PROP_FPS) or 30
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        tracker = SafetyTracker(frame_rate=fps)
+        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}")
+
+        # Initialize trackers
+        face_tracker = FaceTracker()
+        position_tracker = PositionTracker(
+            distance_threshold=CONFIG["POSITION_TRACKING_THRESHOLD"],
+            cooldown=CONFIG["VIOLATION_COOLDOWN"]
+        )
+
+        violations = []
         snapshots = []
+        start_time = time.time()
+        frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
-        last_update = time.time()
+        frame_count = 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
+                
                 ret, frame = cap.read()
                 if not ret:
                     break
-                batch_frames.append(preprocess_frame(frame))
+                
+                frame = preprocess_frame(frame)
+                
+                # 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
-                if CONFIG["FRAME_SKIP"] > 1:
-                    for _ in range(CONFIG["FRAME_SKIP"]-1):
-                        cap.grab()
-                        processed_frames += 1
+                frame_count += 1
 
             if not batch_frames:
                 break
 
+            # Process batch with YOLO model
             results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
-            for i, result in enumerate(results):
+            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
+                current_time = frame_idx / fps
+                
+                # 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 = []
-                for box in result.boxes:
+                
+                for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls)
-                    if label and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.3):
-                        detections.append({
-                            "bbox": box.xywh.cpu().numpy()[0],
-                            "violation": label,
-                            "confidence": conf
-                        })
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                    
+                    if label is None:
+                        continue
+                        
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        continue
 
-                new_violations = tracker.update(detections)
-                
-                for violation in new_violations:
-                    frame_with_det = draw_detections(batch_frames[i].copy(), [violation])
-                    timestamp = f"Time: {violation['timestamp']:.2f}s"
-                    cv2.putText(frame_with_det, timestamp, (10, 30), 
-                               cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2)
+                    bbox = box.xywh.cpu().numpy()[0]
+                    
+                    # 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)
                     
-                    snap_name = f"violation_{violation['violation']}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
-                    snap_path = os.path.join(CONFIG["OUTPUT_DIR"], snap_name)
-                    cv2.imwrite(snap_path, frame_with_det, 
-                               [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
+                    if worker_id is None:
+                        continue  # Skip if this is a duplicate violation
+                        
+                    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)
                     
-                    snapshots.append({
-                        "violation": violation['violation'],
-                        "worker_id": violation['worker_id'],
-                        "timestamp": violation['timestamp'],
-                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snap_name}"
-                    })
-
-            if time.time() - last_update > 1:
-                progress = (processed_frames / total_frames) * 100
-                yield f"Processing... {progress:.1f}%", "", "", "", ""
-                last_update = time.time()
+                    if not existing:
+                        violations.append(detection)
+                        
+                        # Take snapshot for the new violation
+                        snapshot_frame = batch_frames[i].copy()
+                        snapshot_frame = draw_detections(snapshot_frame, [detection])
+                        
+                        # Add timestamp to snapshot
+                        cv2.putText(
+                            snapshot_frame,
+                            f"Time: {current_time:.2f}s",
+                            (10, 30),
+                            cv2.FONT_HERSHEY_SIMPLEX,
+                            0.7,
+                            (255, 255, 255),
+                            2
+                        )
+                        
+                        # Save snapshot with high quality
+                        snapshot_filename = f"violation_{detection['violation']}_worker{detection['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": detection["violation"],
+                            "worker_id": detection["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")
 
         cap.release()
         if os.path.exists(video_path):
             os.remove(video_path)
-
-        violations = [
-            {"worker_id": wid, "violation": v, "timestamp": t, "confidence": 0}  # Confidence placeholder
-            for wid, violations in tracker.violation_history.items()
-            for v, t in violations.items()
-        ]
+            
+        processing_time = time.time() - start_time
+        logger.info(f"Processing complete in {processing_time:.2f}s")
 
         if not violations:
-            yield "No violations found", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
+            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
 
+        # Calculate safety score
         score = calculate_safety_score(violations)
+        
+        # Generate PDF report
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
-        record_id, sf_error = create_salesforce_record(violations, score, pdf_url)
         
-        if pdf_file and record_id:
-            uploaded_url = upload_to_salesforce(connect_to_salesforce(), pdf_file, record_id)
-            if uploaded_url:
-                pdf_url = uploaded_url
+        # Push report to Salesforce
+        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        violation_table = "| Violation | Worker ID | Time (s) |\n|-----------|-----------|----------|\n"
-        violation_table += "\n".join(
-            f"| {CONFIG['DISPLAY_NAMES'][v['violation']]} | {v['worker_id']} | {v['timestamp']:.2f} |"
-            for v in sorted(violations, key=lambda x: x['timestamp'])
-        )
+        # Format violations table for display
+        violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
+        violation_table += "|-----------|-----------|----------|------------|\n"
+        
+        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} | {confidence:.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"
 
-        snapshots_md = "\n\n".join(
-            f"### {CONFIG['DISPLAY_NAMES'][s['violation']]} - Worker {s['worker_id']} at {s['timestamp']:.2f}s\n\n"
-            f"![Snapshot]({s['snapshot_url']})"
-            for s in snapshots
-        ) if snapshots else "No snapshots captured"
+        if not snapshots_text:
+            snapshots_text = "No snapshots captured."
 
         yield (
             violation_table,
             f"Safety Score: {score}%",
-            snapshots_md,
-            f"Salesforce ID: {record_id or 'N/A'} {sf_error or ''}",
-            pdf_url or "N/A"
+            snapshots_text,
+            f"Salesforce Record ID: {report_id or 'N/A'}",
+            final_pdf_url or "N/A"
         )
 
     except Exception as e:
-        logger.error(f"Video processing failed: {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: {str(e)}", "", "", "", ""
+        yield f"Error processing video: {e}", "", "", "", ""
 
-# ========================== # Gradio Interface # ==========================
-def gradio_interface(video):
-    if not video:
-        return "Upload a video file", "", "", "", ""
-    
+def gradio_interface(video_file):
+    """Gradio interface for the video processing"""
+    if not video_file:
+        return "No file uploaded.", "", "No file uploaded.", "", ""
+        
     try:
-        with open(video, "rb") as f:
+        with open(video_file, "rb") as f:
             video_data = f.read()
-        
-        for output in process_video(video_data):
-            yield output
+
+        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
+            yield status, score, snapshots_text, record_id, details_url
+            
     except Exception as e:
-        logger.error(f"Interface error: {e}")
-        yield f"Error: {str(e)}", "", "", "", ""
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
-with gr.Blocks(title="Safety Compliance Analyzer") as app:
-    gr.Markdown("# Worksite Safety Violation Analyzer")
-    gr.Markdown("Upload site videos to detect safety violations (No Helmet, No Harness, etc.)")
-    
-    with gr.Row():
-        video_input = gr.Video(label="Site Video", sources=["upload"])
-        with gr.Column():
-            violations_out = gr.Markdown(label="Detected Violations")
-            score_out = gr.Textbox(label="Safety Score")
-            snapshots_out = gr.Markdown(label="Violation Snapshots")
-            salesforce_out = gr.Textbox(label="Salesforce Record")
-            pdf_out = gr.Textbox(label="Report PDF URL")
-    
-    video_input.change(
-        gradio_interface,
-        inputs=video_input,
-        outputs=[violations_out, score_out, snapshots_out, salesforce_out, pdf_out]
-    )
+# ========================== # 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__":
-    app.launch(server_port=7860, server_name="0.0.0.0")
+    logger.info("Launching Enhanced Safety Analyzer App...")
+    interface.launch()