Update app.py

app.py

@@ -12,295 +12,523 @@ from reportlab.lib.units import inch
 from io import BytesIO
 import base64
 import logging
-from concurrent.futures import ThreadPoolExecutor
+from retrying import retry
 
 # ==========================
-# Optimized Configuration
+# Enhanced Configuration
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
+    "FALLBACK_MODEL": "yolov8n.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
-        1: "no_harness", 
+        1: "no_harness",
         2: "unsafe_posture",
         3: "unsafe_zone",
         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) # Yellow
+    },
+    "DISPLAY_NAMES": {
+        "no_helmet": "No Helmet Violation",
+        "no_harness": "No Harness Violation",
+        "unsafe_posture": "Unsafe Posture Violation",
+        "unsafe_zone": "Unsafe Zone Entry",
+        "improper_tool_use": "Improper Tool Use"
     },
-    "FRAME_SKIP": 8,  # Balanced speed/accuracy
-    "CONFIDENCE_THRESHOLD": 0.35,
-    "MIN_DETECTIONS": 2,
-    "MAX_WORKERS": 4,
-    "PROCESSING_RESOLUTION": (640, 640),
     "SF_CREDENTIALS": {
         "username": "prashanth1ai@safety.com",
         "password": "SaiPrash461",
         "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
         "domain": "login"
     },
-    "MAX_PROCESSING_SECONDS": 30  # Timeout
+    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
+    "FRAME_SKIP": 5,  # Reduced for better detection
+    "MAX_PROCESSING_TIME": 60,
+    "CONFIDENCE_THRESHOLD": 0.25,  # Lower threshold for all violations
+    "IOU_THRESHOLD": 0.4,
+    "MIN_VIOLATION_FRAMES": 3  # Minimum consecutive frames to confirm violation
 }
 
-# Setup
+# 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")
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+logger.info(f"Using device: {device}")
 
-# Load and warm up model
-model = YOLO(CONFIG["MODEL_PATH"]).to(device)
-model.warmup(imgsz=[1, 3, *CONFIG["PROCESSING_RESOLUTION"]])
+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. Detection accuracy may be poor. 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)
+        return model
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+model = load_model()
 
 # ==========================
-# Core Processing Functions
+# Enhanced Helper Functions
 # ==========================
 def draw_detections(frame, detections):
-    """Draw bounding boxes on frame"""
+    """Draw bounding boxes and labels on frame"""
     for det in detections:
