Update app.py

app.py

@@ -24,8 +24,7 @@ CONFIG = {
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
-        2: "unsafe_posture",
-        3: "unsafe_zone"
+        2: "unsafe_posture"
     },
     "SF_CREDENTIALS": {
         "username": "prashanth1ai@safety.com",
@@ -34,6 +33,7 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
+    "MAX_PROCESSING_TIME": 30  # Reduced to 30 seconds
 }
 
 # Setup logging
@@ -196,15 +196,17 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 # ==========================
 # Safety Score Calculation
 # ==========================
-def calculate_safety_score(violations):
+def calculate_safety_score(violations, total_frames):
     penalties = {
-        "no_helmet": 25,
-        "no_harness": 30,
-        "unsafe_posture": 20
+        "no_helmet": 25 / total_frames,  # Normalize penalty by video length
+        "no_harness": 30 / total_frames,
+        "unsafe_posture": 20 / total_frames
     }
-    # Only penalize for detected violations (no_helmet, no_harness, unsafe_posture)
-    score = 100 - sum(penalties.get(v["violation"], 0) for v in violations)
-    return max(score, 0)
+    score = 100
+    for v in violations:
+        penalty = penalties.get(v["violation"], 0)
+        score -= penalty
+    return max(round(score, 2), 0)
 
 # ==========================
 # Video Processing
@@ -222,6 +224,8 @@ def process_video(video_data):
 
         violations, snapshots = [], []
         frame_count = 0
+        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+        start_time = time.time()
 
         while True:
             ret, frame = video.read()
@@ -229,16 +233,13 @@ def process_video(video_data):
                 break
 
             results = model(frame, device=device)
-            seen_violations = set()  # Track unique violations in this frame
+            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"class_{cls}")
-                    # Only consider "no_helmet", "no_harness", "unsafe_posture" as violations
-                    if label not in ["no_helmet", "no_harness", "unsafe_posture"]:
+                    if label in seen_violations or label not in ["no_helmet", "no_harness", "unsafe_posture"]:
                         continue
-                    if label in seen_violations:
-                        continue  # Skip if this violation type was already recorded in this frame
                     seen_violations.add(label)
 
                     violation = {
@@ -262,10 +263,13 @@ def process_video(video_data):
                     })
 
             frame_count += 1
+            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.warning("Processing time limit of 30 seconds exceeded")
+                break
 
         video.release()
         os.remove(video_path)
-        score = calculate_safety_score(violations)
+        score = calculate_safety_score(violations, max(total_frames, 1))
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
         report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
@@ -305,9 +309,11 @@ def gradio_interface(video_file):
             for v in result["violations"]:
                 violation_name = v["violation"]
                 if violation_name == "no_helmet":
-                    violation_name = "no_helmet"
+                    violation_name = "Missing Helmet"
+                elif violation_name == "no_harness":
+                    violation_name = "Missing Harness"
                 else:
-                    violation_name = violation_name.replace("no_", "").replace("unsafe_", "")
+                    violation_name = "Unsafe Posture"
                 row = f"| {violation_name:<13} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
@@ -315,7 +321,7 @@ def gradio_interface(video_file):
         snapshots_text = "No snapshots captured."
         if result["snapshots"]:
             snapshots_text = "\n".join(
-                f"- Snapshot for {s['violation'].replace('no_', '').replace('unsafe_', '')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
+                f"- Snapshot for {s['violation'].replace('no_helmet', 'Missing Helmet').replace('no_harness', 'Missing Harness').replace('unsafe_posture', 'Unsafe Posture')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
 
@@ -341,7 +347,7 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (e.g., no helmet, no harness, unsafe posture)."
+    description="Upload site videos to detect safety violations (e.g., missing helmet, missing harness, unsafe posture)."
 )
 
 if __name__ == "__main__":