Update app.py

app.py

@@ -2,16 +2,17 @@ import os
 import cv2
 import gradio as gr
 import torch
+import numpy as np
 from ultralytics import YOLO
 import time
-import logging
 from concurrent.futures import ThreadPoolExecutor
+import logging
 
 # ==========================
 # Optimized Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",
+    "MODEL_PATH": "yolov8_safety.pt",  # Your trained model
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -20,11 +21,18 @@ CONFIG = {
         3: "unsafe_zone",
         4: "improper_tool_use"
     },
-    "FRAME_SKIP": 10,  # Increased for faster processing
-    "MAX_PROCESSING_TIME": 60,
-    "CONFIDENCE_THRESHOLD": 0.35,  # Balanced threshold
-    "MIN_VIOLATION_FRAMES": 2,  # Reduced from 3 to 2
-    "GPU_ACCELERATION": True  # Enable GPU if available
+    "CLASS_COLORS": {
+        "no_helmet": (0, 0, 255),      # Red
+        "no_harness": (0, 165, 255),   # Orange
+        "unsafe_posture": (0, 255, 0), # Green
+        "unsafe_zone": (255, 0, 0),    # Blue
+        "improper_tool_use": (255, 255, 0) # Yellow
+    },
+    "FRAME_SKIP": 8,                   # Process every 8th frame
+    "MAX_PROCESSING_TIME": 45,         # Max processing time (seconds)
+    "CONFIDENCE_THRESHOLD": 0.35,      # Balanced threshold
+    "MIN_VIOLATION_FRAMES": 2,         # Reduced from 3 to 2
+    "GPU_ACCELERATION": True           # Enable GPU if available
 }
 
 # Setup logging
@@ -38,9 +46,27 @@ logger.info(f"Using device: {device}")
 # Load model
 model = YOLO(CONFIG["MODEL_PATH"]).to(device)
 
+def draw_detections(frame, detections):
+    """Draw bounding boxes with labels on frame."""
+    for det in detections:
+        label = det["violation"]
+        conf = det["confidence"]
+        x, y, w, h = det["bounding_box"]
+        x1, y1 = int(x - w/2), int(y - h/2)
+        x2, y2 = int(x + w/2), int(y + h/2)
+        
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        cv2.putText(frame, 
+                   f"{label}: {conf:.2f}", 
+                   (x1, y1-10), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 
+                   0.5, color, 2)
+    return frame
+
 def process_frame(frame, frame_count, fps):
-    """Process a single frame for violations"""
-    results = model(frame, device=device, verbose=False)  # Disable verbose logging
+    """Process a single frame for violations (optimized)"""
+    results = model(frame, device=device, verbose=False)
     detections = []
     
     for result in results:
@@ -51,21 +77,22 @@ def process_frame(frame, frame_count, fps):
                 detections.append({
                     "frame": frame_count,
                     "violation": CONFIG["VIOLATION_LABELS"][cls],
-                    "confidence": conf,
+                    "confidence": round(conf, 2),
                     "bounding_box": box.xywh.cpu().numpy()[0],
                     "timestamp": frame_count / fps
                 })
     return detections
 
 def process_video(video_path):
-    """Optimized video processing with parallel frame processing"""
+    """Optimized video processing with parallel execution"""
     start_time = time.time()
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS) or 30
     violations = []
+    snapshots = {}
     frame_count = 0
     
-    with ThreadPoolExecutor() as executor:
+    with ThreadPoolExecutor(max_workers=4) as executor:
         futures = []
         while cap.isOpened():
             ret, frame = cap.read()
@@ -73,15 +100,34 @@ def process_video(video_path):
                 break
                 
             if frame_count % CONFIG["FRAME_SKIP"] == 0:
-                futures.append(executor.submit(process_frame, frame.copy(), frame_count, fps))
+                futures.append(executor.submit(
+                    process_frame, 
+                    frame.copy(), 
+                    frame_count, 
+                    fps
+                ))
             
             frame_count += 1
             if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
                 logger.info("Processing time limit reached")
                 break
         
+        # Process results as they complete
         for future in futures:
