Update app.py

app.py

@@ -1,173 +1,275 @@
+Here is my app.py code
 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
 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"
+
+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")
+
+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}")
+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
+print(f"❌ Failed to load model: {e}")
+raise
 
 # ==========================
+
 # Video Processing
+
 # ==========================
-def process_video(video_data, frame_skip=10, max_frames=100):
-    try:
-        print("Processing video data...")
-        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
-
-            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)
-
-                    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
-                    })
 
+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
-            processed_frame_count += 1
-
-            if processed_frame_count >= max_frames:
-                break
-            if time.time() - start_time > 30:
-                print("⏰ Exceeded 30 seconds of processing time.")
-                break
-
-        video.release()
-        os.remove(video_path)
-
-        score = calculate_safety_score(violations)
-        pdf_url = generate_pdf_report(violations, snapshots, score)
-
-        # Clean up snapshots
-        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,
-            "score": score,
-            "pdf_url": pdf_url
-        }
-
-    except Exception as e:
-        print(f"❌ Error processing video: {e}")
-        return {
-            "violations": [],
-            "snapshots": [],
-            "score": 0,
-            "pdf_url": "",
-            "error": str(e)
-        }
+            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)
+
+                # 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)}"
+                })
+
+        frame_count += 1
+        processed_frame_count += 1
+
+        if processed_frame_count >= max_frames:
+            break
+
+        if time.time() - start_time > 30:
+            print("⏰ Exceeded 30 seconds of processing time.")
+            break
+
+    video.release()
+    os.remove(video_path)
+
+    score = calculate_safety_score(violations)
+    pdf_base64 = generate_pdf_report(violations, snapshots, score)
+
+    return {
+        "violations": violations,
+        "snapshots": snapshots,
+        "score": score,
+        "pdf_base64": pdf_base64
+    }
+
+except Exception as e:
+    print(f"❌ Error processing video: {e}")
+    return {
+        "violations": [],
+        "snapshots": [],
+        "score": 0,
+        "pdf_base64": "",
+        "error": str(e)
+    }
+```
 
 # ==========================
+
 # Safety Score Calculation
+
 # ==========================
-def calculate_safety_score(violations):
-    base_score = 100
-    penalties = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 25
-    }
-    for v in violations:
-        base_score -= penalties.get(v["violation"], 0)
-    return max(base_score, 0)
+
+def calculate\_safety\_score(violations):
+base\_score = 100
+penalties = {
+"no\_helmet": 25,
+"no\_harness": 30,
+"unsafe\_posture": 20,
+"unsafe\_zone": 25
+}
+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}")
+
+    # Convert PDF to base64
+    with open(pdf_path, "rb") as f:
+        pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
+    
+    # Clean up
+    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 generate_pdf_report(violations, snapshots, score):
-    try:
-        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
-
-        c.setFont("Helvetica-Bold", 16)
-        c.drawString(50, height - 50, "Worksite Safety Compliance Report")
-
-        c.setFont("Helvetica", 12)
-        c.drawString(50, height - 80, f"Compliance Score: {score}%")
-
-        y = height - 120
-        c.setFont("Helvetica-Bold", 12)
+
+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.Textbox(label="PDF Base64 (for API use)"),
+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()
+