Update app.py

app.py

@@ -40,8 +40,8 @@ CONFIG = {
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": 15,
     "MAX_PROCESSING_TIME": 30,
-    "CONFIDENCE_THRESHOLD": 0.3,  # Lowered threshold for detecting all violations
-    "IOU_THRESHOLD": 0.5  # Added for worker tracking
+    "CONFIDENCE_THRESHOLD": 0.25,  # Lowered for better detection
+    "IOU_THRESHOLD": 0.5  # For worker tracking
 }
 
 # Setup logging
@@ -80,19 +80,12 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    # Convert to top-left and bottom-right coordinates
     x1_min, y1_min = x1 - w1/2, y1 - h1/2
     x1_max, y1_max = x1 + w1/2, y1 + h1/2
     x2_min, y2_min = x2 - w2/2, y2 - h2/2
     x2_max, y2_max = x2 + w2/2, y2 + h2/2
     
-    # Calculate intersection
-    x_min = max(x1_min, x2_min)
-    y_min = max(y1_min, y2_min)
-    x_max = min(x1_max, x2_max)
-    y_max = min(y1_max, y2_max)
-    
-    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    intersection = max(0, x2_min - x1_max) * max(0, y2_min - y1_max)
     area1 = w1 * h1
     area2 = w2 * h2
     union = area1 + area2 - intersection
@@ -141,7 +134,7 @@ def generate_violation_pdf(violations, score):
         else:
             for v in violations:
                 display_name = CONFIG["DISPLAY_NAMES"].get(v["violation"], v["violation"])
-                text = f"{display_name} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+                text = f"{display_name} at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f}, Worker ID: {v['worker_id']})"
                 c.drawString(1 * inch, y_position, text)
                 y_position -= 0.3 * inch
                 if y_position < 1 * inch:
@@ -190,7 +183,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
         violations_text = "\n".join(
-            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']})"
+            f"{CONFIG['DISPLAY_NAMES'].get(v['violation'], v['violation'])} at {v['timestamp']:.2f}s (Confidence: {v['confidence']:.2f}, Worker ID: {v['worker_id']})"
             for v in violations
         ) or "No violations detected."
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
@@ -260,7 +253,6 @@ def process_video(video_data):
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         workers = []  # List to track workers
 
-        # Adding debug logging for violation labels
         logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
         logger.info(f"Using confidence threshold: {CONFIG['CONFIDENCE_THRESHOLD']}")
 
@@ -277,11 +269,9 @@ def process_video(video_data):
                 logger.info("Processing time limit reached")
                 break
 
-            # Run detection on this frame
             results = model(frame, device=device)
             current_detections = []
             
-            # Process detections from the model
             for result in results:
                 boxes = result.boxes
                 logger.info(f"Frame {frame_count}: Found {len(boxes)} potential detections")
@@ -290,10 +280,8 @@ def process_video(video_data):
                     cls, conf = int(box.cls), float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
-                    # Enhanced logging
                     logger.info(f"Detection: class={cls}, conf={conf:.2f}, label={label}")
                     
-                    # Skip if not a known violation or below confidence threshold
                     if label not in CONFIG["VIOLATION_LABELS"].values():
                         logger.info(f"Skipping unknown class: {cls}")
                         continue
@@ -302,7 +290,6 @@ def process_video(video_data):
                         logger.info(f"Skipping low confidence: {conf:.2f} < {CONFIG['CONFIDENCE_THRESHOLD']}")
                         continue
 
-                    # Process valid detection
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     logger.info(f"Valid detection: {label} with confidence: {conf:.2f}")
                     
@@ -314,13 +301,10 @@ def process_video(video_data):
                         "frame": frame_count
                     })
 
-            # Process detections and associate with workers
-            # FIXED: Improved worker tracking logic
             for detection in current_detections:
                 matched_worker = None
                 max_iou = 0
                 
-                # Try to match with existing workers
                 for worker in workers:
                     iou = calculate_iou(detection["bounding_box"], worker["bbox"])
                     if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
@@ -328,9 +312,7 @@ def process_video(video_data):
                         matched_worker = worker
 
                 if matched_worker:
-                    # Update existing worker
                     if detection["violation"] not in matched_worker["violations"]:
-                        # New violation for this worker
                         logger.info(f"New violation for worker {matched_worker['id']}: {detection['violation']}")
                         matched_worker["violations"].add(detection["violation"])
                         violations.append({
@@ -342,7 +324,6 @@ def process_video(video_data):
                             "worker_id": matched_worker["id"]
                         })
                         
