Update app.py

app.py

@@ -1,8 +1,4 @@
 import os
-import sys
-import subprocess
-import logging
-import warnings
 import cv2
 import gradio as gr
 import torch
@@ -15,162 +11,15 @@ from reportlab.pdfgen import canvas
 from reportlab.lib.units import inch
 from io import BytesIO
 import base64
+import logging
 from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
 
-# ========================== # 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__)
-
-# ========================== # ByteTrack Implementation # ==========================
-class BYTETracker:
-    def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
-        self.track_thresh = track_thresh
-        self.track_buffer = track_buffer
-        self.match_thresh = match_thresh
-        self.frame_rate = frame_rate
-        self.next_id = 1
-        self.tracks = {}  # Store active tracks
-        self.worker_history = {}  # Track worker positions over time
-        self.last_positions = {}  # Last known positions of workers
-
-    def update(self, dets, scores, cls):
-        tracks = []
-        current_time = time.time()
-        
-        # 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 history
-                if best_track_id not in self.worker_history:
-                    self.worker_history[best_track_id] = []
-                self.worker_history[best_track_id].append([x, y])
-                self.last_positions[best_track_id] = [x, y]
-                
-                tracks.append({
-                    'id': best_track_id,
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl
-                })
-            else:
-                # Create new track
-                # Check if this detection might be the same worker from a different angle
-                same_worker = False
-                for worker_id, last_pos in self.last_positions.items():
-                    if self._is_same_worker([x, y], last_pos):
-                        self.tracks[worker_id] = {
-                            'bbox': [x, y, w, h],
-                            'score': score,
-                            'cls': cl,
-                            'last_seen': current_time
-                        }
-                        tracks.append({
-                            'id': worker_id,
-                            'bbox': [x, y, w, h],
-                            'score': score,
-                            'cls': cl
-                        })
-                        same_worker = True
-                        break
-                
-                if not same_worker:
-                    self.tracks[self.next_id] = {
-                        'bbox': [x, y, w, h],
-                        'score': score,
-                        'cls': cl,
-                        'last_seen': current_time
-                    }
-                    self.worker_history[self.next_id] = [[x, y]]
-                    self.last_positions[self.next_id] = [x, y]
-                    tracks.append({
-                        'id': self.next_id,
-                        'bbox': [x, y, w, h],
-                        'score': score,
-                        'cls': cl
-                    })
-                    self.next_id += 1
-        
-        # Clean up old tracks
-        current_time = time.time()
-        stale_ids = []
-        for track_id, track_info in self.tracks.items():
-            if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
-                stale_ids.append(track_id)
-        
-        for track_id in stale_ids:
-            del self.tracks[track_id]
-            if track_id in self.worker_history:
-                del self.worker_history[track_id]
-            if track_id in self.last_positions:
-                del self.last_positions[track_id]
-            
-        return tracks
-    
-    def _calculate_iou(self, box1, box2):
-        """Calculate IOU between two boxes"""
-        x1, y1, w1, h1 = box1
-        x2, y2, w2, h2 = box2
-        
-        # Calculate intersection coordinates
-        x_left = max(x1 - w1/2, x2 - w2/2)
-        y_top = max(y1 - h1/2, y2 - h2/2)
-        x_right = min(x1 + w1/2, x2 + w2/2)
-        y_bottom = min(y1 + h1/2, y2 + h2/2)
-        
-        if 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
-        
-        iou = intersection_area / (box1_area + box2_area - intersection_area)
-        return iou
-    
-    def _is_same_worker(self, pos1, pos2, threshold=100):
-        """Check if two positions likely belong to the same worker"""
-        x1, y1 = pos1
-        x2, y2 = pos2
-        distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
-        return distance < threshold
-
-# ========================== # Optimized Configuration # ==========================
+# ==========================
+# Optimized Configuration
+# ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
@@ -183,11 +32,11 @@ CONFIG = {
         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
+        "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)  # Yellow
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -204,26 +53,26 @@ CONFIG = {
     },
     "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
+        "no_helmet": 0.75,  # Increased for stricter helmet detection
+        "no_harness": 0.4,
+        "unsafe_posture": 0.4,
+        "unsafe_zone": 0.4,
+        "improper_tool_use": 0.4
     },
