Update app.py

app.py

@@ -2,24 +2,16 @@ import os
 import cv2
 import gradio as gr
 import torch
-import numpy as np
 from ultralytics import YOLO
 import time
-from simple_salesforce import Salesforce
-from reportlab.lib.pagesizes import letter
-from reportlab.pdfgen import canvas
-from reportlab.lib.units import inch
-from io import BytesIO
-import base64
 import logging
-from retrying import retry
+from concurrent.futures import ThreadPoolExecutor
 
 # ==========================
-# Configuration
+# Optimized Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",  # Your custom-trained model
-    "FALLBACK_MODEL": "yolov8n.pt",
+    "MODEL_PATH": "yolov8_safety.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -28,38 +20,11 @@ CONFIG = {
         3: "unsafe_zone",
         4: "improper_tool_use"
     },
-    "CLASS_COLORS": {  # Bounding box 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) # Yellow
-    },
-    "DISPLAY_NAMES": {
-        "no_helmet": "No Helmet",
-        "no_harness": "No Safety Harness",
-        "unsafe_posture": "Unsafe Posture",
-        "unsafe_zone": "Unsafe Zone Entry",
-        "improper_tool_use": "Improper Tool Use"
-    },
-    "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": 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,  # For worker tracking
-    "MIN_VIOLATION_FRAMES": 3  # Min frames to confirm a violation
+    "FRAME_SKIP": 10,  # Increased for faster processing
+    "MAX_PROCESSING_TIME": 60,
+    "CONFIDENCE_THRESHOLD": 0.35,  # Balanced threshold
+    "MIN_VIOLATION_FRAMES": 2,  # Reduced from 3 to 2
+    "GPU_ACCELERATION": True  # Enable GPU if available
 }
 
 # Setup logging
@@ -67,355 +32,113 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() and CONFIG["GPU_ACCELERATION"] else "cpu")
 logger.info(f"Using device: {device}")
 
-# ==========================
-# Load YOLOv8 Model
-# ==========================
-def load_model():
-    try:
-        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 (lower accuracy)")
-        return model
-    except Exception as 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 on frame."""
-    for det in detections:
-        label = det["violation"]
-        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)
-        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):
-    """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
+# Load model
+model = YOLO(CONFIG["MODEL_PATH"]).to(device)
 
-def generate_violation_pdf(violations, score):
-    """Generate PDF report with violations."""
-    try:
-        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)
-        
-        # Report metadata
-        y_pos = 9.8 * inch
-        report_data = [
-            ("Compliance Score", f"{score}%"),
-            ("Total Violations", len(violations)),
-            ("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
-        
-        # Violation details
-        y_pos -= 0.3 * inch
-        c.setFont("Helvetica-Bold", 12)
-        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_pos, "No violations detected.")
-        else:
-            for v in violations:
-                text = (
-                    f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} "
-                    f"at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f})"
-                )
-                c.drawString(1 * inch, y_pos, text)
-                y_pos -= 0.25 * inch
-                if y_pos < 1 * inch:
-                    c.showPage()
-                    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())
-        
-        return pdf_path, f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}", pdf_file
-    except Exception as e:
-        logger.error(f"PDF generation failed: {e}")
-        return None, None, None
-
-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 (Confidence: {v['confidence']:.2f})"
-            for v in violations
-        ) 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": "Pending Review"
-        }
-        record = sf.Safety_Video_Report__c.create(record_data)
-        record_id = record["id"]
-        
-        # Upload PDF if available
-        pdf_url = ""
-        if pdf_file:
-            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
-    except Exception as e:
-        logger.error(f"Salesforce upload failed: {e}")
-        return None, ""
+def process_frame(frame, frame_count, fps):
+    """Process a single frame for violations"""
+    results = model(frame, device=device, verbose=False)  # Disable verbose logging
+    detections = []
+    
+    for result in results:
+        for box in result.boxes:
+            cls = int(box.cls)
+            conf = float(box.conf)
+            if conf > CONFIG["CONFIDENCE_THRESHOLD"] and cls in CONFIG["VIOLATION_LABELS"]:
+                detections.append({
+                    "frame": frame_count,
+                    "violation": CONFIG["VIOLATION_LABELS"][cls],
+                    "confidence": conf,
+                    "bounding_box": box.xywh.cpu().numpy()[0],
+                    "timestamp": frame_count / fps
+                })
+    return detections
 
-# ==========================
-# Video Processing
-# ==========================
 def process_video(video_path):