-            violations.extend(future.result())
+            frame_detections = future.result()
+            violations.extend(frame_detections)
+            
+            # Capture first occurrence of each violation type
+            for det in frame_detections:
+                if det["violation"] not in snapshots:
+                    snapshots[det["violation"]] = {
+                        "frame": det["frame"],
+                        "timestamp": det["timestamp"],
+                        "image": draw_detections(
+                            cv2.cvtColor(cap.read()[1], cv2.COLOR_BGR2RGB),
+                            [det]
+                        )
+                    }
     
     cap.release()
     
@@ -96,48 +142,94 @@ def process_video(video_path):
         if violation_counts.get((v["violation"], int(v["timestamp"])), 0) >= CONFIG["MIN_VIOLATION_FRAMES"]
     ]
     
-    return filtered_violations, time.time() - start_time
+    # Prepare snapshot outputs
+    snapshot_outputs = []
+    for violation_type, data in snapshots.items():
+        snapshot_path = os.path.join(
+            CONFIG["OUTPUT_DIR"], 
+            f"{violation_type}_{data['frame']}.jpg"
+        )
+        cv2.imwrite(snapshot_path, data["image"])
+        snapshot_outputs.append({
+            "violation": violation_type,
+            "frame": data["frame"],
+            "timestamp": data["timestamp"],
+            "path": snapshot_path
+        })
+    
+    return {
+        "violations": filtered_violations,
+        "snapshots": snapshot_outputs,
+        "processing_time": time.time() - start_time
+    }
+
+def calculate_safety_score(violations):
+    """Calculate safety score (0-100)"""
+    penalty_weights = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
+    }
+    unique_violations = set((v["violation"]) for v in violations)
+    total_penalty = sum(penalty_weights.get(v, 0) for v in unique_violations)
+    return max(100 - total_penalty, 0)
+
+def format_output(result):
+    """Format results for Gradio output"""
+    # Violation table
+    violation_table = (
+        "| Violation Type | Timestamp (s) | Confidence |\n"
+        "|----------------|---------------|------------|\n" +
+        "\n".join(
+            f"| {v['violation']:<14} | {v['timestamp']:.1f}        | {v['confidence']:.2f}      |"
+            for v in result["violations"]
+        ) if result["violations"] else "No violations detected."
+    )
+    
+    # Snapshots
+    snapshots_md = "\n".join(
+        f"**{s['violation']}** at {s['timestamp']:.1f}s: "
+        f"![](file/{s['path']})"
+        for s in result["snapshots"]
+    ) if result["snapshots"] else "No snapshots available."
+    
+    # Safety score
+    safety_score = calculate_safety_score(result["violations"])
+    
+    return (
+        violation_table,
+        f"Safety Score: {safety_score}%",
+        snapshots_md,
+        f"Processed in {result['processing_time']:.1f}s"
+    )
 
 def analyze_video(video_file):
-    """Optimized Gradio interface function"""
+    """Gradio interface function"""
     if not video_file:
         return "No video uploaded", "", "", ""
     
     try:
-        start_time = time.time()
-        violations, processing_time = process_video(video_file)
-        
-        # Generate simple output (removed PDF generation for speed)
-        violation_table = (
-            "| Violation Type | Timestamp (s) | Confidence |\n"
-            "|----------------|---------------|------------|\n" +
-            "\n".join(
-                f"| {v['violation']:<14} | {v['timestamp']:.1f}        | {v['confidence']:.2f}      |"
-                for v in violations
-            ) if violations else "No violations detected."
-        )
-        
-        return (
-            violation_table,
-            f"Processing Time: {processing_time:.1f}s",
-            f"Violations Found: {len(violations)}",
-            f"Analysis Completed in {time.time()-start_time:.1f}s"
-        )
+        result = process_video(video_file)
+        return format_output(result)
     except Exception as e:
+        logger.error(f"Error: {str(e)}")
         return f"Error: {str(e)}", "", "", ""
 
-# Simplified Gradio Interface
+# Gradio Interface
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
         gr.Markdown("## Detected Violations"),
-        gr.Textbox(label="Processing Info"),
-        gr.Textbox(label="Summary"),
-        gr.Textbox(label="Status")
+        gr.Textbox(label="Safety Score"),
+        gr.Markdown("## Violation Snapshots"),
+        gr.Textbox(label="Processing Info")
     ],
-    title="Fast Safety Compliance Analyzer",
-    description="Optimized version for quick safety violation detection"
+    title="AI Safety Compliance Analyzer",
+    description="Optimized for fast detection of safety violations",
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":