-    "MIN_VIOLATION_FRAMES": 1,
-    "VIOLATION_COOLDOWN": 30.0,  # Increased cooldown period
-    "WORKER_TRACKING_DURATION": 5.0,
-    "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,  # Skip more frames for faster processing
-    "BATCH_SIZE": 16,
-    "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 30,
-    "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.7,
-    "SNAPSHOT_QUALITY": 95,  # Higher quality for better visibility
-    "MAX_WORKER_DISTANCE": 100  # Maximum pixel distance to consider same worker
+    "MIN_VIOLATION_FRAMES": 3,
+    "WORKER_TRACKING_DURATION": 3.0,
+    "MAX_PROCESSING_TIME": 60,  # 1 minute limit
+    "FRAME_SKIP": 2,  # Process every 2nd frame for speed
+    "BATCH_SIZE": 16,  # Frames per batch
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1)  # Use all CPU cores except one
 }
 
+# Setup logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
@@ -238,9 +87,7 @@ def load_model():
             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}")
@@ -248,151 +95,118 @@ def load_model():
 
 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
-
+# ==========================
+# Optimized Helper Functions
+# ==========================
 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)
         
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
         
-        # Draw thicker rectangle with border
-        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
-        
-        # 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
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)}: {confidence:.2f}"
+        cv2.putText(frame, display_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    return 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 calculate_iou(box1, box2):
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
     
-    # Count unique violation types per worker
-    worker_violations = {}
-    for v in violations:
-        worker_id = v.get("worker_id", "Unknown")
-        violation_type = v.get("violation", "Unknown")
-        
-        if worker_id not in worker_violations:
-            worker_violations[worker_id] = set()
-        worker_violations[worker_id].add(violation_type)
+    x_left = max(x1 - w1/2, x2 - w2/2)
+    y_top = max(y1 - h1/2, y2 - h2/2)
+    x_right = min(x1 + w1/2, x2 + w2/2)
+    y_bottom = min(y1 + h1/2, y2 + h2/2)
+    
+    if 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
+    union_area = box1_area + box2_area - intersection_area
+    
+    return intersection_area / union_area
+
+def process_frame_batch(frame_batch, frame_indices, fps):
+    batch_results = []
+    results = model(frame_batch, device=device, conf=0.1, verbose=False)
     
-    # Calculate total penalty
-    total_penalty = 0
-    for worker_violations_set in worker_violations.values():
-        worker_penalty = sum(penalties.get(v, 0) for v in worker_violations_set)
-        total_penalty += worker_penalty
+    for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
+        current_time = frame_idx / fps
+        detections = []
+        
+        boxes = result.boxes
+        for box in boxes:
+            cls = int(box.cls)
+            conf = float(box.conf)
+            label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+            
+            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                continue
+
+            bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+            detections.append({
+                "frame": frame_idx,
+                "violation": label,
+                "confidence": round(conf, 2),
+                "bounding_box": bbox,
+                "timestamp": current_time
+            })
+        
+        batch_results.append((frame_idx, detections))
     
-    score = max(0, 100 - total_penalty)
-    return score
+    return batch_results
 
 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
-        
-        # Group violations by worker
-        worker_violations = {}
-        for v in violations:
-            worker_id = v.get("worker_id", "Unknown")
-            if worker_id not in worker_violations:
-                worker_violations[worker_id] = []
-            worker_violations[worker_id].append(v)
-        
+        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
         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")
+
+        y_position = 9.5 * inch
+        report_data = {
+            "Compliance Score": f"{score}%",
+            "Violations Found": len(violations),
+            "Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
-        
-        for key, value in summary_data.items():
+        for key, value in report_data.items():
             c.drawString(1 * inch, y_position, f"{key}: {value}")
-            y_position -= 0.25 * inch
+            y_position -= 0.3 * 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:
+        c.drawString(1 * inch, y_position, "Violation Details:")
+        y_position -= 0.3 * inch
+        if not violations:
+            c.drawString(1 * inch, y_position, "No violations detected.")
+        else:
+            for v in violations:
                 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
-                
+                text = f"{display_name} at {v.get('timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f})"
+                c.drawString(1 * inch, y_position, text)
+                y_position -= 0.3 * inch
                 if y_position < 1 * inch:
                     c.showPage()
                     c.setFont("Helvetica", 10)
                     y_position = 10 * inch
 
+        c.showPage()
         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
@@ -400,9 +214,23 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
+def calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
+    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
+    score = 100 - total_penalty
+    return max(score, 0)
+
+# ==========================
+# Salesforce Integration
+# ==========================
 @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")
@@ -413,12 +241,10 @@ def connect_to_salesforce():
         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())}",
