Update app.py

app.py

@@ -15,60 +15,68 @@ import logging
 from retrying import retry
 
 # ==========================
-# Configuration
+# Enhanced Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",  # Custom-trained model for specific violations
-    "FALLBACK_MODEL_PATH": "yolov8n.pt",  # Fallback if custom model is missing
+    "MODEL_PATH": "yolov8_safety.pt",
+    "FALLBACK_MODEL": "yolov8n.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture"
+        2: "unsafe_posture",
+        3: "unsafe_zone",
+        4: "improper_tool_use"
     },
-    "DISPLAY_NAMES": {  # Mapping for user-friendly violation names
-        "no_helmet": "Missing Helmet",
-        "no_harness": "Missing Harness",
-        "unsafe_posture": "Unsafe Posture"
+    "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
+    },
+    "DISPLAY_NAMES": {
+        "no_helmet": "No Helmet Violation",
+        "no_harness": "No Harness Violation",
+        "unsafe_posture": "Unsafe Posture Violation",
+        "unsafe_zone": "Unsafe Zone Entry",
+        "improper_tool_use": "Improper Tool Use"
     },
     "SF_CREDENTIALS": {
-        "username": "your_username@safety.com",
-        "password": "your_password",
-        "security_token": "your_security_token",
-        "domain": "login"  # Use "test" for sandbox
+        "username": "prashanth1ai@safety.com",
+        "password": "SaiPrash461",
+        "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
+        "domain": "login"
     },
-    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
-    "FRAME_SKIP": 15,  # Process every 15th frame
-    "MAX_PROCESSING_TIME": 25,  # Cap video processing at 25s
-    "CONFIDENCE_THRESHOLD": 0.5  # Minimum confidence for violation detection
+    "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
+    "FRAME_SKIP": 5,  # Reduced for better detection
+    "MAX_PROCESSING_TIME": 60,
+    "CONFIDENCE_THRESHOLD": 0.25,  # Lower threshold for all violations
+    "IOU_THRESHOLD": 0.4,
+    "MIN_VIOLATION_FRAMES": 3  # Minimum consecutive frames to confirm violation
 }
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# Ensure output directory exists
 os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
 
-# ==========================
-# Device Setup
-# ==========================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
-# ==========================
-# Model Loading
-# ==========================
 def load_model():
     try:
-        model_path = CONFIG["MODEL_PATH"]
-        if not os.path.exists(model_path):
-            logger.warning(f"Custom model {model_path} not found. Falling back to {CONFIG['FALLBACK_MODEL_PATH']}")
-            model_path = CONFIG["FALLBACK_MODEL_PATH"]
-        model = YOLO(model_path).to(device)
-        logger.info(f"Model loaded: {model_path}")
-        if model_path == CONFIG["FALLBACK_MODEL_PATH"]:
+        if os.path.isfile(CONFIG["MODEL_PATH"]):
+            model_path = CONFIG["MODEL_PATH"]
+            logger.info(f"Model loaded: {model_path}")
+        else:
+            model_path = CONFIG["FALLBACK_MODEL"]
             logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
+            if not os.path.isfile(model_path):
+                logger.info(f"Downloading fallback model: {model_path}")
+                torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
+        model = YOLO(model_path).to(device)
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -77,9 +85,57 @@ def load_model():
 model = load_model()
 
 # ==========================
-# Salesforce Integration
+# Enhanced Helper Functions
+# ==========================
+def draw_detections(frame, detections):
+    """Draw bounding boxes and labels on frame"""
+    for det in detections:
+        label = det["violation"]
+        confidence = det["confidence"]
+        x, y, w, h = det["bounding_box"]
+        
+        # Convert from center coordinates to corner coordinates
+        x1 = int(x - w/2)
+        y1 = int(y - h/2)
+        x2 = int(x + w/2)
+        y2 = int(y + h/2)
+        
+        color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)}: {confidence:.2f}"
+        cv2.putText(frame, display_text, (x1, y1-10), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    return frame
+
+def calculate_iou(box1, box2):
+    """Calculate Intersection over Union (IoU) for two bounding boxes."""
+    x1, y1, w1, h1 = box1
+    x2, y2, w2, h2 = box2
+    
+    # Convert to top-left and bottom-right coordinates
+    x1_min, y1_min = x1 - w1/2, y1 - h1/2
+    x1_max, y1_max = x1 + w1/2, y1 + h1/2
+    x2_min, y2_min = x2 - w2/2, y2 - h2/2
+    x2_max, y2_max = x2 + w2/2, y2 + h2/2
+    
+    # Calculate intersection
+    x_min = max(x1_min, x2_min)
+    y_min = max(y1_min, y2_min)
+    x_max = min(x1_max, x2_max)
+    y_max = min(y1_max, y2_max)
+    
+    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    area1 = w1 * h1
+    area2 = w2 * h2
+    union = area1 + area2 - intersection
+    
+    return intersection / union if union > 0 else 0
+
+# ==========================
+# Salesforce Integration (unchanged)
 # ==========================
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
@@ -139,7 +195,6 @@ def generate_violation_pdf(violations, score):
         logger.error(f"Error generating PDF: {e}")
         return "", "", None
 
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def upload_pdf_to_salesforce(sf, pdf_file, report_id):
     try:
         if not pdf_file:
@@ -153,7 +208,7 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
             "FirstPublishLocationId": report_id
         }
         content_version = sf.ContentVersion.create(content_version_data)
