Update app.py

app.py

@@ -39,7 +39,6 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
     "FRAME_SKIP": 15,
-    "MAX_PROCESSING_TIME": 30,  # Updated to 30 seconds
     "CONFIDENCE_THRESHOLD": 0.5
 }
 
@@ -210,7 +209,7 @@ def calculate_safety_score(violations):
             score -= penalties[v["violation"]]
     return max(score, 0)
 
-def process_video(video_data):
+def process_video(video_data, progress_callback=None):
     try:
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
@@ -221,11 +220,11 @@ def process_video(video_data):
         if not video.isOpened():
             raise ValueError("Could not open video file")
 
-        violations, snapshots = [], []
-        frame_count = 0
-        start_time = time.time()
+        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = video.get(cv2.CAP_PROP_FPS)
 
+        violations, snapshots = [], []
+        frame_count = 0
         snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
 
         while True:
@@ -235,12 +234,10 @@ def process_video(video_data):
 
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
+                if progress_callback:
+                    progress_callback(min(frame_count / total_frames, 1.0))
                 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:
@@ -279,6 +276,9 @@ def process_video(video_data):
 
             frame_count += 1
 
+            if progress_callback:
+                progress_callback(min(frame_count / total_frames, 1.0))
+
         video.release()
         os.remove(video_path)
 
@@ -316,16 +316,17 @@ def process_video(video_data):
             "message": f"Error processing video: {e}"
         }
 
-def gradio_interface(video_file):
+def gradio_interface(video_file, progress=gr.Progress(track_tqdm=True)):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
+
     try:
         with open(video_file, "rb") as f:
             video_data = f.read()
-        result = process_video(video_data)
+
+        result = process_video(video_data, progress_callback=progress)
 
         if result.get("message"):
-            # If message present (either no violations or error), show it plainly
             return result["message"], "", "", "", ""
 
         violation_table = "No violations detected."
@@ -370,7 +371,8 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet Violation, No Harness Violation, Unsafe Posture Violation). 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.",
+    live=False
 )
 
 if __name__ == "__main__":