Update app.py

app.py

@@ -51,22 +51,22 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "FRAME_SKIP": {
-        "no_helmet": 3,
-        "no_harness": 2,
-        "unsafe_posture": 2,
-        "unsafe_zone": 2,
-        "improper_tool_use": 2
+        "no_helmet": 2,  # Reduced to process more frames
+        "no_harness": 1,
+        "unsafe_posture": 1,
+        "unsafe_zone": 1,
+        "improper_tool_use": 1
     },
-    "MAX_PROCESSING_TIME": 60,
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.3,
-        "no_harness": 0.2,
-        "unsafe_posture": 0.2,
-        "unsafe_zone": 0.2,
-        "improper_tool_use": 0.2
+        "no_helmet": 0.5,  # Increased to reduce false positives
+        "no_harness": 0.15,  # Lowered to improve detection
+        "unsafe_posture": 0.15,
+        "unsafe_zone": 0.15,
+        "improper_tool_use": 0.15
     },
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3
+    "MIN_VIOLATION_FRAMES": 3,
+    "HELMET_CONFIDENCE_THRESHOLD": 0.7  # Require high confidence for No Helmet violation
 }
 
 # Setup logging
@@ -296,35 +296,35 @@ def process_video(video_data):
         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)
         if fps <= 0:
             fps = 30  # Default assumption if FPS cannot be determined
+        video_duration = total_frames / fps
+        logger.info(f"Video duration: {video_duration:.2f} seconds, Total frames: {total_frames}, FPS: {fps}")
 
         workers = []
         violation_history = {label: [] for label in CONFIG["VIOLATION_LABELS"].values()}
         confirmed_violations = {}  # Track confirmed violations per worker
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         helmet_compliance = {}  # Track workers with helmets
-
-        logger.info(f"Processing video with FPS: {fps}")
-        logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
+        detection_counts = {label: 0 for label in CONFIG["VIOLATION_LABELS"].values()}  # Track detection counts
 
         while True:
             ret, frame = video.read()
             if not ret:
                 break
 
-            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
-                logger.info("Processing time limit reached")
-                break
-
             current_time = frame_count / fps
             min_frame_skip = min(CONFIG["FRAME_SKIP"].values())
             if frame_count % min_frame_skip != 0:
                 frame_count += 1
                 continue
 
+            # Yield progress update
+            progress = (frame_count / total_frames) * 100
+            yield f"Processing video... {progress:.1f}% complete (Frame {frame_count}/{total_frames})", "", "", "", ""
+
             # Run detection on this frame
             results = model(frame, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], agnostic_nms=True)
             
@@ -348,11 +348,12 @@ def process_video(video_data):
                     
                     current_detections.append({
                         "frame": frame_count,
-                        "violation": label,  # Corrected key
+                        "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
                         "timestamp": current_time
                     })
+                    detection_counts[label] += 1
 
             logger.debug(f"Frame {frame_count}: Detected {len(current_detections)} violations: {[d['violation'] for d in current_detections]}")
 
@@ -363,7 +364,7 @@ def process_video(video_data):
                     logger.error(f"Invalid detection, missing 'violation' key: {detection}")
                     continue
 
-                # Skip No Helmet detection if worker is compliant
+                # Helmet compliance check
                 if violation_type == "no_helmet":
                     matched_worker = None
                     max_iou = 0
@@ -373,9 +374,18 @@ def process_video(video_data):
                             max_iou = iou
                             matched_worker = worker
                     
-                    if matched_worker and matched_worker["id"] in helmet_compliance:
-                        logger.debug(f"Worker {matched_worker['id']} has helmet, skipping no_helmet violation")
-                        continue
+                    if matched_worker:
+                        worker_id = matched_worker["id"]
+                        # Require high confidence and persistence for No Helmet violation
+                        if worker_id not in helmet_compliance:
+                            helmet_compliance[worker_id] = {"no_helmet_frames": 0, "compliant": False}
+                        helmet_compliance[worker_id]["no_helmet_frames"] += 1
+                        if detection["confidence"] < CONFIG["HELMET_CONFIDENCE_THRESHOLD"]:
+                            helmet_compliance[worker_id]["compliant"] = True
+                            logger.debug(f"Worker {worker_id} marked as helmet compliant due to low no_helmet confidence")
+                        if helmet_compliance[worker_id]["compliant"]:
+                            logger.debug(f"Worker {worker_id} has helmet, skipping no_helmet violation")
+                            continue
                 
                 # Find or create worker
                 matched_worker = None
@@ -399,6 +409,9 @@ def process_video(video_data):
                         "first_seen": current_time,
                         "last_seen": current_time
                     })
+                    # Initialize helmet compliance for new worker
+                    if worker_id not in helmet_compliance:
+                        helmet_compliance[worker_id] = {"no_helmet_frames": 0, "compliant": False}
                 
                 # Skip if this violation type is already confirmed for this worker
                 if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
@@ -407,20 +420,14 @@ def process_video(video_data):
                 
                 detection["worker_id"] = worker_id
                 violation_history[violation_type].append(detection)
-                
-                # Update helmet compliance (simulate by checking if No Helmet confidence is low)
-                if violation_type == "no_helmet" and detection["confidence"] < 0.5:
-                    helmet_compliance[worker_id] = True
-                    logger.debug(f"Worker {worker_id} marked as helmet compliant")
 
             # Clean up old workers
             workers = [w for w in workers if current_time - w["last_seen"] < 5.0]
 
             frame_count += 1
 
