Update app.py

app.py

@@ -1,127 +1,246 @@
-import { LightningElement, api } from 'lwc';
-import uploadVideo from '@salesforce/apex/VideoFileUpload.uploadVideoFile';
-import { ShowToastEvent } from 'lightning/platformShowToastEvent';
-
-export default class VideoUploadWithAI extends LightningElement {
-    @api recordId;
-    videoFile;
-    base64FileData;
-    fileName;
-    complianceScore;
-    violationsFound = false;
-    violationsList = [];
-    snapshots = [];
-    pdfReportUrl;
-    isLoading = false;
-    isDragging = false;
-
-    handleFileChange(event) {
-        const file = event.target.files[0];
-        if (file) {
-            this.fileName = file.name;
-            this.convertToBase64(file);
-        } else {
-            this.showToast('Error', 'No video file selected.', 'error');
-        }
-    }
+import os
+import cv2
+import gradio as gr
+import torch
+import numpy as np
+from ultralytics import YOLO
+import time
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.utils import ImageReader
+import base64
+from PIL import Image
 
-    convertToBase64(file) {
-        const reader = new FileReader();
-        reader.onloadend = () => {
-            this.base64FileData = reader.result.split(',')[1];
-        };
-        reader.readAsDataURL(file);
-    }
+# ==========================
+# Configuration
+# ==========================
+DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
+FALLBACK_MODEL = "yolov8n.pt"
+MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
+OUTPUT_DIR = "output"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-    handleDragOver(event) {
-        event.preventDefault();
-        this.isDragging = true;
-    }
+VIOLATION_LABELS = {
+    0: "no_helmet",
+    1: "no_harness",
+    2: "unsafe_posture",
+    3: "unsafe_zone"
+}
 
-    handleDragEnter(event) {
-        event.preventDefault();
-        this.isDragging = true;
-    }
+# ==========================
+# Device Setup
+# ==========================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"✅ Using device: {device}")
 
-    handleDragLeave(event) {
-        event.preventDefault();
-        this.isDragging = false;
-    }
+# ==========================
+# Load Model
+# ==========================
+try:
+    selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
+    model = YOLO(selected_model)
+    print(f"✅ Model loaded: {selected_model}")
+except Exception as e:
+    print(f"❌ Failed to load model: {e}")
+    raise
 
-    handleDrop(event) {
-        event.preventDefault();
-        this.isDragging = false;
-        const file = event.dataTransfer.files[0];
-        if (file && file.type.startsWith('video/')) {
-            this.fileName = file.name;
-            this.convertToBase64(file);
-        } else {
-            this.showToast('Error', 'Please drop a valid video file.', 'error');
-        }
-    }
+# ==========================
+# Video Processing
+# ==========================
+def process_video(video_data, frame_skip=5, max_frames=100):
+    try:
+        print("Processing video data...")
+        # Save uploaded video data to a temporary file
+        video_path = os.path.join(OUTPUT_DIR, f"temp_{int(time.time())}.mp4")
+        with open(video_path, "wb") as f:
+            f.write(video_data)
+        print(f"Video saved to {video_path}")
 
-    get dropzoneClass() {
-        return this.isDragging
-            ? 'slds-box slds-box_small dropzone dragging'
-            : 'slds-box slds-box_small dropzone';
-    }
+        video = cv2.VideoCapture(video_path)
+        if not video.isOpened():
+            raise ValueError("Could not open video file.")
 
-    async handleSubmit() {
-        if (!this.base64FileData) {
-            this.showToast('Error', 'Please select a video file.', 'error');
-            return;
-        }
+        frame_count = 0
+        violations = []
+        snapshots = []
+        processed_frame_count = 0
+        start_time = time.time()
 
-        if (!this.recordId) {
-            this.showToast('Error', 'No record ID provided. Please ensure this component is placed on a record page.', 'error');
-            return;
-        }
+        while True:
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            if frame_count % frame_skip != 0:
+                frame_count += 1
+                continue
+
+            # Model inference
+            results = model(frame, device=device)
+
+            for result in results:
+                for box in result.boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    xywh = box.xywh.cpu().numpy()[0]
+
+                    label = VIOLATION_LABELS.get(cls, f"class_{cls}")
+                    violation = {
+                        "frame": frame_count,
+                        "violation": label,
+                        "confidence": round(conf, 2),
+                        "bounding_box": [round(x, 2) for x in xywh],
+                        "timestamp": frame_count / video.get(cv2.CAP_PROP_FPS)
+                    }
+                    violations.append(violation)
+
+                    # Save snapshot
+                    snapshot_path = os.path.join(OUTPUT_DIR, f"snapshot_{frame_count}_{label}.jpg")
+                    cv2.imwrite(snapshot_path, frame)
+                    snapshots.append({
+                        "violation": label,
+                        "frame": frame_count,
+                        "snapshot_url": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/output/{os.path.basename(snapshot_path)}"
+                    })
+
+            frame_count += 1
+            processed_frame_count += 1
+
+            if processed_frame_count >= max_frames:
+                break
+
+            if time.time() - start_time > 30:
+                print("⏰ Exceeded 30 seconds of processing time.")
+                break
 
-        this.isLoading = true;
-        try {
-            const result = await uploadVideo({
-                reportId: this.recordId,
-                fileData: this.base64FileData,
-                fileName: this.fileName
-            });
-            console.log('Apex Result:', JSON.stringify(result));
-
-            if (result.status === 'Reviewed') {
-                this.complianceScore = result.complianceScore ? parseFloat(result.complianceScore) : 0;
-                this.violationsList = result.violationsList || [];
-                this.violationsFound = this.violationsList.length > 0;
-                this.snapshots = result.snapshots || [];
-                this.pdfReportUrl = result.pdfReportUrl || '';
-                this.showToast('Success', 'Video processed successfully.', 'success');
-            } else {
-                let errorMessage = result.status || 'Unknown error occurred.';
-                if (errorMessage.includes('INSUFFICIENT_ACCESS') || errorMessage.includes('PERMISSION')) {
-                    errorMessage += ' - Please contact your Salesforce admin to grant necessary permissions.';
-                }
-                this.showToast('Error', errorMessage, 'error');
-            }
-        } catch (error) {
-            console.error('Error in handleSubmit:', error);
-            let errorMessage = 'Failed to process video: ' + (error.message || 'Unknown error.');
-            if (error.message.includes('INSUFFICIENT_ACCESS') || error.message.includes('PERMISSION')) {
-                errorMessage += ' - Please contact your Salesforce admin to grant necessary permissions.';
-            }
-            this.showToast('Error', errorMessage, 'error');
-        } finally {
-            this.isLoading = false;
+        video.release()
+        os.remove(video_path)
+
+        score = calculate_safety_score(violations)
+        pdf_base64 = generate_pdf_report(violations, snapshots, score)
+
+        return {
+            "violations": violations,
+            "snapshots": snapshots,
+            "score": score,
+            "pdf_base64": pdf_base64
         }
-    }
 
-    handleDownloadReport() {
-        if (this.pdfReportUrl) {
-            window.open(this.pdfReportUrl, '_blank');
-        } else {
-            this.showToast('Error', 'No PDF report available to download.', 'error');
+    except Exception as e:
+        print(f"❌ Error processing video: {e}")
+        return {
+            "violations": [],
+            "snapshots": [],
+            "score": 0,
+            "pdf_base64": "",
+            "error": str(e)
         }
-    }
 
-    showToast(title, message, variant) {
-        const event = new ShowToastEvent({ title, message, variant });
-        this.dispatchEvent(event);
+# ==========================
+# Safety Score Calculation
+# ==========================
+def calculate_safety_score(violations):
+    base_score = 100
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 25
     }
-}
+    for v in violations:
+        base_score -= penalties.get(v["violation"], 0)
+    return max(base_score, 0)
+
+# ==========================
+# PDF Report Generation
+# ==========================
+def generate_pdf_report(violations, snapshots, score):
+    try:
+        pdf_path = os.path.join(OUTPUT_DIR, f"report_{int(time.time())}.pdf")
+        c = canvas.Canvas(pdf_path, pagesize=letter)
+        width, height = letter
+
+        # Title
+        c.setFont("Helvetica-Bold", 16)
+        c.drawString(50, height - 50, "Worksite Safety Compliance Report")
+
+        # Compliance Score
+        c.setFont("Helvetica", 12)
+        c.drawString(50, height - 80, f"Compliance Score: {score}%")
+
+        # Violations Table
+        y = height - 120
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(50, y, "Detected Violations:")
+        y -= 20
+
+        for v in violations:
+            c.setFont("Helvetica", 10)
+            text = f"Violation: {v['violation']}, Timestamp: {v['timestamp']:.2f}s, Confidence: {v['confidence']}"
+            c.drawString(50, y, text)
+            y -= 20
+
+            # Add snapshot if available
+            snapshot = next((s for s in snapshots if s["frame"] == v["frame"] and s["violation"] == v["violation"]), None)
+            if snapshot and os.path.exists(snapshot["snapshot_url"].split('/')[-1]):
+                img = ImageReader(snapshot["snapshot_url"].split('/')[-1])
+                c.drawImage(img, 50, y - 100, width=200, height=150)
+                y -= 170
+
+            if y < 50:
+                c.showPage()
+                y = height - 50
+
+        c.save()
+        print(f"PDF generated at {pdf_path}")
+
+        # Convert PDF to base64
+        with open(pdf_path, "rb") as f:
+            pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
+        
+        # Clean up
+        os.remove(pdf_path)
+        print("PDF converted to base64 and file removed")
+        return pdf_base64
+    except Exception as e:
+        print(f"❌ Error generating PDF: {e}")
+        return ""
+
+# ==========================
+# Gradio Interface
+# ==========================
+def gradio_interface(video_file):
+    try:
+        if not video_file:
+            return {"error": "Please upload a video file."}, "", "", []
+
+        with open(video_file, "rb") as f:
+            video_data = f.read()
+
+        result = process_video(video_data)
+        return (
+            result["violations"],
+            f"Safety Score: {result['score']}%",
+            result["pdf_base64"],
+            result["snapshots"]
+        )
+    except Exception as e:
+        print(f"❌ Error in gradio_interface: {e}")
+        return {"error": str(e)}, "", "", []
+
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Video(label="Upload Site Video"),
+    outputs=[
+        gr.JSON(label="Detected Safety Violations"),
+        gr.Textbox(label="Compliance Score"),
+        gr.Textbox(label="PDF Base64 (for API use)"),
+        gr.JSON(label="Snapshots")
+    ],
+    title="Worksite Safety Violation Analyzer",
+    description="Upload short site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
+)
+
+if __name__ == "__main__":
+    print("🚀 Launching Safety Analyzer App...")
+    interface.launch()