Update app.py

app.py

@@ -1,243 +1,127 @@
-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
-from io import BytesIO
-import base64
-from PIL import Image
-
-# ==========================
-# 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)
-
-VIOLATION_LABELS = {
-    0: "no_helmet",
-    1: "no_harness",
-    2: "unsafe_posture",
-    3: "unsafe_zone"
-}
-
-# ==========================
-# Device Setup
-# ==========================
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"✅ Using device: {device}")
-
-# ==========================
-# 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
-
-# ==========================
-# 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}")
-
-        video = cv2.VideoCapture(video_path)
-        if not video.isOpened():
-            raise ValueError("Could not open video file.")
-
-        frame_count = 0
-        violations = []
-        snapshots = []
-        processed_frame_count = 0
-        start_time = time.time()
-
-        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)
+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');
+        }
+    }
 
-                    # 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)}"
-                    })
+    convertToBase64(file) {
+        const reader = new FileReader();
+        reader.onloadend = () => {
+            this.base64FileData = reader.result.split(',')[1];
+        };
+        reader.readAsDataURL(file);
+    }
 
-            frame_count += 1
-            processed_frame_count += 1
+    handleDragOver(event) {
+        event.preventDefault();
+        this.isDragging = true;
+    }
 
-            if processed_frame_count >= max_frames:
-                break
+    handleDragEnter(event) {
+        event.preventDefault();
+        this.isDragging = true;
+    }
 
-            if time.time() - start_time > 30:
-                print("⏰ Exceeded 30 seconds of processing time.")
-                break
+    handleDragLeave(event) {
+        event.preventDefault();
+        this.isDragging = false;
+    }
 
-        video.release()
-        os.remove(video_path)
+    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');
+        }
+    }
 
-        score = calculate_safety_score(violations)
-        pdf_report_path = generate_pdf_report(violations, snapshots, score)
+    get dropzoneClass() {
+        return this.isDragging
+            ? 'slds-box slds-box_small dropzone dragging'
+            : 'slds-box slds-box_small dropzone';
+    }
 
-        return {
-            "violations": violations,
-            "snapshots": snapshots,
-            "score": score,
-            "pdf_report_url": pdf_report_path
+    async handleSubmit() {
+        if (!this.base64FileData) {
+            this.showToast('Error', 'Please select a video file.', 'error');
+            return;
         }
 
-    except Exception as e:
-        print(f"❌ Error processing video: {e}")
-        return {
-            "violations": [],
-            "snapshots": [],
-            "score": 0,
-            "pdf_report_url": "",
-            "error": str(e)
+        if (!this.recordId) {
+            this.showToast('Error', 'No record ID provided. Please ensure this component is placed on a record page.', 'error');
+            return;
         }
 
-# ==========================
-# Safety Score Calculation
-# ==========================
-def calculate_safety_score(violations):
-    base_score = 100
-    penalties = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 25
+        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;
+        }
     }
-    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}")
-        # Return a publicly accessible URL
-        base_url = "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1"
-        pdf_url = f"{base_url}/output/{os.path.basename(pdf_path)}"
-        print(f"PDF URL: {pdf_url}")
-        return pdf_url
-    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_report_url"],
-            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 Report URL"),
-        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)."
-)
+    handleDownloadReport() {
+        if (this.pdfReportUrl) {
+            window.open(this.pdfReportUrl, '_blank');
+        } else {
+            this.showToast('Error', 'No PDF report available to download.', 'error');
+        }
+    }
 
-if __name__ == "__main__":
-    print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    showToast(title, message, variant) {
+        const event = new ShowToastEvent({ title, message, variant });
+        this.dispatchEvent(event);
+    }
+}