Update app.py

app.py

@@ -19,7 +19,6 @@ from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
-import face_recognition
 from collections import defaultdict
 
 # ========================== # Configuration and Setup # ==========================
@@ -28,11 +27,9 @@ os.makedirs('/tmp/Ultralytics', exist_ok=True)
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
-
-# Suppress warnings
 warnings.filterwarnings("ignore")
 
-# ========================== # Enhanced Tracker Implementation # ==========================
+# ========================== # Position-Based 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
@@ -41,13 +38,12 @@ class SafetyTracker:
         self.frame_rate = frame_rate
         self.next_id = 1
         
-        # Trackers for different purposes
-        self.worker_tracks = {}  # Active worker tracks
-        self.violation_history = defaultdict(dict)  # Track violations per worker
-        self.face_encodings = {}  # Store face encodings for helmet violations
-        self.position_history = defaultdict(list)  # Track positions for non-helmet violations
+        # Tracking stores
+        self.worker_tracks = {}  # Active tracks
+        self.violation_history = defaultdict(dict)  # {worker_id: {violation_type: last_detection_time}}
+        self.position_history = defaultdict(list)  # {worker_id: [positions]}
         
-        # Cooldown periods (in seconds)
+        # Violation cooldowns (seconds)
         self.VIOLATION_COOLDOWNS = {
             "no_helmet": 30.0,
             "no_harness": 20.0,
@@ -56,10 +52,9 @@ class SafetyTracker:
             "improper_tool_use": 15.0
         }
 
-    def update(self, detections, frame):
-        """Update tracks with new detections and check for violations"""
+    def update(self, detections):
+        """Update tracks with new detections using position-based matching"""
         current_time = time.time()
-        active_violations = []
         new_violations = []
         
         for det in detections:
@@ -67,20 +62,15 @@ class SafetyTracker:
             label = det['violation']
             confidence = det['confidence']
             
-            # 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)
+            # Match by position
+            worker_id = self._match_by_position(bbox, label)
             
             if worker_id is None:
                 worker_id = self.next_id
                 self.next_id += 1
                 
-            # Check if this is a new violation for this worker
+            # Check if new violation
             if self._is_new_violation(worker_id, label, current_time):
-                # Record the violation
                 violation = {
                     'worker_id': worker_id,
                     'violation': label,
@@ -89,116 +79,64 @@ class SafetyTracker:
                     'timestamp': current_time
                 }
                 new_violations.append(violation)
-                
-                # 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)
             
-            # Keep track of active workers
+            # Update position history
+            x, y, w, h = bbox
+            self.position_history[worker_id].append((x, y))
+            
+            # Update active tracks
             self.worker_tracks[worker_id] = {
                 'bbox': bbox,
                 'last_seen': current_time,
                 'label': label
             }
         
-        # Clean up old tracks
+        # Cleanup old tracks
         self._cleanup_tracks(current_time)
         
         return new_violations
 
-    def _match_by_face(self, bbox, frame):
-        """Match detection by face recognition (for helmet violations)"""
-        x, y, w, h = bbox
-        face_region = frame[max(0, int(y-h/2)):int(y+h/2), max(0, int(x-w/2)):int(x+w/2)]
-        
-        if face_region.size == 0:
-            return None
-            
-        try:
-            # Get face encodings from current detection
-            face_locations = face_recognition.face_locations(face_region)
-            if not face_locations:
-                return None
-                
-            current_encoding = face_recognition.face_encodings(face_region, face_locations)[0]
-            
-            # Compare with known faces
-            for worker_id, encodings in self.face_encodings.items():
-                matches = face_recognition.compare_faces(encodings, current_encoding, tolerance=0.6)
-                if any(matches):
-                    return worker_id
-                    
-        except Exception as e:
-            logger.warning(f"Face recognition error: {e}")
-            
-        return None
-
     def _match_by_position(self, bbox, label):
-        """Match detection by position (for non-helmet violations)"""
+        """Match detection to existing tracks using position"""
         x, y, w, h = bbox
         current_pos = (x, y)
         
         for worker_id, positions in self.position_history.items():
+            # Only match if worker has had this violation type before
             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:
+            # Check distance to historical positions
+            for pos in positions[-5:]:  # Check last 5 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"""
+        """Check if violation is new based on cooldown"""
         if label not in self.violation_history[worker_id]:
             return True
             