-        result = sf.query(f"SELECT Id FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
         if not result['records']:
             logger.error("Failed to retrieve ContentVersion")
             return ""
@@ -164,7 +219,6 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         logger.error(f"Error uploading PDF to Salesforce: {e}")
         return ""
 
-@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
@@ -183,10 +237,10 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         }
         logger.info(f"Creating Salesforce record with data: {record_data}")
         try:
-            record = sf.Safety_Violation_Report__c.create(record_data)
-            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
+            record = sf.Safety_Video_Report__c.create(record_data)
+            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
         except Exception as e:
-            logger.error(f"Failed to create Safety_Violation_Report__c: {e}")
+            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
             logger.warning(f"Fell back to Account record: {record['id']}")
         record_id = record["id"]
@@ -195,36 +249,39 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
             if uploaded_url:
                 try:
-                    sf.Safety_Violation_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
                     logger.info(f"Updated record {record_id} with PDF URL: {uploaded_url}")
                 except Exception as e:
-                    logger.error(f"Failed to update Safety_Violation_Report__c: {e}")
+                    logger.error(f"Failed to update Safety_Video_Report__c: {e}")
                     sf.Account.update(record_id, {"Description": uploaded_url})
                     logger.info(f"Updated Account record {record_id} with PDF URL")
                 pdf_url = uploaded_url
 
         return record_id, pdf_url
     except Exception as e:
-        logger.error(f"Salesforce record creation failed: {e}")
+        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
         return None, ""
 
-# ==========================
-# Safety Score Calculation
-# ==========================
 def calculate_safety_score(violations):
     penalties = {
         "no_helmet": 25,
         "no_harness": 30,
-        "unsafe_posture": 20
+        "unsafe_posture": 20,
+        "unsafe_zone": 35,
+        "improper_tool_use": 25
     }
-    score = 100
+    # Count unique violations per worker
+    unique_violations = set()
     for v in violations:
-        if v["violation"] in penalties:
-            score -= penalties[v["violation"]]
+        key = (v["worker_id"], v["violation"])
+        unique_violations.add(key)
+    
+    total_penalty = sum(penalties.get(violation, 0) for _, violation in unique_violations)
+    score = 100 - total_penalty
     return max(score, 0)
 
 # ==========================
-# Video Processing
+# Enhanced Video Processing
 # ==========================
 def process_video(video_data):
     try:
@@ -237,83 +294,152 @@ def process_video(video_data):
         if not video.isOpened():
             raise ValueError("Could not open video file")
 
-        violations, snapshots = [], []
+        violations = []
+        snapshots = []
         frame_count = 0
         start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
-        max_frames = int(60 * fps)  # Process up to 1 minute
+        if fps <= 0:
+            fps = 30  # Default assumption if FPS cannot be determined
+
+        # Structure to track workers and their violations
+        workers = []
+        violation_history = {label: [] for label in CONFIG["VIOLATION_LABELS"].values()}
+        snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
 
-        # Track one snapshot per violation type
-        snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
+        logger.info(f"Processing video with FPS: {fps}")
+        logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
 
         while True:
             ret, frame = video.read()
-            if not ret or frame_count >= max_frames:
+            if not ret:
                 break
 
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
 
-            # Stop if processing time exceeds 25 seconds
             if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
                 logger.info("Processing time limit reached")
                 break
 
+            current_time = frame_count / fps
+            
+            # Run detection on this frame
             results = model(frame, device=device)
-            seen_violations = set()
+            
+            current_detections = []
             for result in results:
-                for box in result.boxes:
-                    cls, conf = int(box.cls), float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls, f"unknown_class_{cls}")
-                    # Only process specified violations
-                    if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
-                        logger.warning(f"Unexpected detection: {label} (cls: {cls}, conf: {conf}) - ignored")
+                boxes = result.boxes
+                for box in boxes:
+                    cls = int(box.cls)
+                    conf = float(box.conf)
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                    
+                    if label is None:
                         continue
-                    # Apply confidence threshold
+                        
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
-                        logger.info(f"Skipping low-confidence detection: {label} (conf: {conf})")
                         continue
-                    if label in seen_violations:
-                        continue
-                    seen_violations.add(label)
 