-        label = det["label"]
-        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 = det["violation"]
+        confidence = det["confidence"]
+        x, y, w, h = det["bounding_box"]
+        
+        # Convert from center coordinates to corner coordinates
+        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)
-        cv2.putText(frame, f"{label}: {det['confidence']:.2f}", 
-                   (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+        
+        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 process_frame(frame, frame_count, fps):
-    """Process single frame for violations"""
-    resized = cv2.resize(frame, CONFIG["PROCESSING_RESOLUTION"])
-    results = model(resized, verbose=False)
+def calculate_iou(box1, box2):
+    """Calculate Intersection over Union (IoU) for two bounding boxes."""
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
     
-    frame_violations = []
-    for result in results:
-        for box in result.boxes:
-            conf = box.conf.item()
-            if conf > CONFIG["CONFIDENCE_THRESHOLD"]:
-                label = CONFIG["VIOLATION_LABELS"].get(int(box.cls.item()))
-                if label:
-                    frame_violations.append({
-                        "label": label,
-                        "confidence": conf,
-                        "timestamp": frame_count / fps,
-                        "frame": frame_count,
-                        "bbox": box.xywh.cpu().numpy()[0]
-                    })
-    return frame_violations
-
-def process_video(video_path):
-    """Optimized video processing pipeline"""
-    start_time = time.time()
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS) or 30
-    frame_count = 0
-    all_violations = []
-    snapshots = {}
-    last_update = time.time()
+    # Convert to top-left and bottom-right coordinates
+    x1_min, y1_min = x1 - w1/2, y1 - h1/2
+    x1_max, y1_max = x1 + w1/2, y1 + h1/2
+    x2_min, y2_min = x2 - w2/2, y2 - h2/2
+    x2_max, y2_max = x2 + w2/2, y2 + h2/2
     
-    with ThreadPoolExecutor(max_workers=CONFIG["MAX_WORKERS"]) as executor:
-        futures = []
-        
-        while cap.isOpened():
-            if time.time() - start_time > CONFIG["MAX_PROCESSING_SECONDS"]:
-                logger.warning("Processing timeout reached")
-                break
-                
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            if frame_count % CONFIG["FRAME_SKIP"] == 0:
-                futures.append(executor.submit(
-                    process_frame, 
-                    frame.copy(), 
-                    frame_count,
-                    fps
-                ))
-            
-            frame_count += 1
-            
-            # Process completed frames
-            while futures and futures[0].done():
-                for violation in futures.pop(0).result():
-                    all_violations.append(violation)
-                    if violation["label"] not in snapshots:
-                        snapshots[violation["label"]] = {
-                            "frame": draw_detections(frame.copy(), [violation]),
-                            "timestamp": violation["timestamp"],
-                            "confidence": violation["confidence"]
-                        }
-    
-    cap.release()
+    # Calculate intersection
+    x_min = max(x1_min, x2_min)
+    y_min = max(y1_min, y2_min)
+    x_max = min(x1_max, x2_max)
+    y_max = min(y1_max, y2_max)
     
-    # Filter violations
-    confirmed_violations = []
-    violation_counts = {}
-    for v in all_violations:
-        key = v["label"]
-        violation_counts[key] = violation_counts.get(key, 0) + 1
-        
-    for v in all_violations:
-        if violation_counts[v["label"]] >= CONFIG["MIN_DETECTIONS"]:
-            confirmed_violations.append(v)
-    
-    # Save snapshots
-    snapshot_paths = []
-    for label, data in snapshots.items():
-        if violation_counts.get(label, 0) >= CONFIG["MIN_DETECTIONS"]:
-            path = os.path.join(CONFIG["OUTPUT_DIR"], f"{label}_{int(time.time())}.jpg")
-            cv2.imwrite(path, data["frame"])
-            snapshot_paths.append({
-                "label": label,
-                "path": path,
-                "timestamp": data["timestamp"],
-                "confidence": data["confidence"]
-            })
+    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    area1 = w1 * h1
+    area2 = w2 * h2
+    union = area1 + area2 - intersection
     
-    return confirmed_violations, snapshot_paths, time.time() - start_time
+    return intersection / union if union > 0 else 0
 
 # ==========================
-# Reporting Functions
+# Salesforce Integration (unchanged)
 # ==========================
-def generate_report(violations, snapshots, processing_time):
-    """Generate PDF report"""
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
+def connect_to_salesforce():
+    try:
+        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        logger.info("Connected to Salesforce")
+        sf.describe()
+        return sf
+    except Exception as e:
+        logger.error(f"Salesforce connection failed: {e}")
+        raise
+
+def generate_violation_pdf(violations, score):
     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)
-        
-        # Header
-        c.setFont("Helvetica-Bold", 14)
-        c.drawString(1*inch, 10.5*inch, "Safety Violation Report")
         c.setFont("Helvetica", 12)
-        
-        # Summary
-        y_pos = 9.8*inch
-        c.drawString(1*inch, y_pos, f"Processing Time: {processing_time:.1f}s")
-        y_pos -= 0.4*inch
-        c.drawString(1*inch, y_pos, f"Total Violations: {len(violations)}")
-        y_pos -= 0.6*inch
-        
-        # Violations
-        c.setFont("Helvetica-Bold", 12)
-        c.drawString(1*inch, y_pos, "Violation Details:")
-        y_pos -= 0.3*inch
+        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
         c.setFont("Helvetica", 10)
-        
-        for v in violations[:50]:  # Limit to first 50
-            text = f"{v['label']} at {v['timestamp']:.1f}s (Confidence: {v['confidence']:.2f})"
-            c.drawString(1*inch, y_pos, text)
-            y_pos -= 0.2*inch
-            if y_pos < 1*inch:
-                c.showPage()
-                y_pos = 10*inch
-        
+
+        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 report_data.items():
+            c.drawString(1 * inch, y_position, f"{key}: {value}")
+            y_position -= 0.3 * inch
+
+        y_position -= 0.3 * inch
+        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["violation"], v["violation"])
+                text = f"{display_name} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+                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)
-        return pdf_file
+
+        with open(pdf_path, "wb") as f:
+            f.write(pdf_file.getvalue())
+        public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
+        logger.info(f"PDF generated: {public_url}")
+        return pdf_path, public_url, pdf_file
+    except Exception as e:
+        logger.error(f"Error generating PDF: {e}")
+        return "", "", None
+
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    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):
+    try:
+        sf = connect_to_salesforce()
+        violations_text = "\n".join(
+            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            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 = {
+            "Compliance_Score__c": score,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
+            "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']}")
+        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"Report generation failed: {e}")
-        return None
+        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
+        return None, ""
+
+def calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
+    # Count unique violations per worker
+    unique_violations = set()
+    for v in violations:
+        key = (v["worker_id"], v["violation"])
+        unique_violations.add(key)
+    
+    total_penalty = sum(penalties.get(violation, 0) for _, violation in unique_violations)
+    score = 100 - total_penalty
+    return max(score, 0)
 
 # ==========================