-    """Analyze video for safety violations."""
-    try:
-        cap = cv2.VideoCapture(video_path)
-        fps = cap.get(cv2.CAP_PROP_FPS) or 30
-        frame_count = 0
-        violations = []
-        snapshots = []
-        workers = []
-        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        
+    """Optimized video processing with parallel frame processing"""
+    start_time = time.time()
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 30
+    violations = []
+    frame_count = 0
+    
+    with ThreadPoolExecutor() as executor:
+        futures = []
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
-            
-            if frame_count % CONFIG["FRAME_SKIP"] != 0:
-                frame_count += 1
-                continue
-            
-            # Run detection
-            results = model(frame, device=device)
-            current_time = frame_count / fps
-            
-            for result in results:
-                for box in result.boxes:
-                    cls = int(box.cls)
-                    conf = float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls)
-                    
-                    if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3):
-                        continue
-                    
-                    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(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
-                    else:
-                        worker_id = len(workers) + 1
-                        workers.append({"id": worker_id, "bbox": bbox})
-                    
-                    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
+                
+            if frame_count % CONFIG["FRAME_SKIP"] == 0:
+                futures.append(executor.submit(process_frame, frame.copy(), frame_count, fps))
             
             frame_count += 1
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.info("Processing time limit reached")
+                break
         
-        cap.release()
-        
-        # 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": filtered_violations,
-            "snapshots": snapshots,
-            "score": safety_score,
-            "message": ""
-        }
-    except Exception as e:
-        logger.error(f"Video processing failed: {e}")
-        return {
-            "violations": [],
-            "snapshots": [],
-            "score": 100,
-            "message": f"Error: {str(e)}"
-        }
+        for future in futures:
+            violations.extend(future.result())
+    
+    cap.release()
+    
+    # Filter violations by frequency
+    violation_counts = {}
+    for v in violations:
+        key = (v["violation"], int(v["timestamp"]))
+        violation_counts[key] = violation_counts.get(key, 0) + 1
+    
+    filtered_violations = [
+        v for v in violations 
+        if violation_counts.get((v["violation"], int(v["timestamp"])), 0) >= CONFIG["MIN_VIOLATION_FRAMES"]
+    ]
+    
+    return filtered_violations, time.time() - start_time
 
-# ==========================
-# Gradio Interface
-# ==========================
 def analyze_video(video_file):
-    """Gradio interface function."""
+    """Optimized Gradio interface function"""
     if not video_file:
-        return "No video uploaded", "", "", "", ""
+        return "No video uploaded", "", "", ""
     
     try:
-        # Process video
-        result = process_video(video_file)
-        if result["message"]:
-            return result["message"], "", "", "", ""
+        start_time = time.time()
+        violations, processing_time = process_video(video_file)
         
-        # 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
+        # Generate simple output (removed PDF generation for speed)
         violation_table = (
-            "| Violation Type         | Timestamp (s) | Confidence | Worker ID |\n"
-            "|------------------------|---------------|------------|-----------|\n" +
+            "| Violation Type | Timestamp (s) | Confidence |\n"
+            "|----------------|---------------|------------|\n" +
             "\n".join(
-                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."
+                f"| {v['violation']:<14} | {v['timestamp']:.1f}        | {v['confidence']:.2f}      |"
+                for v in violations
+            ) if violations else "No violations detected."
         )
         
-        snapshots_md = "\n".join(
-            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 available."
-        
         return (
             violation_table,
-            f"Safety Score: {result['score']}%",
-            snapshots_md,
-            f"Salesforce Record: {record_id or 'N/A'}",
-            sf_url or pdf_url or "N/A"
+            f"Processing Time: {processing_time:.1f}s",
+            f"Violations Found: {len(violations)}",
+            f"Analysis Completed in {time.time()-start_time:.1f}s"
         )
     except Exception as e:
-        return f"Error: {str(e)}", "", "", "", ""
+        return f"Error: {str(e)}", "", "", ""
 
-# Launch Gradio App
+# Simplified Gradio Interface
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
         gr.Markdown("## Detected Violations"),
-        gr.Textbox(label="Safety Score"),
-        gr.Markdown("## Violation Snapshots"),
-        gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Report URL")
+        gr.Textbox(label="Processing Info"),
+        gr.Textbox(label="Summary"),
+        gr.Textbox(label="Status")
     ],
-    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."
-    )
+    title="Fast Safety Compliance Analyzer",
+    description="Optimized version for quick safety violation detection"
 )
 
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface.launch()