@@ -428,11 +254,9 @@ def upload_pdf_to_salesforce(sf, pdf_file, 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
@@ -441,23 +265,12 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         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"
-        
-        if not violations_text:
-            violations_text = "No violations detected."
-            
+        violations_text = "\n".join(
+            f"{CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} at {v.get('timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f})"
+            for v in violations
+        ) or "No violations detected."
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
 
         record_data = {
@@ -467,9 +280,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             "Status__c": "Pending",
             "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)
             logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
@@ -477,7 +288,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             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:
@@ -497,47 +307,46 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
         return None, ""
 
+# ==========================
+# Fast Video Processing
+# ==========================
 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']}")
-
+        # Create temp video file
         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}")
 
+        # Open video file
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            os.remove(video_path)
             raise ValueError("Could not open video file")
 
+        # Get video properties
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = cap.get(cv2.CAP_PROP_FPS) or 30
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 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 = BYTETracker(
-            track_thresh=CONFIG["TRACK_THRESH"],
-            track_buffer=CONFIG["TRACK_BUFFER"],
-            match_thresh=CONFIG["MATCH_THRESH"],
-            frame_rate=fps
-        )
+        logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Track unique violations by worker ID
-        unique_violations = {}  # {worker_id: {violation_type: first_detection_time}}
+        workers = []
+        violations = []
+        helmet_violations = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
-        processed_frames = 0
 
-        while processed_frames < total_frames:
+        # Process frames in batches
+        while True:
             batch_frames = []
             batch_indices = []
             
+            # Collect frames for this batch
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
                 if frame_idx >= total_frames:
@@ -547,8 +356,6 @@ def process_video(video_data):
                 if not ret:
                     break
                 
-                frame = preprocess_frame(frame)
-                
                 # Skip frames if needed
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
@@ -556,172 +363,127 @@ def process_video(video_data):
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
-                processed_frames += 1
 
+            # Break if no more frames
             if not batch_frames:
                 break
 
-            # Process batch with YOLO model
+            # Run batch detection
             results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
+            # Process results for each frame in batch
             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})", "", "", "", ""
+                # Update progress periodically
+                if time.time() - start_time > 1.0:  # Update every second
+                    progress = (frame_idx / total_frames) * 100
+                    yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
                     start_time = time.time()
 
+                # Process detections in this frame
                 boxes = result.boxes
-                track_inputs = []
-                
                 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):
+                    if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
                         continue
 
-                    bbox = box.xywh.cpu().numpy()[0]
-                    track_inputs.append({
-                        "bbox": bbox,
-                        "conf": conf,
-                        "cls": cls
-                    })
+                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    detection = {
+                        "frame": frame_idx,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    }
 
