Update app.py

app.py

@@ -5,14 +5,20 @@ import torch
 import numpy as np
 from ultralytics import YOLO
 import time
-from concurrent.futures import ThreadPoolExecutor
+from simple_salesforce import Salesforce
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.units import inch
+from io import BytesIO
+import base64
 import logging
+from retrying import retry
 
 # ==========================
 # Optimized Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",  # Your trained model
+    "MODEL_PATH": "yolov8_safety.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
@@ -22,17 +28,17 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "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
+        "no_helmet": (0, 0, 255),
+        "no_harness": (0, 165, 255),
+        "unsafe_posture": (0, 255, 0),
+        "unsafe_zone": (255, 0, 0),
+        "improper_tool_use": (255, 255, 0)
     },
-    "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
+    "FRAME_SKIP": 10,  # Increased from 5 to 10 for faster processing
+    "MAX_PROCESSING_TIME": 30,  # Reduced from 60 to 30 seconds
+    "CONFIDENCE_THRESHOLD": 0.35,  # Slightly increased for faster filtering
+    "IOU_THRESHOLD": 0.5,
+    "MIN_VIOLATION_FRAMES": 2  # Reduced from 3 to 2 frames
 }
 
 # Setup logging
@@ -40,195 +46,177 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
-device = torch.device("cuda" if torch.cuda.is_available() and CONFIG["GPU_ACCELERATION"] else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 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 (optimized)"""
-    results = model(frame, device=device, verbose=False)
-    detections = []
-    
-    for result in results:
-        for box in result.boxes:
-            cls = int(box.cls)
-            conf = float(box.conf)
-            if conf > CONFIG["CONFIDENCE_THRESHOLD"] and cls in CONFIG["VIOLATION_LABELS"]:
-                detections.append({
-                    "frame": frame_count,
-                    "violation": CONFIG["VIOLATION_LABELS"][cls],
-                    "confidence": round(conf, 2),
-                    "bounding_box": box.xywh.cpu().numpy()[0],
-                    "timestamp": frame_count / fps
-                })
-    return detections
+# Load model with caching
+model = None
+def load_model():
+    global model
+    if model is None:
+        model = YOLO(CONFIG["MODEL_PATH"]).to(device)
+    return model
 
+# ==========================
+# Optimized Processing Functions
+# ==========================
 def process_video(video_path):
-    """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(max_workers=4) as executor:
-        futures = []
+    """Optimized video processing with early termination"""
+    try:
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30
+        frame_count = 0
+        violations = []
+        workers = []
+        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
+        start_time = time.time()
+        
         while cap.isOpened():
             ret, frame = cap.read()
             if not ret:
                 break
-                
-            if frame_count % CONFIG["FRAME_SKIP"] == 0:
-                futures.append(executor.submit(
-                    process_frame, 
-                    frame.copy(), 
-                    frame_count, 
-                    fps
-                ))
             
-            frame_count += 1
+            # Skip frames for faster processing
+            if frame_count % CONFIG["FRAME_SKIP"] != 0:
+                frame_count += 1
+                continue
+            
+            # Early termination if processing takes too long
             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:
-            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]
-                        )
+            # Run detection on current frame
+            results = load_model()(frame, device=device, verbose=False)  # Disable verbose logging
+            
+            for result in results:
+                for box in result.boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls)
+                    
+                    if not label or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                        continue
+                    
+                    bbox = box.xywh.cpu().numpy()[0]
+                    detection = {
+                        "frame": frame_count,
+                        "violation": label,
+                        "confidence": conf,
+                        "bounding_box": bbox,
+                        "timestamp": frame_count / fps
                     }
