Update app.py

app.py

@@ -2,7 +2,6 @@ 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
@@ -47,11 +46,6 @@ except Exception as e:
 # Video Processing
 # ==========================
 def process_video(video_data, frame_skip=10, max_frames=100):
-    """
-    Processes uploaded video data to detect safety violations using YOLO model.
-    frame_skip: number of frames to skip between inferences to speed up processing.
-    max_frames: max frames to process before stopping.
-    """
     try:
         print("Processing video data...")
         video_path = os.path.join(OUTPUT_DIR, f"temp_{int(time.time())}.mp4")
@@ -78,7 +72,6 @@ def process_video(video_data, frame_skip=10, max_frames=100):
                 frame_count += 1
                 continue
 
-            # Model inference
             results = model(frame, device=device)
 
             for result in results:
@@ -97,20 +90,18 @@ def process_video(video_data, frame_skip=10, max_frames=100):
                     }
                     violations.append(violation)
 
-                    # Save snapshot locally with a filename pattern
                     snapshot_filename = f"snapshot_{frame_count}_{label}.jpg"
                     snapshot_path = os.path.join(OUTPUT_DIR, snapshot_filename)
                     cv2.imwrite(snapshot_path, frame)
                     snapshots.append({
                         "violation": label,
                         "frame": frame_count,
-                        "snapshot_url": snapshot_filename  # just filename, local path
+                        "snapshot_url": snapshot_filename
                     })
 
             frame_count += 1
             processed_frame_count += 1
 
-            # Stop if max frames or 30 seconds elapsed
             if processed_frame_count >= max_frames:
                 break
             if time.time() - start_time > 30:
@@ -121,9 +112,9 @@ def process_video(video_data, frame_skip=10, max_frames=100):
         os.remove(video_path)
 
         score = calculate_safety_score(violations)
-        pdf_base64 = generate_pdf_report(violations, snapshots, score)
+        pdf_url = generate_pdf_report(violations, snapshots, score)
 
-        # Clean up snapshot images
+        # Clean up snapshots
         for snap in snapshots:
             snap_file = os.path.join(OUTPUT_DIR, snap["snapshot_url"])
             if os.path.exists(snap_file):
@@ -133,7 +124,7 @@ def process_video(video_data, frame_skip=10, max_frames=100):
             "violations": violations,
             "snapshots": snapshots,
             "score": score,
-            "pdf_base64": pdf_base64
+            "pdf_url": pdf_url
         }
 
     except Exception as e:
@@ -142,7 +133,7 @@ def process_video(video_data, frame_skip=10, max_frames=100):
             "violations": [],
             "snapshots": [],
             "score": 0,
-            "pdf_base64": "",
+            "pdf_url": "",
             "error": str(e)
         }
 
@@ -166,98 +157,17 @@ def calculate_safety_score(violations):
 # ==========================
 def generate_pdf_report(violations, snapshots, score):
     try:
-        pdf_path = os.path.join(OUTPUT_DIR, f"report_{int(time.time())}.pdf")
+        timestamp = int(time.time())
+        pdf_filename = f"report_{timestamp}.pdf"
+        pdf_path = os.path.join(OUTPUT_DIR, pdf_filename)
         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 header
         y = height - 120
         c.setFont("Helvetica-Bold", 12)
-        c.drawString(50, y, "Detected Violations:")
-        y -= 20
-
-        for v in violations:
-            if y < 150:
-                c.showPage()
-                y = height - 50
-
-            c.setFont("Helvetica", 10)
-            text = f"Violation: {v['violation']}, Timestamp: {v['timestamp']:.2f}s, Confidence: {v['confidence']}"
-            c.drawString(50, y, text)
-            y -= 20
-
-            # Find matching snapshot by frame and violation label
-            snapshot = next((s for s in snapshots if s["frame"] == v["frame"] and s["violation"] == v["violation"]), None)
-            if snapshot:
-                snapshot_file = os.path.join(OUTPUT_DIR, snapshot["snapshot_url"])
-                if os.path.exists(snapshot_file):
-                    img = ImageReader(snapshot_file)
-                    img_width = 200
-                    img_height = 150
-                    if y - img_height < 50:
-                        c.showPage()
-                        y = height - 50
-                    c.drawImage(img, 50, y - img_height, width=img_width, height=img_height)
-                    y -= img_height + 20
-
-        c.save()
-        print(f"PDF generated at {pdf_path}")
-
-        # Convert PDF to base64 for returning to frontend
-        with open(pdf_path, "rb") as f:
-            pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
-
-        # Clean up the PDF file after encoding
-        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.File(label="Download PDF Report"),  # Changed from Textbox to File for PDF download
-        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()