Update app.py

app.py

@@ -5,13 +5,21 @@ 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"  # Use nano model if custom one is missing
+FALLBACK_MODEL = "yolov8n.pt"
 MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
+OUTPUT_DIR = "output"  # Directory to store snapshots and PDFs
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 VIOLATION_LABELS = {
     0: "no_helmet",
@@ -27,22 +35,28 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"✅ Using device: {device}")
 
 # ==========================
-# Load Model (Use YOLOv8n for Faster Inference)
+# Load Model
 # ==========================
 selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
 model = YOLO(selected_model)
 
 # ==========================
-# Video Processing with Optimizations
+# Video Processing
 # ==========================
-def process_video(video_path, frame_skip=5, max_frames=100):
+def process_video(video_data, frame_skip=5, max_frames=100):
     try:
+        # 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)
+
         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()
 
@@ -55,7 +69,7 @@ def process_video(video_path, frame_skip=5, max_frames=100):
                 frame_count += 1
                 continue
 
-            # Model inference for detecting violations
+            # Model inference
             results = model(frame, device=device)
 
             for result in results:
@@ -65,11 +79,22 @@ def process_video(video_path, frame_skip=5, max_frames=100):
                     xywh = box.xywh.cpu().numpy()[0]
 
                     label = VIOLATION_LABELS.get(cls, f"class_{cls}")
-                    violations.append({
+                    violation = {
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in xywh]
+                        "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": snapshot_path
                     })
 
             frame_count += 1
@@ -78,22 +103,32 @@ def process_video(video_path, frame_skip=5, max_frames=100):
             if processed_frame_count >= max_frames:
                 break
 
-            elapsed_time = time.time() - start_time
-            if elapsed_time > 30:
+            if time.time() - start_time > 30:
                 print("⏰ Exceeded 30 seconds of processing time.")
                 break
 
         video.release()
-        score = calculate_safety_score(violations)
+        os.remove(video_path)  # Clean up temporary video file
 
-        # Generate the PDF report URL (using an external method or library)
-        pdf_report_url = generate_pdf_report(violations, score)
+        score = calculate_safety_score(violations)
+        pdf_report_path = generate_pdf_report(violations, snapshots, score)
 
-        return violations, score, pdf_report_url
+        return {
+            "violations": violations,
+            "snapshots": snapshots,
+            "score": score,
+            "pdf_report_url": pdf_report_path
+        }
 
     except Exception as e:
         print(f"❌ Error processing video: {e}")
-        return [], f"Error: {e}", None
+        return {
+            "violations": [],
+            "snapshots": [],
+            "score": 0,
+            "pdf_report_url": "",
+            "error": str(e)
+        }
 
 # ==========================
 # Safety Score Calculation
@@ -113,29 +148,76 @@ def calculate_safety_score(violations):
 # ==========================
 # PDF Report Generation
 # ==========================
-def generate_pdf_report(violations, score):
-    # Create a PDF with violation details (This can be done using an external library or template)
-    pdf_url = "http://path_to_pdf_report"  # URL to the generated PDF (replace with actual URL)
-    return pdf_url
+def generate_pdf_report(violations, snapshots, score):
+    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"]):
+            img = ImageReader(snapshot["snapshot_url"])
+            c.drawImage(img, 50, y - 100, width=200, height=150)
+            y -= 170
+
+        if y < 50:
+            c.showPage()
+            y = height - 50
+
+    c.save()
+    return pdf_path
 
 # ==========================
 # Gradio Interface
 # ==========================
 def gradio_interface(video_file):
     if not video_file:
-        return "Please upload a video file.", ""
+        return {"error": "Please upload a video file."}, "", ""
+
+    with open(video_file, "rb") as f:
+        video_data = f.read()
 
-    violations, score, pdf_url = process_video(video_file)
-    return violations, f"Safety Score: {score}%", pdf_url
+    result = process_video(video_data)
+    return (
+        result["violations"],
+        f"Safety Score: {result['score']}%",
+        result["pdf_report_url"],
+        result["snapshots"]
+    )
 
 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")],
+    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)."
 )
 
 if __name__ == "__main__":
     print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()