-        video.release()
-        os.remove(video_path)
-        
+        logger.info(f"Detection counts: {detection_counts}")
+
         # Process violation history to confirm persistent violations
         for violation_type, detections in violation_history.items():
             if not detections:
@@ -440,9 +447,14 @@ def process_video(video_data):
                         continue
                         
                     # Skip No Helmet if worker is compliant
-                    if violation_type == "no_helmet" and worker_id in helmet_compliance:
-                        logger.debug(f"Skipping no_helmet for worker {worker_id} due to helmet compliance")
-                        continue
+                    if violation_type == "no_helmet":
+                        if worker_id in helmet_compliance and helmet_compliance[worker_id]["compliant"]:
+                            logger.debug(f"Skipping no_helmet for worker {worker_id} due to helmet compliance")
+                            continue
+                        # Require persistent No Helmet detections
+                        if helmet_compliance[worker_id]["no_helmet_frames"] < CONFIG["MIN_VIOLATION_FRAMES"] * 2:
+                            logger.debug(f"Skipping no_helmet for worker {worker_id}, not enough persistent detections")
+                            continue
                     
                     best_detection = max(worker_dets, key=lambda x: x["confidence"])
                     violations.append(best_detection)
@@ -455,22 +467,29 @@ def process_video(video_data):
                         cap = cv2.VideoCapture(video_path)
                         cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
                         ret, snapshot_frame = cap.read()
-                        cap.release()
+                        if not ret:
+                            logger.error(f"Failed to capture snapshot for {violation_type} at frame {best_detection['frame']}")
+                            cap.release()
+                            continue
+                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
                         
-                        if ret:
-                            snapshot_frame = draw_detections(snapshot_frame, [best_detection])
-                            
-                            snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
-                            snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                            cv2.imwrite(snapshot_path, snapshot_frame)
-                            snapshots.append({
-                                "violation": violation_type,
-                                "frame": best_detection["frame"],
-                                "snapshot_path": snapshot_path,
-                                "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                            })
-                            snapshot_taken[violation_type] = True
-                            logger.info(f"Snapshot taken for {violation_type} at frame {best_detection['frame']}")
+                        snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                        cv2.imwrite(snapshot_path, snapshot_frame)
+                        snapshots.append({
+                            "violation": violation_type,
+                            "frame": best_detection["frame"],
+                            "snapshot_path": snapshot_path,
+                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                        })
+                        snapshot_taken[violation_type] = True
+                        logger.info(f"Snapshot taken for {violation_type} at frame {best_detection['frame']}")
+                        cap.release()
+
+        # Clean up video file after snapshots are captured
+        video.release()
+        os.remove(video_path)
+        logger.info(f"Video file {video_path} removed")
 
         if not violations:
             logger.info("No persistent violations detected")
@@ -514,44 +533,46 @@ def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
     try:
-        yield "Processing video... please wait.", "", "", "", ""
-
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        result = process_video(video_data)
-
-        if result.get("message"):
-            yield result["message"], "", "", "", ""
-            return
-
-        violation_table = "No violations detected."
-        if result["violations"]:
-            header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
-            separator = "|------------------------|---------------|------------|-----------|\n"
-            rows = []
-            violation_name_map = CONFIG["DISPLAY_NAMES"]
-            for v in result["violations"]:
-                display_name = violation_name_map.get(v.get("violation", "Unknown"), "Unknown")
-                row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |"
-                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.get('violation', 'Unknown'), 'Unknown')} at frame {s.get('frame', 0)}: ![]({s.get('snapshot_base64', '')})"
-                for s in result["snapshots"]
-            )
+        # Use generator to yield progress updates
+        for status, violations_table, score, snapshots_text, record_id, details_url in process_video(video_data):
+            if status.startswith("Processing video"):
+                yield status, "", "", "", ""
+                continue
 
-        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"
-        )
+            if status.get("message"):
+                yield status["message"], "", "", "", ""
+                return
+
+            violation_table = "No violations detected."
+            if status["violations"]:
+                header = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
+                separator = "|------------------------|---------------|------------|-----------|\n"
+                rows = []
+                violation_name_map = CONFIG["DISPLAY_NAMES"]
+                for v in status["violations"]:
+                    display_name = violation_name_map.get(v.get("violation", "Unknown"), "Unknown")
+                    row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |"
+                    rows.append(row)
+                violation_table = header + separator + "\n".join(rows)
+
+            snapshots_text = "No snapshots captured."
+            if status["snapshots"]:
+                violation_name_map = CONFIG["DISPLAY_NAMES"]
+                snapshots_text = "\n".join(
+                    f"- Snapshot for {violation_name_map.get(s.get('violation', 'Unknown'), 'Unknown')} at frame {s.get('frame', 0)}: ![]({s.get('snapshot_base64', '')})"
+                    for s in status["snapshots"]
+                )
+
+            yield (
+                violation_table,
+                f"Safety Score: {status['score']}%",
+                snapshots_text,
+                f"Salesforce Record ID: {status['salesforce_record_id'] or 'N/A'}",
+                status["violation_details_url"] or "N/A"
+            )
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""