-    
-    cap.release()
-    
-    # Filter violations by frequency
-    violation_counts = {}
-    for v in violations:
-        key = (v["violation"], int(v["timestamp"]))
-        violation_counts[key] = violation_counts.get(key, 0) + 1
-    
-    filtered_violations = [
-        v for v in violations 
-        if violation_counts.get((v["violation"], int(v["timestamp"])), 0) >= CONFIG["MIN_VIOLATION_FRAMES"]
-    ]
-    
-    # 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"
-    )
+                    
+                    # Simplified worker tracking
+                    matched = False
+                    for worker in workers:
+                        if calculate_iou(worker["bbox"], bbox) > CONFIG["IOU_THRESHOLD"]:
+                            worker["bbox"] = bbox
+                            detection["worker_id"] = worker["id"]
+                            matched = True
+                            break
+                    
+                    if not matched:
+                        worker_id = len(workers) + 1
+                        workers.append({"id": worker_id, "bbox": bbox})
+                        detection["worker_id"] = worker_id
+                    
+                    violations.append(detection)
+                    
+                    # Capture snapshot if first detection of this type
+                    if not snapshot_taken[label]:
+                        snapshot_path = os.path.join(
+                            CONFIG["OUTPUT_DIR"],
+                            f"{label}_{frame_count}.jpg"
+                        )
+                        cv2.imwrite(snapshot_path, frame)
+                        snapshot_taken[label] = True
+            
+            frame_count += 1
+        
+        cap.release()
+        
+        # Filter violations (require min frames)
+        filtered_violations = []
+        violation_counts = {}
+        for v in violations:
+            key = (v["worker_id"], v["violation"])
+            violation_counts[key] = violation_counts.get(key, 0) + 1
+        
+        for v in violations:
+            if violation_counts[(v["worker_id"], v["violation"])] >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                filtered_violations.append(v)
+        
+        return {
+            "violations": filtered_violations,
+            "snapshots": [k for k, v in snapshot_taken.items() if v],
+            "message": ""
+        }
+    except Exception as e:
+        logger.error(f"Video processing failed: {e}")
+        return {
+            "violations": [],
+            "snapshots": [],
+            "message": f"Error: {str(e)}"
+        }
 
+# ==========================
+# Fast Gradio Interface
+# ==========================
 def analyze_video(video_file):
-    """Gradio interface function"""
+    """Optimized interface with progress updates"""
     if not video_file:
-        return "No video uploaded", "", "", ""
+        return "No video uploaded", "", "", "", ""
+    
+    # Immediate feedback
+    yield "Processing started...", "", "", "", ""
     
     try:
+        # Process video in background
         result = process_video(video_file)
-        return format_output(result)
+        
+        if result["message"]:
+            yield result["message"], "", "", "", ""
+            return
+        
+        # Generate simplified output
+        violation_table = (
+            "| Violation Type | Timestamp | Confidence | Worker ID |\n"
+            "|----------------|-----------|------------|-----------|\n" +
+            "\n".join(
+                f"| {v['violation']} | {v['timestamp']:.1f}s | {v['confidence']:.2f} | {v['worker_id']} |"
+                for v in result["violations"]
+            ) if result["violations"] else "No violations detected."
+        )
+        
+        snapshots_md = "\n".join(
+            f"**{violation}** detected" 
+            for violation in result["snapshots"]
+        ) if result["snapshots"] else "No snapshots available."
+        
+        yield (
+            violation_table,
+            f"Analysis complete in {time.time() - start_time:.1f}s",
+            snapshots_md,
+            "Salesforce integration placeholder",
+            "Report URL placeholder"
+        )
     except Exception as e:
-        logger.error(f"Error: {str(e)}")
-        return f"Error: {str(e)}", "", "", ""
+        yield f"Error: {str(e)}", "", "", "", ""
 
-# Gradio Interface
+# Launch optimized interface
 interface = gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload Site Video"),
     outputs=[
         gr.Markdown("## Detected Violations"),
-        gr.Textbox(label="Safety Score"),
+        gr.Textbox(label="Processing Time"),
         gr.Markdown("## Violation Snapshots"),
-        gr.Textbox(label="Processing Info")
+        gr.Textbox(label="Salesforce Record"),
+        gr.Textbox(label="Report URL")
     ],
-    title="AI Safety Compliance Analyzer",
-    description="Optimized for fast detection of safety violations",
+    title="Fast Safety Compliance Analyzer",
+    description="Optimized for quick safety violation detection in worksite videos.",
     allow_flagging="never"
 )