Update app.py

app.py

@@ -18,46 +18,40 @@ from retrying import retry
 # 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_PATH": "yolov8n.pt",
     "OUTPUT_DIR": "static/output",
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
         2: "unsafe_posture"
     },
-    "DISPLAY_NAMES": {  # Mapping for user-friendly violation names
-        "no_helmet": "Missing Helmet",
-        "no_harness": "Missing Harness",
-        "unsafe_posture": "Unsafe Posture"
+    "DISPLAY_NAMES": {
+        "no_helmet": "No Helmet Violation",
+        "no_harness": "No Harness Violation",
+        "unsafe_posture": "Unsafe Posture Violation"
     },
     "SF_CREDENTIALS": {
         "username": "your_username@safety.com",
         "password": "your_password",
         "security_token": "your_security_token",
-        "domain": "login"  # Use "test" for sandbox
+        "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
-    "FRAME_SKIP": 15,  # Process every 15th frame
-    "CONFIDENCE_THRESHOLD": 0.5  # Minimum confidence for violation detection
+    "FRAME_SKIP": 15,
+    "MAX_PROCESSING_TIME": 30,  # Updated to 30 seconds
+    "CONFIDENCE_THRESHOLD": 0.5
 }
 
 # 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"]
@@ -75,9 +69,6 @@ def load_model():
 
 model = load_model()
 
-# ==========================
-# Salesforce Integration
-# ==========================
 @retry(stop_max_attempt_number=2, wait_fixed=1000)
 def connect_to_salesforce():
     try:
@@ -207,9 +198,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.error(f"Salesforce record creation failed: {e}")
         return None, ""
 
-# ==========================
-# Safety Score Calculation
-# ==========================
 def calculate_safety_score(violations):
     penalties = {
         "no_helmet": 25,
@@ -222,9 +210,6 @@ def calculate_safety_score(violations):
             score -= penalties[v["violation"]]
     return max(score, 0)
 
-# ==========================
-# Video Processing
-# ==========================
 def process_video(video_data):
     try:
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
@@ -238,30 +223,32 @@ def process_video(video_data):
 
         violations, snapshots = [], []
         frame_count = 0
+        start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
 
-        # 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:
-                break
+                break  # End of video
 
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
 
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.info("Processing time limit of 30 seconds reached")
+                break
+
             results = model(frame, device=device)
             seen_violations = set()
             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
+                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
                         continue
-                    # Apply confidence threshold
                     if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
                     if label in seen_violations:
@@ -277,7 +264,6 @@ def process_video(video_data):
                     }
                     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)
@@ -304,7 +290,7 @@ def process_video(video_data):
                 "score": 100,
                 "salesforce_record_id": None,
                 "violation_details_url": "",
-                "message": "No violations detected here."
+                "message": "No violations detected in the video."
             }
 
         score = calculate_safety_score(violations)
@@ -327,12 +313,9 @@ def process_video(video_data):
             "score": 100,
             "salesforce_record_id": None,
             "violation_details_url": "",
-            "message": "Error processing video."
+            "message": f"Error processing video: {e}"
         }
 
-# ==========================
-# Gradio Interface
-# ==========================
 def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
@@ -342,34 +325,26 @@ def gradio_interface(video_file):
         result = process_video(video_data)
 
         if result.get("message"):
-            # Show message (like "No violations detected here.")
-            return result["message"], f"Safety Score: {result['score']}%", "", "N/A", "N/A"
+            # If message present (either no violations or error), show it plainly
+            return result["message"], "", "", "", ""
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation         | Timestamp | Confidence | Violation Details                 |\n"
-            separator = "|-------------------|-----------|------------|---------------------------------|\n"
+            header = "| Violation              | Timestamp (s) | Confidence | Bounding Box             |\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"])
-                # Provide clearer human-readable violation explanation
-                if v["violation"] == "no_helmet":
-                    details = "Employee not wearing helmet"
-                elif v["violation"] == "no_harness":
-                    details = "Employee not wearing proper harness"
-                elif v["violation"] == "unsafe_posture":
-                    details = "Employee in unsafe posture/zone"
-                else:
-                    details = "Violation detected"
-
-                row = f"| {display_name:<17} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}     | {details:<31} |"
+                display_name = violation_name_map.get(v["violation"], v["violation"])
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['bounding_box']} |"
                 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"]
             )
 
@@ -395,7 +370,7 @@ 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 Violation, No Harness Violation, Unsafe Posture Violation). Non-violations are ignored."
 )
 
 if __name__ == "__main__":