Update app.py

app.py

@@ -12,263 +12,235 @@ 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=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
+                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)}"
+                    })
 
-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
-            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)
-    }
-```
+            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)
+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
 
-        # 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 ""
-```
+        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 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)}, "", "", []
-```
+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)."
+    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()
-
+if __name__ == "__main__":
+    print("🚀 Launching Safety Analyzer App...")
+    interface.launch()