Update app.py

app.py

@@ -15,7 +15,7 @@ import logging
 from retrying import retry
 
 # ==========================
-# Enhanced Configuration
+# OPTIMIZED CONFIGURATION
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
@@ -36,10 +36,10 @@ CONFIG = {
         "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",
+        "no_helmet": "No Helmet",
+        "no_harness": "No Harness",
+        "unsafe_posture": "Unsafe Posture",
+        "unsafe_zone": "Unsafe Zone",
         "improper_tool_use": "Improper Tool Use"
     },
     "SF_CREDENTIALS": {
@@ -49,217 +49,206 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 5,  # Reduced for better detection
+    "FRAME_SKIP": 3,
     "MAX_PROCESSING_TIME": 60,
-    "CONFIDENCE_THRESHOLD": 0.25,  # Lower threshold for all violations
+    "CONFIDENCE_THRESHOLD": {
+        "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  # Minimum consecutive frames to confirm violation
+    "MIN_VIOLATION_FRAMES": 3,
+    "MIN_VIOLATION_DURATION": 1.5
 }
 
-# Setup logging
+# Initialize 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
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        if os.path.isfile(CONFIG["MODEL_PATH"]):
-            model_path = CONFIG["MODEL_PATH"]
-            logger.info(f"Model loaded: {model_path}")
+        if os.path.exists(CONFIG["MODEL_PATH"]):
+            model = YOLO(CONFIG["MODEL_PATH"]).to(device)
+            logger.info("Loaded custom safety model")
         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)
+            model = YOLO(CONFIG["FALLBACK_MODEL"]).to(device)
+            logger.warning("Using fallback model - recommend training yolov8_safety.pt")
         return model
     except Exception as e:
-        logger.error(f"Failed to load model: {e}")
+        logger.error(f"Model loading failed: {str(e)}")
         raise
 
 model = load_model()
 
-# ==========================
-# Enhanced Helper Functions
-# ==========================
 def draw_detections(frame, detections):
-    """Draw bounding boxes and labels on frame"""
+    """Draw bounding boxes with labels and confidence scores"""
     for det in detections:
         label = det["violation"]
-        confidence = det["confidence"]
-        x, y, w, h = det["bounding_box"]
+        x, y, w, h = [int(v) for v in det["bounding_box"]]
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         
-        # 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)
+        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)
         
-        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)
+        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)
     return frame
 
 def calculate_iou(box1, box2):
-    """Calculate Intersection over Union (IoU) for two bounding boxes."""
-    x1, y1, w1, h1 = box1
-    x2, y2, w2, h2 = 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]
     
-    # 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
+    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])
     
-    # 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)
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
     
-    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
-    area1 = w1 * h1
-    area2 = w2 * h2
-    union = area1 + area2 - intersection
+    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 / union if union > 0 else 0
-
-# ==========================
-# Salesforce Integration (unchanged)
-# ==========================
-@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
+    return intersection_area / float(box1_area + box2_area - intersection_area)
 
 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)
+        c.setFont("Helvetica-Bold", 14)
+        c.drawString(1 * inch, 10.5 * inch, "Worksite Safety Violation Report")
         c.setFont("Helvetica", 12)
-        c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
-        c.setFont("Helvetica", 10)
-
-        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():
+        
+        y_position = 10 * inch
+        report_data = [
+            ("Compliance Score", f"{score}%"),
+            ("Total Violations", len(violations)),
+            ("Report Date", time.strftime("%Y-%m-%d %H:%M:%S"))
+        ]
+        
+        for key, value in report_data:
             c.drawString(1 * inch, y_position, f"{key}: {value}")
-            y_position -= 0.3 * inch
-
+            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
+        
+        c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Violation Details:")
         y_position -= 0.3 * inch
+        c.setFont("Helvetica", 10)
+        
         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
+                violation_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']})"
+                )
+                c.drawString(1 * inch, y_position, violation_text)
+                y_position -= 0.25 * inch
                 if y_position < 1 * inch:
                     c.showPage()
-                    c.setFont("Helvetica", 10)
                     y_position = 10 * inch
-
-        c.showPage()
+                    c.setFont("Helvetica", 10)
+        
         c.save()
         pdf_file.seek(0)
-
+        
         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}")
+        logger.info(f"Generated PDF report: {public_url}")
         return pdf_path, public_url, pdf_file
+        
     except Exception as e:
-        logger.error(f"Error generating PDF: {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:
-        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",
+        content_version = sf.ContentVersion.create({
+            "Title": f"Safety_Report_{int(time.time())}",
+            "PathOnClient": "safety_report.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
+        })
+        return f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
     except Exception as e:
-        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        logger.error(f"PDF upload failed: {str(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']})"
+            f"- {CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} "
+            f"at {v['timestamp']:.2f}s (Worker {v['worker_id']}, Confidence: {v['confidence']:.2f})"
             for v in violations
-        ) or "No violations detected."
-        pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
-
+        ) or "No violations detected"
+        
         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
+            "Status__c": "New"
         }
-        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']}")
+            record_id = record["id"]
+            logger.info(f"Created Salesforce record: {record_id}")
         except Exception as e:
-            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
+            logger.error(f"Failed to create Safety Report: {str(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"]
-
+            record_id = record["id"]
+            logger.warning(f"Created fallback Account record: {record_id}")
+        
+        pdf_url = ""
         if pdf_file:
-            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
-            if uploaded_url:
+            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": 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
+                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": pdf_url})
+                except:
+                    sf.Account.update(record_id, {"Description": pdf_url})
+        
+        return record_id, pdf_url if pdf_url else f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}"
+    
     except Exception as e:
-        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
+        logger.error(f"Salesforce integration failed: {str(e)}")
         return None, ""
 
 def calculate_safety_score(violations):
@@ -270,265 +259,199 @@ def calculate_safety_score(violations):
         "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)
+    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)
 
-# ==========================
-# Enhanced Video Processing
-# ==========================
 def process_video(video_data):
     try:
-        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
-        with open(video_path, "wb") as f:
+        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)
-        logger.info(f"Video saved: {video_path}")
-
-        video = cv2.VideoCapture(video_path)
-        if not video.isOpened():
-            raise ValueError("Could not open video file")
-
+        
+        cap = cv2.VideoCapture(temp_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 = []
         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()
-        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()
+        
+        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.info("Processing time limit reached")
+                logger.warning("Processing timeout reached")
                 break
-
+                
             current_time = frame_count / fps
+            results = model(frame, device=device, verbose=False)
             
-            # Run detection on this frame
-            results = model(frame, device=device)
-            
-            current_detections = []
             for result in results:
-                boxes = result.boxes
-                for box in boxes:
+                for box in result.boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls)
                     
-                    if label is None:
+                    if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3):
                         continue
                         
-                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
-                        continue
-
-                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    bbox = box.xywh.cpu().numpy()[0].tolist()
+                    
+                    matched_worker = None
+                    max_iou = 0
+                    for worker in workers:
+                        iou = calculate_iou(bbox, worker["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
+                        })
                     
-                    current_detections.append({
+                    violation_history[label].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
+                        "timestamp": current_time,
+                        "worker_id": worker_id
                     })
-                
-                # Add to violation history
-                detection["worker_id"] = worker_id
-                violation_history[detection["violation"]].append(detection)
-
+            
             frame_count += 1
-
-        video.release()
-        os.remove(video_path)
         
-        # Process violation history to confirm persistent violations
         for violation_type, detections in violation_history.items():
             if not detections:
                 continue
                 
-            # Group by worker
-            worker_violations = {}
+            worker_groups = {}
             for det in detections:
-                if det["worker_id"] not in worker_violations:
-                    worker_violations[det["worker_id"]] = []
-                worker_violations[det["worker_id"]].append(det)
+                if det["worker_id"] not in worker_groups:
+                    worker_groups[det["worker_id"]] = []
+                worker_groups[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)
+            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)
                     
-                    # 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"])
+                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_det["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)
+                            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_detection["frame"],
-                                "snapshot_path": snapshot_path,
-                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                                "frame": best_det["frame"],
+                                "path": path,
+                                "url": f"{CONFIG['PUBLIC_URL_BASE']}{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."
-            }
-
+        
+        cap.release()
+        os.remove(temp_video_path)
+        
         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)
-
+        record_id, sf_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,
+            "salesforce_record_id": record_id,
+            "violation_details_url": sf_url or pdf_url,
             "message": ""
         }
+        
     except Exception as e:
-        logger.error(f"Error processing video: {e}", exc_info=True)
+        logger.error(f"Video processing failed: {str(e)}")
         return {
             "violations": [],
             "snapshots": [],
             "score": 100,
             "salesforce_record_id": None,
             "violation_details_url": "",
-            "message": f"Error processing video: {e}"
+            "message": f"Error: {str(e)}"
         }
 
-# ==========================
-# Gradio Interface
-# ==========================
 def gradio_interface(video_file):
-    if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
-        yield "Processing video... please wait.", "", "", "", ""
-
+        yield "Analyzing video...", "", "", "", ""
+        
         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"]
+            result = process_video(f.read())
+        
+        violation_table = (
+            "| Violation Type      | Timestamp | 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']} |"
+                for v in result["violations"]
             )
-
+        ) if result["violations"] else "No violations detected"
+        
+        snapshots_md = "\n".join(
+            f"![{s['violation']} at frame {s['frame']}]({s['url']})"
+            for s in result["snapshots"]
+        ) if result["snapshots"] else "No 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"
+            snapshots_md,
+            f"Salesforce ID: {result['salesforce_record_id'] or 'None'}",
+            result["violation_details_url"] or "None"
         )
     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: {str(e)}")
+        yield f"Error: {str(e)}", "", "", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
-    outputs=[ 
-        gr.Markdown(label="Detected Safety Violations"),
+    outputs=[
+        gr.Markdown(label="Violations Detected"),
         gr.Textbox(label="Compliance Score"),
-        gr.Markdown(label="Snapshots"),
-        gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL")
+        gr.Markdown(label="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). Non-violations are ignored.",
-    allow_flagging="never"
+    title="AI Safety Compliance Monitor",
+    description="Detects 5 violation types: No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
-    interface.launch()
+    interface.launch(share=True)