-                    violation = {
+                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
+                    
+                    current_detections.append({
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
-                        "timestamp": frame_count / fps
-                    }
-                    violations.append(violation)
-
-                    # Save only one snapshot per violation type
-                    if not snapshot_taken[label]:
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
-                        cv2.imwrite(snapshot_path, frame)
-                        with open(snapshot_path, "rb") as img_file:
-                            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
-                        snapshots.append({
-                            "violation": label,
-                            "frame": frame_count,
-                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
-                            "snapshot_base64": f"data:image/jpeg;base64,{img_base64}"
-                        })
-                        snapshot_taken[label] = True
+                        "bounding_box": bbox,
+                        "timestamp": current_time
+                    })
+
+            # Process detections and associate with workers
+            for detection in current_detections:
+                # Find matching worker
+                matched_worker = None
+                max_iou = 0
+                
+                for worker in workers:
+                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                        max_iou = iou
+                        matched_worker = worker
+
+                if matched_worker:
+                    # Update worker's position
+                    matched_worker["bbox"] = detection["bounding_box"]
+                    matched_worker["last_seen"] = current_time
+                    worker_id = matched_worker["id"]
+                else:
+                    # New worker
+                    worker_id = len(workers) + 1
+                    workers.append({
+                        "id": worker_id,
+                        "bbox": detection["bounding_box"],
+                        "first_seen": current_time,
+                        "last_seen": current_time
+                    })
+                
+                # Add to violation history
+                detection["worker_id"] = worker_id
+                violation_history[detection["violation"]].append(detection)
 
             frame_count += 1
 
         video.release()
         os.remove(video_path)
-
+        
+        # Process violation history to confirm persistent violations
+        for violation_type, detections in violation_history.items():
+            if not detections:
+                continue
+                
+            # Group by worker
+            worker_violations = {}
+            for det in detections:
+                if det["worker_id"] not in worker_violations:
+                    worker_violations[det["worker_id"]] = []
+                worker_violations[det["worker_id"]].append(det)
+            
+            # Check each worker's violations for persistence
+            for worker_id, worker_dets in worker_violations.items():
+                if len(worker_dets) >= CONFIG["MIN_VIOLATION_FRAMES"]:
+                    # Take the highest confidence detection
+                    best_detection = max(worker_dets, key=lambda x: x["confidence"])
+                    violations.append(best_detection)
+                    
+                    # Capture snapshot if not already taken
+                    if not snapshot_taken[violation_type]:
+                        # Get the frame for this violation
+                        cap = cv2.VideoCapture(video_path)
+                        cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                        ret, snapshot_frame = cap.read()
+                        cap.release()
+                        
+                        if ret:
+                            # Draw detections on snapshot
+                            snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                            
+                            snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
+                            snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                            cv2.imwrite(snapshot_path, snapshot_frame)
+                            snapshots.append({
+                                "violation": violation_type,
+                                "frame": best_detection["frame"],
+                                "snapshot_path": snapshot_path,
+                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                            })
+                            snapshot_taken[violation_type] = True
+
+        # Final processing
         if not violations:
-            logger.info("No violations detected")
+            logger.info("No persistent violations detected")
             return {
                 "violations": [],
                 "snapshots": [],
                 "score": 100,
                 "salesforce_record_id": None,
-                "violation_details_url": ""
+                "violation_details_url": "",
+                "message": "No violations detected in the video."
             }
 
         score = calculate_safety_score(violations)
@@ -325,16 +451,18 @@ def process_video(video_data):
             "snapshots": snapshots,
             "score": score,
             "salesforce_record_id": report_id,
-            "violation_details_url": final_pdf_url
+            "violation_details_url": final_pdf_url,
+            "message": ""
         }
     except Exception as e:
-        logger.error(f"Error processing video: {e}")
+        logger.error(f"Error processing video: {e}", exc_info=True)
         return {
             "violations": [],
             "snapshots": [],
             "score": 100,
             "salesforce_record_id": None,
-            "violation_details_url": ""
+            "violation_details_url": "",
+            "message": f"Error processing video: {e}"
         }
 
 # ==========================
@@ -344,29 +472,38 @@ def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
+        yield "Processing video... please wait.", "", "", "", ""
+
         with open(video_file, "rb") as f:
             video_data = f.read()
+
         result = process_video(video_data)
 
+        if result.get("message"):
+            yield result["message"], "", "", "", ""
+            return
+
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation        | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
-            separator = "|------------------|-----------|------------|--------------------------|-------------------------|\n"
+            header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+            separator = "|------------------------|---------------|------------|-----------|\n"
             rows = []
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
             for v in result["violations"]:
-                display_name = CONFIG["DISPLAY_NAMES"].get(v["violation"], v["violation"])
-                row = f"| {display_name:<16} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
+                display_name = violation_name_map.get(v["violation"], v["violation"])
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
         snapshots_text = "No snapshots captured."
         if result["snapshots"]:
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
             snapshots_text = "\n".join(
-                f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
 
-        return (
+        yield (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_text,
@@ -374,13 +511,13 @@ def gradio_interface(video_file):
             result["violation_details_url"] or "N/A"
         )
     except Exception as e:
-        logger.error(f"Error in Gradio interface: {e}")
-        return f"Error: {str(e)}", "", "Error in processing.", "", ""
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
-    outputs=[
+    outputs=[ 
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
         gr.Markdown(label="Snapshots"),
@@ -388,9 +525,10 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (Missing Helmet, Missing Harness, Unsafe Posture). Non-violations are ignored."
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Safety Analyzer App...")
-    interface.launch()
+    logger.info("Launching Enhanced Safety Analyzer App...")
+    interface.launch()