Update app.py

app.py

@@ -33,7 +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
+    "MAX_PROCESSING_TIME": 60  # Set to 60 seconds for 1-minute video
 }
 
 # Setup logging
@@ -198,9 +198,9 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 # ==========================
 def calculate_safety_score(violations, total_frames):
     penalties = {
-        "no_helmet": 25 / total_frames,  # Normalize penalty by video length
-        "no_harness": 30 / total_frames,
-        "unsafe_posture": 20 / total_frames
+        "no_helmet": 25 / max(total_frames, 1),  # Normalize penalty by video length
+        "no_harness": 30 / max(total_frames, 1),
+        "unsafe_posture": 20 / max(total_frames, 1)
     }
     score = 100
     for v in violations:
@@ -226,10 +226,12 @@ def process_video(video_data):
         frame_count = 0
         total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
         start_time = time.time()
+        fps = video.get(cv2.CAP_PROP_FPS)
+        max_frames = int(fps * CONFIG["MAX_PROCESSING_TIME"])  # Max frames for 60 seconds
 
         while True:
             ret, frame = video.read()
-            if not ret:
+            if not ret or frame_count >= max_frames:
                 break
 
             results = model(frame, device=device)
@@ -247,7 +249,7 @@ def process_video(video_data):
                         "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
-                        "timestamp": frame_count / video.get(cv2.CAP_PROP_FPS)
+                        "timestamp": frame_count / max(fps, 1)
                     }
                     violations.append(violation)
 
@@ -264,12 +266,12 @@ 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")
+                logger.warning("Processing time limit of 60 seconds exceeded")
                 break
 
         video.release()
         os.remove(video_path)
-        score = calculate_safety_score(violations, max(total_frames, 1))
+        score = calculate_safety_score(violations, max(total_frames, frame_count))
         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)
 
@@ -303,8 +305,8 @@ def gradio_interface(video_file):
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation     | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
-            separator = "|---------------|-----------|------------|--------------------------|-------------------------|\n"
+            header = "| Violation        | Timestamp | Confidence | Bounding Box             | Violation Details       |\n"
+            separator = "|------------------|-----------|------------|--------------------------|-------------------------|\n"
             rows = []
             for v in result["violations"]:
                 violation_name = v["violation"]
@@ -314,7 +316,7 @@ def gradio_interface(video_file):
                     violation_name = "Missing Harness"
                 else:
                     violation_name = "Unsafe Posture"
-                row = f"| {violation_name:<13} | {v['timestamp']:.2f}s  | {v['confidence']:.2f}      | {v['bounding_box']} | {result['violation_details_url']} |"
+                row = f"| {violation_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)
 
@@ -347,7 +349,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., missing helmet, missing harness, unsafe posture)."
+    description="Upload a video (up to 1 minute) to detect safety violations (missing helmet, missing harness, unsafe posture)."
 )
 
 if __name__ == "__main__":