Update app.py

app.py

@@ -38,10 +38,9 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 1,  # Process every frame (set to 1 for full video processing)
+    "FRAME_SKIP": 15,
     "MAX_PROCESSING_TIME": 30,
-    "CONFIDENCE_THRESHOLD": 0.5,
-    "TEMPORAL_THRESHOLD": 1.0  # Time threshold in seconds to avoid counting the same violation
+    "CONFIDENCE_THRESHOLD": 0.5
 }
 
 # Setup logging
@@ -210,13 +209,11 @@ def calculate_safety_score(violations):
 
 def process_video(video_data):
     try:
-        # Save video to temporary file
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
             f.write(video_data)
         logger.info(f"Video saved: {video_path}")
 
-        # Read the video
         video = cv2.VideoCapture(video_path)
         if not video.isOpened():
             raise ValueError("Could not open video file")
@@ -227,14 +224,13 @@ def process_video(video_data):
         fps = video.get(cv2.CAP_PROP_FPS)
 
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        last_detected = {label: 0 for label in CONFIG["VIOLATION_LABELS"].values()}  # Tracks last violation time
+        seen_violations = set()  # to track violations and avoid repeating
 
         while True:
             ret, frame = video.read()
             if not ret:
-                break  # Break if the video has ended
+                break
 
-            # Process every frame (or based on FRAME_SKIP)
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
@@ -243,37 +239,27 @@ def process_video(video_data):
                 logger.info("Processing time limit reached")
                 break
 
-            # Model inference
             results = model(frame, device=device)
-            seen_violations = set()  # Track violations detected in the current frame
-
             for result in results:
                 for box in result.boxes:
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                    
-                    # Skip if it's not a relevant violation or if confidence is too low
-                    if label not in CONFIG["VIOLATION_LABELS"].values() or conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                    if label not in CONFIG["VIOLATION_LABELS"].values():
                         continue
-                    
-                    # Skip if the same violation is detected again within the temporal threshold
-                    if time.time() - last_detected[label] < CONFIG["TEMPORAL_THRESHOLD"]:
+                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
+                    if label in seen_violations:
+                        continue
+                    seen_violations.add(label)
 
-                    # Update last detected time
-                    last_detected[label] = time.time()
-
-                    # Save the violation data
                     violation = {
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in box.xywh.cpu().numpy()[0]],
                         "timestamp": frame_count / fps
                     }
                     violations.append(violation)
 
-                    # Snapshot for the first occurrence of each violation
                     if not snapshot_taken[label]:
                         snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], f"snapshot_{frame_count}_{label}.jpg")
                         cv2.imwrite(snapshot_path, frame)
@@ -292,7 +278,6 @@ def process_video(video_data):
         video.release()
         os.remove(video_path)
 
-        # If no violations were detected, return a message
         if not violations:
             logger.info("No violations detected")
             return {
@@ -304,7 +289,6 @@ def process_video(video_data):
                 "message": "No violations detected in the video."
             }
 
-        # Calculate compliance score
         score = calculate_safety_score(violations)
         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)
@@ -330,8 +314,7 @@ def process_video(video_data):
 
 def gradio_interface(video_file):
     if not video_file:
-        return "", "", "", "", ""  # Return empty outputs when no file uploaded
-    
+        return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
         yield "Processing video... please wait.", "", "", "", ""
 
@@ -340,26 +323,25 @@ def gradio_interface(video_file):
 
         result = process_video(video_data)
 
-        # Only output violations if detected, else output nothing or minimal
-        if not result["violations"]:
-            # No violations detected — return empty or minimal outputs
-            yield "", "", "", "", ""
+        if result.get("message"):
+            yield result["message"], "", "", "", ""
             return
 
-        # Build violation table only if violations exist
-        header = "| Violation              | Timestamp (s) | Confidence |\n"
-        separator = "|------------------------|---------------|------------|\n"
-        rows = []
-        violation_name_map = CONFIG["DISPLAY_NAMES"]
-        for v in result["violations"]:
-            display_name = violation_name_map.get(v["violation"], v["violation"])
-            row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      |"
-            rows.append(row)
-        violation_table = header + separator + "\n".join(rows)
-
-        # Prepare snapshots if available
-        snapshots_text = ""
+        violation_table = "No violations detected."
+        if result["violations"]:
+            header = "| Violation              | Timestamp (s) | Confidence | \n"
+            separator = "|------------------------|---------------|------------|\n"
+            rows = []
+            violation_name_map = CONFIG["DISPLAY_NAMES"]
+            for v in result["violations"]:
+                display_name = violation_name_map.get(v["violation"], v["violation"])
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      |"
+                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 {violation_name_map.get(s['violation'], s['violation'])} at frame {s['frame']}: ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
@@ -374,7 +356,7 @@ def gradio_interface(video_file):
         )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "", "", ""
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
 interface = gr.Interface(
     fn=gradio_interface,