Update app.py

app.py

@@ -15,10 +15,10 @@ import logging
 from retrying import retry
 
 # ==========================
-# OPTIMIZED CONFIGURATION
+# Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",
+    "MODEL_PATH": "yolov8_safety.pt",  # Your custom-trained model
     "FALLBACK_MODEL": "yolov8n.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
@@ -28,7 +28,7 @@ CONFIG = {
         3: "unsafe_zone",
         4: "improper_tool_use"
     },
-    "CLASS_COLORS": {
+    "CLASS_COLORS": {  # Bounding box colors
         "no_helmet": (0, 0, 255),      # Red
         "no_harness": (0, 165, 255),   # Orange
         "unsafe_posture": (0, 255, 0), # Green
@@ -37,267 +37,221 @@ CONFIG = {
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet",
-        "no_harness": "No Harness",
+        "no_harness": "No Safety Harness",
         "unsafe_posture": "Unsafe Posture",
-        "unsafe_zone": "Unsafe Zone",
+        "unsafe_zone": "Unsafe Zone Entry",
         "improper_tool_use": "Improper Tool Use"
     },
-    "SF_CREDENTIALS": {
+    "SF_CREDENTIALS": {  # Salesforce credentials
         "username": "prashanth1ai@safety.com",
         "password": "SaiPrash461",
         "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 3,
-    "MAX_PROCESSING_TIME": 60,
-    "CONFIDENCE_THRESHOLD": {
+    "FRAME_SKIP": 5,  # Process every 5th frame (balance speed vs. accuracy)
+    "MAX_PROCESSING_TIME": 60,  # Max processing time (seconds)
+    "CONFIDENCE_THRESHOLD": {  # Per-class thresholds
         "no_helmet": 0.4,
         "no_harness": 0.3,
         "unsafe_posture": 0.25,
         "unsafe_zone": 0.3,
         "improper_tool_use": 0.35
     },
-    "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3,
-    "MIN_VIOLATION_DURATION": 1.5
+    "IOU_THRESHOLD": 0.4,  # For worker tracking
+    "MIN_VIOLATION_FRAMES": 3  # Min frames to confirm a violation
 }
 
-# Initialize logging
+# 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 configuration
+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}")
 
+# ==========================
+# Load YOLOv8 Model
+# ==========================
 def load_model():
     try:
-        if os.path.exists(CONFIG["MODEL_PATH"]):
+        if os.path.isfile(CONFIG["MODEL_PATH"]):
             model = YOLO(CONFIG["MODEL_PATH"]).to(device)
             logger.info("Loaded custom safety model")
         else:
             model = YOLO(CONFIG["FALLBACK_MODEL"]).to(device)
-            logger.warning("Using fallback model - recommend training yolov8_safety.pt")
+            logger.warning("Using fallback model (lower accuracy)")
         return model
     except Exception as e:
-        logger.error(f"Model loading failed: {str(e)}")
+        logger.error(f"Model load failed: {e}")
         raise
 
 model = load_model()
 
+# ==========================
+# Core Detection Functions
+# ==========================
 def draw_detections(frame, detections):
-    """Draw bounding boxes with labels and confidence scores"""
+    """Draw bounding boxes with labels on frame."""
     for det in detections:
         label = det["violation"]
-        x, y, w, h = [int(v) for v in det["bounding_box"]]
-        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
-        
+        conf = det["confidence"]
+        x, y, w, h = det["bounding_box"]
         x1, y1 = int(x - w/2), int(y - h/2)
         x2, y2 = int(x + w/2), int(y + h/2)
         
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-        
-        label_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)}: {det['confidence']:.2f}"
-        (text_width, text_height), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(frame, (x1, y1 - text_height - 10), (x1 + text_width, y1), color, -1)
-        cv2.putText(frame, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+        cv2.putText(frame, f"{label}: {conf:.2f}", (x1, y1-10), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     return frame
 
 def calculate_iou(box1, box2):
-    """Calculate Intersection over Union for two bounding boxes"""
-    box1 = [box1[0] - box1[2]/2, box1[1] - box1[3]/2, box1[0] + box1[2]/2, box1[1] + box1[3]/2]
-    box2 = [box2[0] - box2[2]/2, box2[1] - box2[3]/2, box2[0] + box2[2]/2, box2[1] + box2[3]/2]
-    
-    x_left = max(box1[0], box2[0])
-    y_top = max(box1[1], box2[1])
-    x_right = min(box1[2], box2[2])
-    y_bottom = min(box1[3], box2[3])
-    
-    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 = (box1[2] - box1[0]) * (box1[3] - box1[1])
-    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
-    
-    return intersection_area / float(box1_area + box2_area - intersection_area)
+    """Compute Intersection-over-Union for tracking."""
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
+    x_min = max(x1 - w1/2, x2 - w2/2)
+    y_min = max(y1 - h1/2, y2 - h2/2)
+    x_max = min(x1 + w1/2, x2 + w2/2)
+    y_max = min(y1 + h1/2, y2 + h2/2)
+    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    union = w1 * h1 + w2 * h2 - intersection
+    return intersection / union if union > 0 else 0
+
+# ==========================
+# Salesforce Integration
+# ==========================
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
+def connect_to_salesforce():
+    try:
+        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
+        logger.info("Salesforce connection successful")
+        return sf
+    except Exception as e:
+        logger.error(f"Salesforce login failed: {e}")
+        raise
 
 def generate_violation_pdf(violations, score):
+    """Generate PDF report with 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)
         c.setFont("Helvetica-Bold", 14)
         c.drawString(1 * inch, 10.5 * inch, "Worksite Safety Violation Report")
         c.setFont("Helvetica", 12)
         
-        y_position = 10 * inch
+        # Report metadata
+        y_pos = 9.8 * inch
         report_data = [
             ("Compliance Score", f"{score}%"),
             ("Total Violations", len(violations)),
-            ("Report Date", time.strftime("%Y-%m-%d %H:%M:%S"))
+            ("Date", time.strftime("%Y-%m-%d")),
+            ("Time", time.strftime("%H:%M:%S"))
         ]
+        for label, value in report_data:
+            c.drawString(1 * inch, y_pos, f"{label}: {value}")
+            y_pos -= 0.4 * inch
         
-        for key, value in report_data:
-            c.drawString(1 * inch, y_position, f"{key}: {value}")
-            y_position -= 0.4 * inch
-            
-        y_position -= 0.2 * inch
-        c.line(1 * inch, y_position, 7.5 * inch, y_position)
-        y_position -= 0.3 * inch
-        
+        # Violation details
+        y_pos -= 0.3 * inch
         c.setFont("Helvetica-Bold", 12)
-        c.drawString(1 * inch, y_position, "Violation Details:")
-        y_position -= 0.3 * inch
+        c.drawString(1 * inch, y_pos, "Violation Details:")
         c.setFont("Helvetica", 10)
+        y_pos -= 0.3 * inch
         
         if not violations:
-            c.drawString(1 * inch, y_position, "No violations detected.")
+            c.drawString(1 * inch, y_pos, "No violations detected.")
         else:
             for v in violations:
-                violation_text = (
+                text = (
                     f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} "
-                    f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f}, "
-                    f"Worker: {v['worker_id']})"
+                    f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})"
                 )
-                c.drawString(1 * inch, y_position, violation_text)
-                y_position -= 0.25 * inch
-                if y_position < 1 * inch:
+                c.drawString(1 * inch, y_pos, text)
+                y_pos -= 0.25 * inch
+                if y_pos < 1 * inch:
                     c.showPage()
-                    y_position = 10 * inch
-                    c.setFont("Helvetica", 10)
+                    y_pos = 10 * inch
         
         c.save()
         pdf_file.seek(0)
         
+        # Save PDF
+        pdf_filename = f"violation_report_{int(time.time())}.pdf"
+        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         with open(pdf_path, "wb") as f:
             f.write(pdf_file.getvalue())
-            
-        public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
-        logger.info(f"Generated PDF report: {public_url}")
-        return pdf_path, public_url, pdf_file
         
+        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}", pdf_file
     except Exception as e:
-        logger.error(f"PDF generation failed: {str(e)}")
-        return "", "", None
-
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
-def connect_to_salesforce():
-    try:
-        sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
-        logger.info("Connected to Salesforce")
-        return sf
-    except Exception as e:
-        logger.error(f"Salesforce connection failed: {str(e)}")
-        raise
-
-def upload_pdf_to_salesforce(sf, pdf_file, report_id):
-    try:
-        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
-        content_version = sf.ContentVersion.create({
-            "Title": f"Safety_Report_{int(time.time())}",
-            "PathOnClient": "safety_report.pdf",
-            "VersionData": encoded_pdf,
-            "FirstPublishLocationId": report_id
-        })
-        return f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
-    except Exception as e:
-        logger.error(f"PDF upload failed: {str(e)}")
-        return ""
+        logger.error(f"PDF generation failed: {e}")
+        return None, None, None
 
-def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+def push_report_to_salesforce(violations, score, pdf_file):
+    """Upload report to Salesforce."""
     try:
         sf = connect_to_salesforce()
         
+        # Create violation details text
         violations_text = "\n".join(
-            f"- {CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} "
-            f"at {v['timestamp']:.2f}s (Worker {v['worker_id']}, Confidence: {v['confidence']:.2f})"
+            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} "
+            f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})"
             for v in violations
-        ) or "No violations detected"
+        ) or "No violations detected."
         
+        # Create Salesforce record
         record_data = {
             "Compliance_Score__c": score,
             "Violations_Found__c": len(violations),
             "Violations_Details__c": violations_text,
-            "Status__c": "New"
+            "Status__c": "Pending Review"
         }
+        record = sf.Safety_Video_Report__c.create(record_data)
+        record_id = record["id"]
         
-        try:
-            record = sf.Safety_Video_Report__c.create(record_data)
-            record_id = record["id"]
-            logger.info(f"Created Salesforce record: {record_id}")
-        except Exception as e:
-            logger.error(f"Failed to create Safety Report: {str(e)}")
-            record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
-            record_id = record["id"]
-            logger.warning(f"Created fallback Account record: {record_id}")
-        
+        # Upload PDF if available
         pdf_url = ""
         if pdf_file:
-            pdf_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
-            if pdf_url:
-                try:
-                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": pdf_url})
-                except:
-                    sf.Account.update(record_id, {"Description": pdf_url})
+            encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode("utf-8")
+            content_version = sf.ContentVersion.create({
+                "Title": f"Safety_Report_{record_id}",
+                "PathOnClient": f"report_{record_id}.pdf",
+                "VersionData": encoded_pdf,
+                "FirstPublishLocationId": record_id
+            })
+            pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
         
-        return record_id, pdf_url if pdf_url else f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}"
-    
+        return record_id, pdf_url
     except Exception as e:
-        logger.error(f"Salesforce integration failed: {str(e)}")
+        logger.error(f"Salesforce upload failed: {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
-    }
-    unique_violations = {(v["worker_id"], v["violation"]) for v in violations}
-    total_penalty = sum(penalties.get(v[1], 0) for v in unique_violations)
-    return max(100 - total_penalty, 0)
-
-def process_video(video_data):
+# ==========================
+# Video Processing
+# ==========================
+def process_video(video_path):
+    """Analyze video for safety violations."""
     try:
-        temp_video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
-        with open(temp_video_path, "wb") as f:
-            f.write(video_data)
-        
-        cap = cv2.VideoCapture(temp_video_path)
+        cap = cv2.VideoCapture(video_path)
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        
-        workers = []
+        frame_count = 0
         violations = []
         snapshots = []
-        violation_history = {k: [] for k in CONFIG["VIOLATION_LABELS"].values()}
-        snapshot_taken = {k: False for k in CONFIG["VIOLATION_LABELS"].values()}
-        
-        frame_count = 0
-        start_time = time.time()
+        workers = []
+        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         
         while cap.isOpened():
             ret, frame = cap.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.warning("Processing timeout reached")
-                break
-                
+            
+            # Run detection
+            results = model(frame, device=device)
             current_time = frame_count / fps
-            results = model(frame, device=device, verbose=False)
             
             for result in results:
                 for box in result.boxes:
@@ -307,150 +261,160 @@ def process_video(video_data):
                     
                     if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3):
                         continue
-                        
-                    bbox = box.xywh.cpu().numpy()[0].tolist()
                     
+                    bbox = box.xywh.cpu().numpy()[0]
+                    detection = {
+                        "frame": frame_count,
+                        "violation": label,
+                        "confidence": conf,
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    }
+                    
+                    # Track worker
                     matched_worker = None
                     max_iou = 0
                     for worker in workers:
-                        iou = calculate_iou(bbox, worker["bbox"])
+                        iou = calculate_iou(worker["bbox"], bbox)
                         if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
                             max_iou = iou
                             matched_worker = worker
-                            
+                    
                     if matched_worker:
                         worker_id = matched_worker["id"]
                         matched_worker["bbox"] = bbox
-                        matched_worker["last_seen"] = current_time
                     else:
                         worker_id = len(workers) + 1
-                        workers.append({
-                            "id": worker_id,
-                            "bbox": bbox,
-                            "first_seen": current_time,
-                            "last_seen": current_time
-                        })
+                        workers.append({"id": worker_id, "bbox": bbox})
                     
-                    violation_history[label].append({
-                        "frame": frame_count,
-                        "violation": label,
-                        "confidence": round(conf, 2),
-                        "bounding_box": bbox,
-                        "timestamp": current_time,
-                        "worker_id": worker_id
-                    })
+                    detection["worker_id"] = worker_id
+                    violations.append(detection)
+                    
+                    # Capture snapshot if first detection of this type
+                    if not snapshot_taken[label]:
+                        snapshot_path = os.path.join(
+                            CONFIG["OUTPUT_DIR"],
+                            f"{label}_{frame_count}.jpg"
+                        )
+                        cv2.imwrite(snapshot_path, draw_detections(frame.copy(), [detection]))
+                        snapshots.append({
+                            "violation": label,
+                            "frame": frame_count,
+                            "path": snapshot_path,
+                            "url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}"
+                        })
+                        snapshot_taken[label] = True
             
             frame_count += 1
         
-        for violation_type, detections in violation_history.items():
-            if not detections:
-                continue
-                
-            worker_groups = {}
-            for det in detections:
-                if det["worker_id"] not in worker_groups:
-                    worker_groups[det["worker_id"]] = []
-                worker_groups[det["worker_id"]].append(det)
-            
-            for worker_id, worker_dets in worker_groups.items():
-                if len(worker_dets) < 2:
-                    continue
-                    
-                duration = worker_dets[-1]["timestamp"] - worker_dets[0]["timestamp"]
-                if duration >= CONFIG["MIN_VIOLATION_DURATION"]:
-                    best_det = max(worker_dets, key=lambda x: x["confidence"])
-                    violations.append(best_det)
-                    
-                    if not snapshot_taken[violation_type]:
-                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_det["frame"])
-                        ret, snapshot_frame = cap.read()
-                        if ret:
-                            snapshot_frame = draw_detections(snapshot_frame, [best_det])
-                            filename = f"{violation_type}_{best_det['frame']}.jpg"
-                            path = os.path.join(CONFIG["OUTPUT_DIR"], filename)
-                            cv2.imwrite(path, snapshot_frame)
-                            snapshots.append({
-                                "violation": violation_type,
-                                "frame": best_det["frame"],
-                                "path": path,
-                                "url": f"{CONFIG['PUBLIC_URL_BASE']}{filename}"
-                            })
-                            snapshot_taken[violation_type] = True
-        
         cap.release()
-        os.remove(temp_video_path)
         
-        score = calculate_safety_score(violations)
-        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
-        record_id, sf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
+        # Filter violations (require min frames)
+        filtered_violations = []
+        violation_counts = {}
+        for v in violations:
+            key = (v["worker_id"], v["violation"])
+            violation_counts[key] = violation_counts.get(key, 0) + 1
+        
+        for v in violations:
+            if violation_counts[(v["worker_id"], v["violation"])] >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                filtered_violations.append(v)
+        
+        # Calculate safety score
+        penalty_weights = {
+            "no_helmet": 25,
+            "no_harness": 30,
+            "unsafe_posture": 20,
+            "unsafe_zone": 35,
+            "improper_tool_use": 25
+        }
+        unique_violations = set((v["worker_id"], v["violation"]) for v in filtered_violations)
+        total_penalty = sum(penalty_weights.get(v, 0) for _, v in unique_violations)
+        safety_score = max(100 - total_penalty, 0)
         
         return {
-            "violations": violations,
+            "violations": filtered_violations,
             "snapshots": snapshots,
-            "score": score,
-            "salesforce_record_id": record_id,
-            "violation_details_url": sf_url or pdf_url,
+            "score": safety_score,
             "message": ""
         }
-        
     except Exception as e:
-        logger.error(f"Video processing failed: {str(e)}")
+        logger.error(f"Video processing failed: {e}")
         return {
             "violations": [],
             "snapshots": [],
             "score": 100,
-            "salesforce_record_id": None,
-            "violation_details_url": "",
             "message": f"Error: {str(e)}"
         }
 
-def gradio_interface(video_file):
+# ==========================
+# Gradio Interface
+# ==========================
+def analyze_video(video_file):
+    """Gradio interface function."""
+    if not video_file:
+        return "No video uploaded", "", "", "", ""
+    
     try:
-        yield "Analyzing video...", "", "", "", ""
+        # Process video
+        result = process_video(video_file)
+        if result["message"]:
+            return result["message"], "", "", "", ""
         
-        with open(video_file, "rb") as f:
-            result = process_video(f.read())
+        # Generate report
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(
+            result["violations"], 
+            result["score"]
+        )
+        record_id, sf_url = push_report_to_salesforce(
+            result["violations"], 
+            result["score"], 
+            pdf_file
+        )
         
+        # Format outputs
         violation_table = (
-            "| Violation Type      | Timestamp | Confidence | Worker ID |\n"
-            "|---------------------|-----------|------------|-----------|\n" +
+            "| Violation Type         | Timestamp (s) | Confidence | Worker ID |\n"
+            "|------------------------|---------------|------------|-----------|\n" +
             "\n".join(
-                f"| {CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation']):<19} | "
-                f"{v['timestamp']:.2f}    | "
-                f"{v['confidence']:.2f}      | "
-                f"{v['worker_id']} |"
+                f"| {CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation']):<22} | "
+                f"{v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
                 for v in result["violations"]
-            )
-        ) if result["violations"] else "No violations detected"
+            ) if result["violations"] else "No violations detected."
+        )
         
         snapshots_md = "\n".join(
-            f"![{s['violation']} at frame {s['frame']}]({s['url']})"
+            f"**{CONFIG['DISPLAY_NAMES'].get(s['violation'], s['violation'])}** "
+            f"(Frame {s['frame']}): ![]({s['url']})"
             for s in result["snapshots"]
-        ) if result["snapshots"] else "No snapshots"
+        ) if result["snapshots"] else "No snapshots available."
         
-        yield (
+        return (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_md,
-            f"Salesforce ID: {result['salesforce_record_id'] or 'None'}",
-            result["violation_details_url"] or "None"
+            f"Salesforce Record: {record_id or 'N/A'}",
+            sf_url or pdf_url or "N/A"
         )
     except Exception as e:
-        logger.error(f"Interface error: {str(e)}")
-        yield f"Error: {str(e)}", "", "", "", ""
+        return f"Error: {str(e)}", "", "", "", ""
 
+# Launch Gradio App
 interface = gr.Interface(
-    fn=gradio_interface,
+    fn=analyze_video,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
-        gr.Markdown(label="Violations Detected"),
-        gr.Textbox(label="Compliance Score"),
-        gr.Markdown(label="Evidence Snapshots"),
-        gr.Textbox(label="Salesforce Record"),
+        gr.Markdown("## Detected Violations"),
+        gr.Textbox(label="Safety Score"),
+        gr.Markdown("## Violation Snapshots"),
+        gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Report URL")
     ],
-    title="AI Safety Compliance Monitor",
-    description="Detects 5 violation types: No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use"
+    title="AI Safety Compliance Analyzer",
+    description=(
+        "Upload worksite video to detect safety violations. "
+        "Supported violations: Missing Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use."
+    )
 )
 
 if __name__ == "__main__":