-# Salesforce Integration
+# Enhanced Video Processing
 # ==========================
-def upload_to_salesforce(violations, snapshots, processing_time):
-    """Upload results to Salesforce"""
+def process_video(video_data):
     try:
-        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        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}")
+
+        video = cv2.VideoCapture(video_path)
+        if not video.isOpened():
+            raise ValueError("Could not open video file")
+
+        violations = []
+        snapshots = []
+        frame_count = 0
+        start_time = time.time()
+        fps = video.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 30  # Default assumption if FPS cannot be determined
+
+        # Structure to track workers and their violations
+        workers = []
+        violation_history = {label: [] for label in CONFIG["VIOLATION_LABELS"].values()}
+        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
+
+        logger.info(f"Processing video with FPS: {fps}")
+        logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
+
+        while True:
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            if frame_count % CONFIG["FRAME_SKIP"] != 0:
+                frame_count += 1
+                continue
+
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.info("Processing time limit reached")
+                break
+
+            current_time = frame_count / fps
+            
+            # Run detection on this frame
+            results = model(frame, device=device)
+            
+            current_detections = []
+            for result in results:
+                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:
+                        continue
+                        
+                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                        continue
+
+                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    
+                    current_detections.append({
+                        "frame": frame_count,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    })
+
+            # Process detections and associate with workers
+            for detection in current_detections:
+                # Find matching worker
+                matched_worker = None
+                max_iou = 0
+                
+                for worker in workers:
+                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                        max_iou = iou
+                        matched_worker = worker
+
+                if matched_worker:
+                    # Update worker's position
+                    matched_worker["bbox"] = detection["bounding_box"]
+                    matched_worker["last_seen"] = current_time
+                    worker_id = matched_worker["id"]
+                else:
+                    # New worker
+                    worker_id = len(workers) + 1
+                    workers.append({
+                        "id": worker_id,
+                        "bbox": detection["bounding_box"],
+                        "first_seen": current_time,
+                        "last_seen": current_time
+                    })
+                
+                # Add to violation history
+                detection["worker_id"] = worker_id
+                violation_history[detection["violation"]].append(detection)
+
+            frame_count += 1
+
+        video.release()
+        os.remove(video_path)
         