-                        # Save snapshot for this violation type if not already taken
                         if not snapshot_taken[detection["violation"]]:
                             snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
                             snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
@@ -352,14 +333,14 @@ def process_video(video_data):
                                 "violation": detection["violation"],
                                 "frame": frame_count,
                                 "snapshot_path": snapshot_path,
-                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
+                                "timestamp": detection["timestamp"],
+                                "confidence": detection["confidence"]
                             })
                     
-                    # Update worker position
                     matched_worker["bbox"] = detection["bounding_box"]
                     matched_worker["last_frame"] = frame_count
                 else:
-                    # New worker detected
                     worker_id = len(workers) + 1
                     logger.info(f"New worker {worker_id} with violation: {detection['violation']}")
                     workers.append({
@@ -378,7 +359,6 @@ def process_video(video_data):
                         "worker_id": worker_id
                     })
                     
-                    # Save snapshot for this violation type if not already taken
                     if not snapshot_taken[detection["violation"]]:
                         snapshot_filename = f"{detection['violation']}_{frame_count}.jpg"
                         snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
@@ -388,10 +368,11 @@ def process_video(video_data):
                             "violation": detection["violation"],
                             "frame": frame_count,
                             "snapshot_path": snapshot_path,
-                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
+                            "timestamp": detection["timestamp"],
+                            "confidence": detection["confidence"]
                         })
 
-            # Clean up workers that haven't been seen for a while
             active_workers = [w for w in workers if frame_count - w["last_frame"] < CONFIG["FRAME_SKIP"] * 5]
             if len(active_workers) != len(workers):
                 logger.info(f"Cleaned up {len(workers) - len(active_workers)} inactive workers")
@@ -402,7 +383,6 @@ def process_video(video_data):
         video.release()
         os.remove(video_path)
         
-        # Final log of violations detected
         violation_types = {}
         for v in violations:
             violation_types[v["violation"]] = violation_types.get(v["violation"], 0) + 1
@@ -445,9 +425,9 @@ def process_video(video_data):
 
 def gradio_interface(video_file):
     if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", ""
+        return "No file uploaded.", "", "No file uploaded.", "", "", []
     try:
-        yield "Processing video... please wait.", "", "", "", ""
+        yield "Processing video... please wait.", "", "", "", "", []
 
         with open(video_file, "rb") as f:
             video_data = f.read()
@@ -455,52 +435,56 @@ def gradio_interface(video_file):
         result = process_video(video_data)
 
         if result.get("message"):
-            yield result["message"], "", "", "", ""
+            yield result["message"], "", "", "", "", []
             return
 
         violation_table = "No violations detected."
         if result["violations"]:
-            header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
-            separator = "|------------------------|---------------|------------|-----------|\n"
+            header = "| Violation              | Timestamp (s) | Confidence | Worker ID | Frame |\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}      | {v['worker_id']} |"
+                row = f"| {display_name:<22} | {v['timestamp']:.2f}        | {v['confidence']:.2f}      | {v['worker_id']}        | {v['frame']} |"
                 rows.append(row)
             violation_table = header + separator + "\n".join(rows)
 
         snapshots_text = "No snapshots captured."
+        snapshot_images = []
         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']})"
+                f"- Snapshot for {violation_name_map.get(s['violation'], s['violation'])} at {s['timestamp']:.2f}s (Frame {s['frame']}, Confidence: {s['confidence']:.2f}): ![]({s['snapshot_base64']})"
                 for s in result["snapshots"]
             )
+            snapshot_images = [s["snapshot_base64"] for s in result["snapshots"]]
 
         yield (
             violation_table,
             f"Safety Score: {result['score']}%",
             snapshots_text,
             f"Salesforce Record ID: {result['salesforce_record_id'] or 'N/A'}",
-            result["violation_details_url"] or "N/A"
+            result["violation_details_url"] or "N/A",
+            snapshot_images
         )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
-        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+        yield f"Error: {str(e)}", "", "Error in processing.", "", "", []
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
-    outputs=[ 
+    outputs=[
         gr.Markdown(label="Detected Safety Violations"),
         gr.Textbox(label="Compliance Score"),
         gr.Markdown(label="Snapshots"),
         gr.Textbox(label="Salesforce Record ID"),
-        gr.Textbox(label="Violation Details URL")
+        gr.Textbox(label="Violation Details URL"),
+        gr.Gallery(label="Violation Snapshots")
     ],
     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, No Harness, Unsafe Posture). Non-violations are ignored.",
     allow_flagging="never"
 )