Update app.py

app.py

@@ -39,6 +39,7 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo1/resolve/main/static/output/",
     "FRAME_SKIP": 15,
+    "MAX_PROCESSING_TIME": 30,  # seconds
     "CONFIDENCE_THRESHOLD": 0.5
 }
 
@@ -209,7 +210,7 @@ def calculate_safety_score(violations):
             score -= penalties[v["violation"]]
     return max(score, 0)
 
-def process_video(video_data, progress_callback=None):
+def process_video(video_data):
     try:
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
@@ -220,11 +221,11 @@ def process_video(video_data, progress_callback=None):
         if not video.isOpened():
             raise ValueError("Could not open video file")
 
-        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps = video.get(cv2.CAP_PROP_FPS)
-
         violations, snapshots = [], []
         frame_count = 0
+        start_time = time.time()
+        fps = video.get(cv2.CAP_PROP_FPS)
+
         snapshot_taken = {"no_helmet": False, "no_harness": False, "unsafe_posture": False}
 
         while True:
@@ -234,10 +235,12 @@ def process_video(video_data, progress_callback=None):
 
             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:
@@ -252,22 +255,11 @@ def process_video(video_data, progress_callback=None):
                         continue
                     seen_violations.add(label)
 
-                    xywh_array = box.xywh.cpu().numpy()
-                    if xywh_array.size == 0 or xywh_array.shape[1] < 4:
-                        xyxy_array = box.xyxy.cpu().numpy()
-                        if xyxy_array.size == 0 or xyxy_array.shape[1] < 4:
-                            logger.warning("No valid bounding box info found, skipping box")
-                            continue
-                        xmin, ymin, xmax, ymax = xyxy_array[0]
-                        width = xmax - xmin
-                        height = ymax - ymin
-                        xywh_array = np.array([[xmin, ymin, width, height]])
-
                     violation = {
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in xywh_array[0]],
+                        "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
                         "timestamp": frame_count / fps
                     }
                     violations.append(violation)
@@ -287,9 +279,6 @@ def process_video(video_data, progress_callback=None):
 
             frame_count += 1
 
-            if progress_callback:
-                progress_callback(min(frame_count / total_frames, 1.0))
-
         video.release()
         os.remove(video_path)
 
@@ -327,18 +316,22 @@ def process_video(video_data, progress_callback=None):
             "message": f"Error processing video: {e}"
         }
 
-def gradio_interface(video_file, progress=gr.Progress(track_tqdm=True)):
+def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
-
     try:
+        # Show processing message early in the UI
+        yield "Processing video... please wait.", "", "", "", ""
+
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        result = process_video(video_data, progress_callback=progress)
+        result = process_video(video_data)
 
         if result.get("message"):
-            return result["message"], "", "", "", ""
+            # If message present (either no violations or error), show it plainly
+            yield result["message"], "", "", "", ""
+            return
 
         violation_table = "No violations detected."
         if result["violations"]:
@@ -360,7 +353,7 @@ def gradio_interface(video_file, progress=gr.Progress(track_tqdm=True)):
                 for s in result["snapshots"]
             )
 
-        return (
+        yield (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_text,
@@ -369,7 +362,7 @@ def gradio_interface(video_file, progress=gr.Progress(track_tqdm=True)):
         )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}")
-        return f"Error: {str(e)}", "", "Error in processing.", "", ""
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,
@@ -383,7 +376,7 @@ interface = gr.Interface(
     ],
     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.",
-    live=False
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":