-        last_detection = self.violation_history[worker_id][label]
+        last_time = self.violation_history[worker_id][label]
         cooldown = self.VIOLATION_COOLDOWNS.get(label, 10.0)
-        
-        return (current_time - last_detection) > cooldown
-
-    def _store_face_encoding(self, worker_id, bbox, frame):
-        """Store face encoding for a worker"""
-        x, y, w, h = bbox
-        face_region = frame[max(0, int(y-h/2)):int(y+h/2), max(0, int(x-w/2)):int(x+w/2)]
-        
-        if face_region.size == 0:
-            return
-            
-        try:
-            face_locations = face_recognition.face_locations(face_region)
-            if face_locations:
-                encoding = face_recognition.face_encodings(face_region, face_locations)[0]
-                if worker_id not in self.face_encodings:
-                    self.face_encodings[worker_id] = []
-                self.face_encodings[worker_id].append(encoding)
-        except Exception as e:
-            logger.warning(f"Error storing face encoding: {e}")
+        return (current_time - last_time) > cooldown
 
     def _cleanup_tracks(self, current_time):
-        """Clean up old tracks and face encodings"""
-        # Remove inactive workers
+        """Remove inactive tracks"""
         inactive_ids = [
-            worker_id for worker_id, track in self.worker_tracks.items()
+            id for 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)
-            
-            # 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 # ==========================
+        for id in inactive_ids:
+            self.worker_tracks.pop(id, None)
+            self.position_history.pop(id, None)
+            # Keep violation history for longer
+            if (current_time - max(self.violation_history[id].values(), default=0)) > 300:
+                self.violation_history.pop(id, None)
+
+# ========================== # App Configuration # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
@@ -211,11 +149,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),
+        "no_harness": (0, 165, 255),
+        "unsafe_posture": (0, 255, 0),
+        "unsafe_zone": (255, 0, 0),
+        "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -238,502 +176,293 @@ CONFIG = {
         "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
+    "BATCH_SIZE": 8,  # Reduced for stability
+    "SNAPSHOT_QUALITY": 90
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
+# ========================== # Core Functions # ==========================
 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
+        model_path = CONFIG["MODEL_PATH"] if os.path.exists(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
+        logger.info(f"Loading model: {model_path}")
+        if not os.path.exists(model_path):
+            logger.info("Downloading fallback model...")
+            torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
+        return YOLO(model_path).to(device)
     except Exception as e:
-        logger.error(f"Failed to load model: {e}")
+        logger.error(f"Model loading failed: {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
+    """Basic image enhancement"""
+    return cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
 
 def draw_detections(frame, detections):
-    """Draw bounding boxes and labels on detection frame with improved visibility"""
-    result_frame = frame.copy()
-    
+    """Draw bounding boxes with labels"""
+    result = 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))
-        
-        # Draw thicker rectangle with border
-        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+        x, y, w, h = det['bbox']
+        x1, y1 = int(x-w/2), int(y-h/2)
+        x2, y2 = int(x+w/2), int(y+h/2)
+        label = CONFIG["DISPLAY_NAMES"].get(det['violation'], det['violation'])
+        color = CONFIG["CLASS_COLORS"].get(det['violation'], (0,0,255))
         
-        # 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
+        cv2.rectangle(result, (x1, y1), (x2, y2), color, 3)
+        cv2.putText(result, f"{label} (Worker {det['worker_id']})", 
+                   (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2)
+    return result
 
 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,
+    penalty_map = {
+        "no_helmet": 25, "no_harness": 30, 
+        "unsafe_posture": 20, "unsafe_zone": 35, 
         "improper_tool_use": 25
     }
-    
-    # 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
+    unique_violations = {v['violation'] for v in violations}
+    return max(0, 100 - sum(penalty_map.get(v,0) for v in unique_violations))
 
+# ========================== # Reporting Functions # ==========================
 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)
+        pdf_buffer = BytesIO()
+        c = canvas.Canvas(pdf_buffer, pagesize=letter)
         
-        # Title
+        # Header
         c.setFont("Helvetica-Bold", 16)
-        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
-        
-        # Basic Information
+        c.drawString(1*inch, 10*inch, "Safety Violation Report")
         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')}")