-        # Create record
-        record_data = {
-            "Processing_Time__c": processing_time,
-            "Violation_Count__c": len(violations),
-            "Status__c": "Completed"
+        # Process violation history to confirm persistent violations
+        for violation_type, detections in violation_history.items():
+            if not detections:
+                continue
+                
+            # Group by worker
+            worker_violations = {}
+            for det in detections:
+                if det["worker_id"] not in worker_violations:
+                    worker_violations[det["worker_id"]] = []
+                worker_violations[det["worker_id"]].append(det)
+            
+            # Check each worker's violations for persistence
+            for worker_id, worker_dets in worker_violations.items():
+                if len(worker_dets) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                    # Take the highest confidence detection
+                    best_detection = max(worker_dets, key=lambda x: x["confidence"])
+                    violations.append(best_detection)
+                    
+                    # Capture snapshot if not already taken
+                    if not snapshot_taken[violation_type]:
+                        # Get the frame for this violation
+                        cap = cv2.VideoCapture(video_path)
+                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                        ret, snapshot_frame = cap.read()
+                        cap.release()
+                        
+                        if ret:
+                            # Draw detections on snapshot
+                            snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                            
+                            snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
+                            snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                            cv2.imwrite(snapshot_path, snapshot_frame)
+                            snapshots.append({
+                                "violation": violation_type,
+                                "frame": best_detection["frame"],
+                                "snapshot_path": snapshot_path,
+                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                            })
+                            snapshot_taken[violation_type] = True
+
+        # Final processing
+        if not violations:
+            logger.info("No persistent violations detected")
+            return {
+                "violations": [],
+                "snapshots": [],
+                "score": 100,
+                "salesforce_record_id": None,
+                "violation_details_url": "",
+                "message": "No violations detected in the video."
+            }
+
+        score = calculate_safety_score(violations)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
+
+        return {
+            "violations": violations,
+            "snapshots": snapshots,
+            "score": score,
+            "salesforce_record_id": report_id,
+            "violation_details_url": final_pdf_url,
+            "message": ""
         }
-        record = sf.Safety_Video_Report__c.create(record_data)
-        
-        # Upload report
-        pdf = generate_report(violations, snapshots, processing_time)
-        if pdf:
-            encoded = base64.b64encode(pdf.getvalue()).decode("utf-8")
-            sf.ContentVersion.create({
-                "Title": f"Safety_Report_{record['id']}",
-                "PathOnClient": "report.pdf",
-                "VersionData": encoded,
-                "FirstPublishLocationId": record['id']
-            })
-        
-        return record['id']
     except Exception as e:
-        logger.error(f"Salesforce upload failed: {e}")
-        return None
+        logger.error(f"Error processing video: {e}", exc_info=True)
+        return {
+            "violations": [],
+            "snapshots": [],
+            "score": 100,
+            "salesforce_record_id": None,
+            "violation_details_url": "",
+            "message": f"Error processing video: {e}"
+        }
 
 # ==========================
 # Gradio Interface
 # ==========================
-def analyze_video(video_file):
-    """Main processing function"""
+def gradio_interface(video_file):
     if not video_file:
-        return "No video uploaded", "", "", ""
-    
+        return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
-        # Process video
-        violations, snapshots, proc_time = process_video(video_file)
-        
-        # Generate outputs
-        violation_table = (
-            "| Violation | Time (s) | Confidence |\n"
-            "|-----------|----------|------------|\n" +
-            "\n".join(
-                f"| {v['label']} | {v['timestamp']:.1f} | {v['confidence']:.2f} |"
-                for v in violations[:20]  # Show first 20
+        yield "Processing video... please wait.", "", "", "", ""
+
+        with open(video_file, "rb") as f:
+            video_data = f.read()
+
+        result = process_video(video_data)
+
+        if result.get("message"):
+            yield result["message"], "", "", "", ""
+            return
+
+        violation_table = "No violations detected."
+        if result["violations"]:
+            header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+            separator = "|------------------------|---------------|------------|-----------|\n"
+            rows = []
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
+            for v in result["violations"]:
+                display_name = violation_name_map.get(v["violation"], v["violation"])
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
+                rows.append(row)
+            violation_table = header + separator + "\n".join(rows)
+
+        snapshots_text = "No snapshots captured."
+        if result["snapshots"]:
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
+            snapshots_text = "\n".join(
+                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                for s in result["snapshots"]
             )
+
+        yield (
+            violation_table,
+            f"Safety Score: {result['score']}%",
+            snapshots_text,
+            f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
+            result["violation_details_url"] or "N/A"
         )
-        
-        snapshot_markdown = "\n".join(
-            f"**{s['label']}** ({s['timestamp']:.1f}s, {s['confidence']:.2f}):\n"
-            f"![]({s['path']})"
-            for s in snapshots
-        )
-        
-        # Salesforce upload (async)
-        sf_id = "Not uploaded"
-        if violations:
-            sf_id = upload_to_salesforce(violations, snapshots, proc_time) or "Upload failed"
-        
-        return (
-            f"Processed in {proc_time:.1f}s\n{violation_table}",
-            snapshot_markdown,
-            f"Found {len(violations)} violations",
-            f"Salesforce ID: {sf_id}"
-        )
-        
     except Exception as e:
-        return f"Error: {str(e)}", "", "", ""
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
-# Launch interface
 interface = gr.Interface(
-    fn=analyze_video,
+    fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
-    outputs=[
-        gr.Markdown("## Violation Results"),
-        gr.Markdown("## Evidence Snapshots"),
-        gr.Textbox(label="Summary"),
-        gr.Textbox(label="Salesforce Record")
+    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="AI Safety Compliance Inspector",
-    description=(
-        "Upload worksite video for automated safety violation detection. "
-        "Detects: Missing helmets, improper harnessing, unsafe zones, and more."
-    ),
+    title="Worksite Safety Violation Analyzer",
+    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"
 )
 
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    logger.info("Launching Enhanced Safety Analyzer App...")
+    interface.launch()