Update app.py

app.py

@@ -18,7 +18,7 @@ from PIL import Image
 DEFAULT_MODEL_PATH = "models/yolov8_safety.pt"
 FALLBACK_MODEL = "yolov8n.pt"
 MODEL_PATH = os.getenv("SAFETY_MODEL_PATH", DEFAULT_MODEL_PATH)
-OUTPUT_DIR = "output"  # Directory to store snapshots and PDFs
+OUTPUT_DIR = "output"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 VIOLATION_LABELS = {
@@ -37,18 +37,25 @@ print(f"✅ Using device: {device}")
 # ==========================
 # Load Model
 # ==========================
-selected_model = MODEL_PATH if os.path.isfile(MODEL_PATH) else FALLBACK_MODEL
-model = YOLO(selected_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():
@@ -94,7 +101,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                     snapshots.append({
                         "violation": label,
                         "frame": frame_count,
-                        "snapshot_url": snapshot_path
+                        "snapshot_url": f"https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/output/{os.path.basename(snapshot_path)}"
                     })
 
             frame_count += 1
@@ -108,7 +115,7 @@ def process_video(video_data, frame_skip=5, max_frames=100):
                 break
 
         video.release()
-        os.remove(video_path)  # Clean up temporary video file
+        os.remove(video_path)
 
         score = calculate_safety_score(violations)
         pdf_report_path = generate_pdf_report(violations, snapshots, score)
@@ -149,61 +156,74 @@ def calculate_safety_score(violations):
 # PDF Report Generation
 # ==========================
 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)
+    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
 
-        # 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
+        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):
-    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"]
-    )
+    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,