+        c.drawString(1*inch, 9.5*inch, f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        c.drawString(1*inch, 9*inch, f"Safety Score: {score}%")
         
-        # Safety Score
+        # Violations list
+        y = 8.5*inch
         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.drawString(1*inch, y, "Violations Detected:")
+        y -= 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:
+            text = f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]} at {v['timestamp']:.1f}s"
+            c.drawString(1.2*inch, y, text)
+            y -= 0.2*inch
+            if y < 1*inch:
                 c.showPage()
-                c.setFont("Helvetica", 10)
-                y_position = 10 * inch
-
+                y = 10*inch
+        
         c.save()
-        pdf_file.seek(0)
-
-        # Save PDF file
+        pdf_buffer.seek(0)
+        
+        # Save to file
+        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], f"report_{int(time.time())}.pdf")
         with open(pdf_path, "wb") as f:
-            f.write(pdf_file.getvalue())
+            f.write(pdf_buffer.getvalue())
             
-        public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
-        logger.info(f"PDF generated: {public_url}")
-        return pdf_path, public_url, pdf_file
+        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}", pdf_buffer
     except Exception as e:
-        logger.error(f"Error generating PDF: {e}")
-        return "", "", None
+        logger.error(f"PDF generation failed: {e}")
+        return None, None, 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()
+        logger.info("Salesforce connection established")
         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"""
+def push_report_to_salesforce(violations, score, pdf_path, pdf_buffer):
     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."
-            
-        pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
-
+        # Create record
         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
+            "Violations_Details__c": "\n".join(
+                f"Worker {v['worker_id']}: {CONFIG['DISPLAY_NAMES'][v['violation']]}"
+                for v in violations
+            ),
+            "Status__c": "New"
         }
         
-        logger.info(f"Creating Salesforce record with data: {record_data}")
+        record = sf.Safety_Video_Report__c.create(record_data)
+        record_id = record['id']
+        logger.info(f"Created Salesforce record: {record_id}")
+        
+        # Upload PDF if available
+        pdf_url = ""
+        if pdf_buffer:
+            encoded = base64.b64encode(pdf_buffer.getvalue()).decode()
+            content_version = sf.ContentVersion.create({
+                "Title": f"Safety_Report_{record_id}",
+                "PathOnClient": "report.pdf",
+                "VersionData": encoded,
+                "FirstPublishLocationId": record_id
+            })
+            pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+            logger.info(f"PDF uploaded: {pdf_url}")
         
-        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)
+        logger.error(f"Salesforce upload failed: {e}")
         return None, ""
 
+# ========================== # Video Processing # ==========================
 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")
+        
+        # Save video
+        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"input_{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():
-            os.remove(video_path)
-            raise ValueError("Could not open video file")
-
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            raise ValueError("Failed to open video")
+            
         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}")
-
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         tracker = SafetyTracker(frame_rate=fps)
         snapshots = []
-        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
-                
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                
-                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
-                frame_counter += 1
-
-            if not batch_frames:
+        
+        frame_count = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
                 break
-
-            # Process batch with YOLO model
-            results = model(batch_frames, device=device, conf=0.1, verbose=False)
-            
-            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 = []
+            if frame_count % CONFIG["FRAME_SKIP"] != 0:
+                frame_count += 1
+                continue
                 
-                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]
+            # Process frame
+            frame = preprocess_frame(frame)
+            results = model(frame, verbose=False)[0]
+            
+            # Get detections
+            detections = []
+            for box in results.boxes:
+                cls = int(box.cls)
+                label = CONFIG["VIOLATION_LABELS"].get(cls)
+                if label and box.conf > CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.3):
                     detections.append({
-                        "bbox": bbox,
-                        "violation": label,
-                        "confidence": conf
+                        'bbox': box.xywh[0].cpu().numpy(),
+                        'violation': label,
+                        'confidence': float(box.conf)
                     })
-
-                if not detections:
-                    continue
-                    
-                # Update tracker with new detections
-                new_violations = tracker.update(detections, batch_frames[i])
-                
-                # Process new violations
-                for violation in new_violations:
-                    # Take snapshot for the new violation
-                    snapshot_frame = batch_frames[i].copy()
-                    snapshot_frame = draw_detections(snapshot_frame, [violation])
-                    
-                    # Add timestamp to snapshot
-                    cv2.putText(
-                        snapshot_frame,
-                        f"Time: {violation['timestamp']:.2f}s",
-                        (10, 30),
-                        cv2.FONT_HERSHEY_SIMPLEX,
-                        0.7,
-                        (255, 255, 255),
-                        2
-                    )
-                    
-                    # Save snapshot with high quality
-                    snapshot_filename = f"violation_{violation['violation']}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
-                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                    
-                    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()
-        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")
-
-        # 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
+            # Update tracker
+            new_violations = tracker.update(detections)
+            
+            # Capture snapshots for new violations
+            for violation in new_violations:
+                snapshot = draw_detections(frame.copy(), [violation])
+                timestamp = time.strftime("%Y%m%d_%H%M%S")
+                img_path = os.path.join(
+                    CONFIG["OUTPUT_DIR"],
+                    f"violation_{violation['worker_id']}_{violation['violation']}_{timestamp}.jpg"
+                )
+                cv2.imwrite(img_path, snapshot, [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
+                snapshots.append({
+                    'path': img_path,
+                    'url': f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(img_path)}",
+                    'violation': violation
                 })
-
+            
+            # Update progress
+            if frame_count % 10 == 0:
+                progress = min(100, frame_count / total_frames * 100)
+                yield f"Processing... {progress:.1f}%", "", "", "", ""
+                
+            frame_count += 1
+            
+        cap.release()
+        os.remove(video_path)
+        
+        # Generate report
+        violations = [
+            {
+                'worker_id': worker_id,
+                'violation': violation_type,
+                'timestamp': detection_time
+            }
+            for worker_id, violations in tracker.violation_history.items()
+            for violation_type, detection_time in violations.items()
+        ]
+        
         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"
+            yield "No violations found", "Safety Score: 100%", "No snapshots", "N/A", "N/A"
             return
-
-        # Calculate safety score
+            
         score = calculate_safety_score(violations)
+        pdf_path, pdf_url, pdf_buffer = generate_violation_pdf(violations, score)
+        record_id, salesforce_url = push_report_to_salesforce(violations, score, pdf_path, pdf_buffer)
         
-        # Generate PDF report
-        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        # Format output
+        violations_table = "| Violation | Worker ID | Time |\n|-----------|-----------|------|\n"
+        violations_table += "\n".join(
+            f"| {CONFIG['DISPLAY_NAMES'][v['violation']]} | {v['worker_id']} | {v['timestamp']:.1f}s |"
+            for v in violations
+        )
         
-        # 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"
+        snapshots_md = "\n\n".join(
+            f"### {CONFIG['DISPLAY_NAMES'][s['violation']['violation']]} (Worker {s['violation']['worker_id']})\n"
+            f"![Snapshot]({s['url']})"
+            for s in snapshots
+        )
         
-        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 (
-            violation_table,
+            violations_table,
             f"Safety Score: {score}%",
-            snapshots_text,
-            f"Salesforce Record ID: {report_id or 'N/A'}",
-            final_pdf_url or "N/A"
+            snapshots_md or "No snapshots",
+            f"Salesforce ID: {record_id or 'N/A'}",
+            salesforce_url or pdf_url or "N/A"
         )
-
+        
     except Exception as e:
-        logger.error(f"Error processing video: {e}", exc_info=True)
+        logger.error(f"Processing failed: {e}")
         if 'video_path' in locals() and os.path.exists(video_path):
             os.remove(video_path)
-        yield f"Error processing video: {e}", "", "", "", ""
+        yield f"Error: {str(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.", "", ""
+        return "Upload a video file", "", "", "", ""
         
     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
+            
+        for output in process_video(video_data):
+            yield output
             
     except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+        logger.error(f"Interface error: {e}")
+        yield f"Error: {str(e)}", "", "", "", ""
 
-# ========================== # Gradio Interface # ==========================
 interface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Video(label="Upload Site Video"),
+    inputs=gr.Video(label="Upload Safety 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")
+        gr.Markdown("## Detected Violations"),
+        gr.Textbox(label="Safety Score"),
+        gr.Markdown("## Evidence Snapshots"),
+        gr.Textbox(label="Salesforce Record"),
+        gr.Textbox(label="Report 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"
+    title="AI Safety Compliance Analyzer",
+    description="Detects PPE and safety violations in worksite videos"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()