Update app.py

app.py

@@ -15,10 +15,11 @@ 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": {
         0: "no_helmet",
@@ -27,18 +28,38 @@ CONFIG = {
         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)
+    "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
     },
-    "FRAME_SKIP": 10,  # Increased from 5 to 10 for faster processing
-    "MAX_PROCESSING_TIME": 30,  # Reduced from 60 to 30 seconds
-    "CONFIDENCE_THRESHOLD": 0.35,  # Slightly increased for faster filtering
-    "IOU_THRESHOLD": 0.5,
-    "MIN_VIOLATION_FRAMES": 2  # Reduced from 3 to 2 frames
+    "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
 }
 
 # Setup logging
@@ -49,45 +70,188 @@ 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 model with caching
-model = None
+# ==========================
+# Load YOLOv8 Model
+# ==========================
 def load_model():
-    global model
-    if model is None:
-        model = YOLO(CONFIG["MODEL_PATH"]).to(device)
-    return 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
+
+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, ""
 
 # ==========================
-# Optimized Processing Functions
+# Video Processing
 # ==========================
 def process_video(video_path):
-    """Optimized video processing with early termination"""
+    """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()}
-        start_time = time.time()
         
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
             
-            # Skip frames for faster processing
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
             
-            # Early termination if processing takes too long
-            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
-                logger.info("Processing time limit reached")
-                break
-            
-            # Run detection on current frame
-            results = load_model()(frame, device=device, verbose=False)  # Disable verbose logging
+            # Run detection
+            results = model(frame, device=device)
+            current_time = frame_count / fps
             
             for result in results:
                 for box in result.boxes:
@@ -95,7 +259,7 @@ def process_video(video_path):
                     conf = float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls)
                     
-                    if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                    if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3):
                         continue
                     
                     bbox = box.xywh.cpu().numpy()[0]
@@ -104,23 +268,26 @@ def process_video(video_path):
                         "violation": label,
                         "confidence": conf,
                         "bounding_box": bbox,
-                        "timestamp": frame_count / fps
+                        "timestamp": current_time
                     }
                     
-                    # Simplified worker tracking
-                    matched = False
+                    # Track worker
+                    matched_worker = None
+                    max_iou = 0
                     for worker in workers:
-                        if calculate_iou(worker["bbox"], bbox) > CONFIG["IOU_THRESHOLD"]:
-                            worker["bbox"] = bbox
-                            detection["worker_id"] = worker["id"]
-                            matched = True
-                            break
+                        iou = calculate_iou(worker["bbox"], bbox)
+                        if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                            max_iou = iou
+                            matched_worker = worker
                     
-                    if not matched:
+                    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
                     
+                    detection["worker_id"] = worker_id
                     violations.append(detection)
                     
                     # Capture snapshot if first detection of this type
@@ -129,7 +296,13 @@ def process_video(video_path):
                             CONFIG["OUTPUT_DIR"],
                             f"{label}_{frame_count}.jpg"
                         )
-                        cv2.imwrite(snapshot_path, frame)
+                        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
@@ -147,9 +320,22 @@ def process_video(video_path):
             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": [k for k, v in snapshot_taken.items() if v],
+            "snapshots": snapshots,
+            "score": safety_score,
             "message": ""
         }
     except Exception as e:
@@ -157,68 +343,79 @@ def process_video(video_path):
         return {
             "violations": [],
             "snapshots": [],
+            "score": 100,
             "message": f"Error: {str(e)}"
         }
 
 # ==========================
-# Fast Gradio Interface
+# Gradio Interface
 # ==========================
 def analyze_video(video_file):
-    """Optimized interface with progress updates"""
+    """Gradio interface function."""
     if not video_file:
         return "No video uploaded", "", "", "", ""
     
-    # Immediate feedback
-    yield "Processing started...", "", "", "", ""
-    
     try:
-        # Process video in background
+        # Process video
         result = process_video(video_file)
-        
         if result["message"]:
-            yield result["message"], "", "", "", ""
-            return
+            return result["message"], "", "", "", ""
+        
+        # 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
+        )
         
-        # Generate simplified output
+        # Format outputs
         violation_table = (
-            "| Violation Type | Timestamp | Confidence | Worker ID |\n"
-            "|----------------|-----------|------------|-----------|\n" +
+            "| Violation Type         | Timestamp (s) | Confidence | Worker ID |\n"
+            "|------------------------|---------------|------------|-----------|\n" +
             "\n".join(
-                f"| {v['violation']} | {v['timestamp']:.1f}s | {v['confidence']:.2f} | {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."
         )
         
         snapshots_md = "\n".join(
-            f"**{violation}** detected" 
-            for violation in result["snapshots"]
+            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."
         
-        yield (
+        return (
             violation_table,
-            f"Analysis complete in {time.time() - start_time:.1f}s",
+            f"Safety Score: {result['score']}%",
             snapshots_md,
-            "Salesforce integration placeholder",
-            "Report URL placeholder"
+            f"Salesforce Record: {record_id or 'N/A'}",
+            sf_url or pdf_url or "N/A"
         )
     except Exception as e:
-        yield f"Error: {str(e)}", "", "", "", ""
+        return f"Error: {str(e)}", "", "", "", ""
 
-# Launch optimized interface
+# Launch Gradio App
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
         gr.Markdown("## Detected Violations"),
-        gr.Textbox(label="Processing Time"),
+        gr.Textbox(label="Safety Score"),
         gr.Markdown("## Violation Snapshots"),
-        gr.Textbox(label="Salesforce Record"),
+        gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Report URL")
     ],
-    title="Fast Safety Compliance Analyzer",
-    description="Optimized for quick safety violation detection in worksite videos.",
-    allow_flagging="never"
+    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__":
-    interface.launch()
+    interface.launch(share=True)