Update app.py

app.py

@@ -9,7 +9,6 @@ from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 from reportlab.lib.utils import ImageReader
 import base64
-from PIL import Image
 
 # ==========================
 # Configuration
@@ -47,10 +46,14 @@ except Exception as e:
 # ==========================
 # Video Processing
 # ==========================
-def process_video(video_data, frame_skip=5, max_frames=100):
+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...")
-        # 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)
@@ -94,21 +97,22 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                     }
                     violations.append(violation)
 
-                    # Save snapshot
-                    snapshot_path = os.path.join(OUTPUT_DIR, f"snapshot_{frame_count}_{label}.jpg")
+                    # 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": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/output/{os.path.basename(snapshot_path)}"
+                        "snapshot_url": snapshot_filename  # just filename, local path
                     })
 
             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:
                 print("⏰ Exceeded 30 seconds of processing time.")
                 break
@@ -119,6 +123,12 @@ def process_video(video_data, frame_skip=5, max_frames=100):
         score = calculate_safety_score(violations)
         pdf_base64 = generate_pdf_report(violations, snapshots, score)
 
+        # Clean up snapshot images
+        for snap in snapshots:
+            snap_file = os.path.join(OUTPUT_DIR, snap["snapshot_url"])
+            if os.path.exists(snap_file):
+                os.remove(snap_file)
+
         return {
             "violations": violations,
             "snapshots": snapshots,
@@ -168,37 +178,44 @@ def generate_pdf_report(violations, snapshots, score):
         c.setFont("Helvetica", 12)
         c.drawString(50, height - 80, f"Compliance Score: {score}%")
 
-        # Violations Table
+        # 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
 
-            # Add snapshot if available
+            # 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 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
+            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
+        # 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
+
+        # Clean up the PDF file after encoding
         os.remove(pdf_path)
         print("PDF converted to base64 and file removed")
         return pdf_base64
@@ -234,7 +251,7 @@ interface = gr.Interface(
     outputs=[
         gr.JSON(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
-        gr.Textbox(label="PDF Base64 (for API use)"),
+        gr.File(label="Download PDF Report"),  # Changed from Textbox to File for PDF download
         gr.JSON(label="Snapshots")
     ],
     title="Worksite Safety Violation Analyzer",
@@ -243,4 +260,4 @@ interface = gr.Interface(
 
 if __name__ == "__main__":
     print("🚀 Launching Safety Analyzer App...")
-    interface.launch()
+    interface.launch()