Update app.py

app.py

@@ -15,68 +15,60 @@ import logging
 from retrying import retry
 
 # ==========================
-# Enhanced Configuration
+# Configuration
 # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",
-    "FALLBACK_MODEL": "yolov8n.pt",
+    "MODEL_PATH": "yolov8_safety.pt",  # Custom-trained model for specific violations
+    "FALLBACK_MODEL_PATH": "yolov8n.pt",  # Fallback if custom model is missing
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture",
-        3: "unsafe_zone",
-        4: "improper_tool_use"
+        2: "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"
+    "DISPLAY_NAMES": {  # Mapping for user-friendly violation names
+        "no_helmet": "Missing Helmet",
+        "no_harness": "Missing Harness",
+        "unsafe_posture": "Unsafe Posture"
     },
     "SF_CREDENTIALS": {
-        "username": "prashanth1ai@safety.com",
-        "password": "SaiPrash461",
-        "security_token": "AP4AQnPoidIKPvSvNEfAHyoK",
-        "domain": "login"
+        "username": "your_username@safety.com",
+        "password": "your_password",
+        "security_token": "your_security_token",
+        "domain": "login"  # Use "test" for sandbox
     },
-    "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
+    "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
 }
 
 # 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:
-        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_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"]:
+            logger.warning("Using fallback model. Detection accuracy may be poor. Train yolov8_safety.pt for best results.")
         return model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
@@ -85,57 +77,9 @@ def load_model():
 model = load_model()
 
 # ==========================
-# 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)
+# Salesforce Integration
 # ==========================
-@retry(stop_max_attempt_number=3, wait_fixed=2000)
+@retry(stop_max_attempt_number=2, wait_fixed=1000)
 def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
@@ -195,6 +139,7 @@ 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:
@@ -208,7 +153,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, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        result = sf.query(f"SELECT Id FROM ContentVersion WHERE Id = '{content_version['id']}'")
         if not result['records']:
             logger.error("Failed to retrieve ContentVersion")
             return ""
@@ -219,6 +164,7 @@ 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()
@@ -237,10 +183,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_Video_Report__c.create(record_data)
-            logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
+            record = sf.Safety_Violation_Report__c.create(record_data)
+            logger.info(f"Created Safety_Violation_Report__c record: {record['id']}")
         except Exception as e:
-            logger.error(f"Failed to create Safety_Video_Report__c: {e}")
+            logger.error(f"Failed to create Safety_Violation_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"]
@@ -249,39 +195,36 @@ 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_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    sf.Safety_Violation_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_Video_Report__c: {e}")
+                    logger.error(f"Failed to update Safety_Violation_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}", exc_info=True)
+        logger.error(f"Salesforce record creation failed: {e}")
         return None, ""
 
+# ==========================
+# Safety Score Calculation
+# ==========================
 def calculate_safety_score(violations):
     penalties = {
         "no_helmet": 25,
         "no_harness": 30,
-        "unsafe_posture": 20,
-        "unsafe_zone": 35,
-        "improper_tool_use": 25
+        "unsafe_posture": 20
     }
-    # Count unique violations per worker
-    unique_violations = set()
+    score = 100
     for v in violations:
-        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
+        if v["violation"] in penalties:
+            score -= penalties[v["violation"]]
     return max(score, 0)
 
 # ==========================
-# Enhanced Video Processing
+# Video Processing
 # ==========================
 def process_video(video_data):
     try:
@@ -294,152 +237,83 @@ 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)
-        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()}
+        max_frames = int(60 * fps)  # Process up to 1 minute
 
-        logger.info(f"Processing video with FPS: {fps}")
-        logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
+        # Track one snapshot per violation type
+        snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
 
         while True:
             ret, frame = video.read()
-            if not ret:
+            if not ret or frame_count >= max_frames:
                 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)
-            
-            current_detections = []
+            seen_violations = set()
             for result in results:
-                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:
+                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")
                         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)
 
-                    bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
-                    
-                    current_detections.append({
+                    violation = {
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "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)
+                        "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
 
             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 persistent violations detected")
+            logger.info("No violations detected")
             return {
                 "violations": [],
                 "snapshots": [],
                 "score": 100,
                 "salesforce_record_id": None,
-                "violation_details_url": "",
-                "message": "No violations detected in the video."
+                "violation_details_url": ""
             }
 
         score = calculate_safety_score(violations)
@@ -451,18 +325,16 @@ def process_video(video_data):
             "snapshots": snapshots,
             "score": score,
             "salesforce_record_id": report_id,
-            "violation_details_url": final_pdf_url,
-            "message": ""
+            "violation_details_url": final_pdf_url
         }
     except Exception as e:
-        logger.error(f"Error processing video: {e}", exc_info=True)
+        logger.error(f"Error processing video: {e}")
         return {
             "violations": [],
             "snapshots": [],
             "score": 100,
             "salesforce_record_id": None,
-            "violation_details_url": "",
-            "message": f"Error processing video: {e}"
+            "violation_details_url": ""
         }
 
 # ==========================
@@ -472,38 +344,29 @@ 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 (s) | Confidence | Worker ID |\n"
-            separator = "|------------------------|---------------|------------|-----------|\n"
+            header = "| Violation        | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
+            separator = "|------------------|-----------|------------|--------------------------|-------------------------|\n"
             rows = []
-            violation_name_map = CONFIG["DISPLAY_NAMES"]
             for v in result["violations"]:
-                display_name = violation_name_map.get(v["violation"], v["violation"])
-                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']} |"
+                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']} |"
                 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 {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
 
-        yield (
+        return (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_text,
@@ -511,13 +374,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}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+        logger.error(f"Error in Gradio interface: {e}")
+        return 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"),
@@ -525,10 +388,9 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    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"
+    description="Upload site videos to detect safety violations (Missing Helmet, Missing Harness, Unsafe Posture). Non-violations are ignored."
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
-    interface.launch()
+    logger.info("Launching Safety Analyzer App...")
+    interface.launch()