-                if not track_inputs:
-                    continue
-                    
-                tracked_objects = tracker.update(
-                    np.array([t["bbox"] for t in track_inputs]),
-                    np.array([t["conf"] for t in track_inputs]),
-                    np.array([t["cls"] for t in track_inputs])
-                )
-
-                # 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 label is None:
-                        continue
-                        
-                    # Initialize worker if not seen before
-                    if worker_id not in unique_violations:
-                        unique_violations[worker_id] = {}
-                    
-                    # Check if this violation type has been recorded for this worker
-                    if label not in unique_violations[worker_id]:
-                        # This is a new violation type for this worker
-                        unique_violations[worker_id][label] = current_time
-                        
-                        # Create detection object
-                        detection = {
-                            "worker_id": worker_id,
-                            "violation": label,
-                            "confidence": round(conf, 2),
-                            "bounding_box": bbox,
-                            "timestamp": current_time
-                        }
-                        
-                        # Take snapshot for the new violation
-                        snapshot_frame = batch_frames[i].copy()
-                        snapshot_frame = draw_detections(snapshot_frame, [detection])
-                        
-                        # Add timestamp to snapshot
-                        cv2.putText(
-                            snapshot_frame,
-                            f"Time: {current_time:.2f}s",
-                            (10, 30),
-                            cv2.FONT_HERSHEY_SIMPLEX,
-                            0.7,
-                            (255, 255, 255),
-                            2
-                        )
-                        
-                        # Save snapshot with high quality
-                        snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        
-                        cv2.imwrite(
-                            snapshot_path,
-                            snapshot_frame,
-                            [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
-                        )
-                        
-                        snapshots.append({
-                            "violation": label,
-                            "worker_id": worker_id,
-                            "timestamp": current_time,
-                            "snapshot_path": snapshot_path,
-                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    # Worker tracking
+                    worker_id = None
+                    max_iou = 0
+                    for idx, worker in enumerate(workers):
+                        iou = calculate_iou(bbox, worker["bbox"])
+                        if iou > max_iou and iou > 0.4:  # IOU threshold
+                            max_iou = iou
+                            worker_id = worker["id"]
+                            workers[idx]["bbox"] = bbox
+                            workers[idx]["last_seen"] = current_time
+
+                    if worker_id is None:
+                        worker_id = len(workers) + 1
+                        workers.append({
+                            "id": worker_id,
+                            "bbox": bbox,
+                            "first_seen": current_time,
+                            "last_seen": current_time
                         })
-                        
-                        logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
+
+                    detection["worker_id"] = worker_id
+
+                    # Track helmet violations with stricter criteria
+                    if detection["violation"] == "no_helmet":
+                        # Only include high-confidence no_helmet detections
+                        if conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                            if worker_id not in helmet_violations:
+                                helmet_violations[worker_id] = []
+                            helmet_violations[worker_id].append(detection)
+                    else:
+                        violations.append(detection)
+
+                # Remove inactive workers
+                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
         cap.release()
-        if os.path.exists(video_path):
-            os.remove(video_path)
-            
+        os.remove(video_path)
         processing_time = time.time() - start_time
-        logger.info(f"Processing complete in {processing_time:.2f}s")
-
-        # Convert tracked violations to final violation list
-        violations = []
-        for worker_id, worker_violations in unique_violations.items():
-            for label, detection_time in worker_violations.items():
-                violation = {
-                    "worker_id": worker_id,
-                    "violation": label,
-                    "timestamp": detection_time
-                }
-                violations.append(violation)
+        logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
+
+        # Confirm helmet violations (require multiple detections)
+        for worker_id, detections in helmet_violations.items():
+            if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                # Select the detection with the highest confidence
+                best_detection = max(detections, key=lambda x: x["confidence"])
+                violations.append(best_detection)
+
+                # Capture snapshot for confirmed no_helmet violation
+                cap = cv2.VideoCapture(video_path)
+                cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                ret, snapshot_frame = cap.read()
+                if ret:
+                    snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                    snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
+                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                    cv2.imwrite(snapshot_path, snapshot_frame)
+                    snapshots.append({
+                        "violation": "no_helmet",
+                        "frame": best_detection["frame"],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    })
+                cap.release()
 
+        # Generate results
         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
 
-        # Calculate safety score
         score = calculate_safety_score(violations)
-        
-        # 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) | Confidence |\n"
-        violation_table += "|-----------|-----------|----------|------------|\n"
-        
-        for v in sorted(violations, key=lambda x: (x.get("worker_id", "Unknown"), x.get("timestamp", 0.0))):
+        violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+        violation_table += "|------------------------|---------------|------------|-----------|\n"
+        for v in sorted(violations, key=lambda x: x["timestamp"]):
             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"
+            row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
+            violation_table += row
 
-        if not snapshots_text:
-            snapshots_text = "No snapshots captured."
+        snapshots_text = "\n".join(
+            f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+            for s in snapshots
+        ) if snapshots else "No snapshots captured."
 
         yield (
             violation_table,
@@ -733,27 +495,24 @@ def process_video(video_data):
 
     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}", "", "", "", ""
 
+# ==========================
+# Gradio Interface
+# ==========================
 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_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
-            
     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"),
@@ -765,7 +524,7 @@ interface = gr.Interface(
         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.",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
     allow_flagging="never"
 )