Update app.py

app.py

@@ -30,11 +30,11 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "no_helmet": (0, 0, 255),      # Red
-        "no_harness": (0, 165, 255),   # Orange
-        "unsafe_posture": (0, 255, 0), # Green
-        "unsafe_zone": (255, 0, 0),    # Blue
-        "improper_tool_use": (255, 255, 0) # Yellow
+        "no_helmet": (0, 0, 255),
+        "no_harness": (0, 165, 255),
+        "unsafe_posture": (0, 255, 0),
+        "unsafe_zone": (255, 0, 0),
+        "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -50,13 +50,7 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": {
-        "no_helmet": 2,
-        "no_harness": 1,
-        "unsafe_posture": 1,
-        "unsafe_zone": 1,
-        "improper_tool_use": 1
-    },
+    "FRAME_SKIP": 5,  # Increased to process fewer frames (1 frame every 5 frames)
     "CONFIDENCE_THRESHOLDS": {
         "no_helmet": 0.5,
         "no_harness": 0.15,
@@ -65,8 +59,9 @@ CONFIG = {
         "improper_tool_use": 0.15
     },
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_FRAMES": 3,
-    "HELMET_CONFIDENCE_THRESHOLD": 0.7
+    "MIN_VIOLATION_FRAMES": 2,  # Reduced to ensure violations are detected with fewer frames
+    "HELMET_CONFIDENCE_THRESHOLD": 0.7,
+    "MAX_PROCESSING_TIME": 30  # Enforce 30-second processing limit
 }
 
 # Setup logging
@@ -101,7 +96,6 @@ model = load_model()
 # Enhanced Helper Functions
 # ==========================
 def draw_detections(frame, detections):
-    """Draw bounding boxes and labels on frame"""
     for det in detections:
         label = det.get("violation", "Unknown")
         confidence = det.get("confidence", 0.0)
@@ -121,7 +115,6 @@ def draw_detections(frame, detections):
     return frame
 
 def calculate_iou(box1, box2):
-    """Calculate Intersection over Union (IoU) for two bounding boxes."""
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
@@ -299,34 +292,41 @@ def process_video(video_data):
         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
+            fps = 30
         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
+        confirmed_violations = {}
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        helmet_compliance = {}  # Track workers with helmets
-        detection_counts = {label: 0 for label in CONFIG["VIOLATION_LABELS"].values()}  # Track detection counts
+        helmet_compliance = {}
+        detection_counts = {label: 0 for label in CONFIG["VIOLATION_LABELS"].values()}
+        start_time = time.time()
+
+        # Calculate frames to process within 30 seconds
+        target_frames = int(total_frames / CONFIG["FRAME_SKIP"])
+        frame_indices = np.linspace(0, total_frames - 1, target_frames, dtype=int)
+
+        processed_frames = 0
+        for idx in frame_indices:
+            elapsed_time = time.time() - start_time
+            if elapsed_time > CONFIG["MAX_PROCESSING_TIME"]:
+                logger.info(f"Processing time limit of {CONFIG['MAX_PROCESSING_TIME']} seconds reached. Processed {processed_frames}/{target_frames} frames.")
+                break
 
-        while True:
+            video.set(cv2.CAP_PROP_POS_FRAMES, idx)
             ret, frame = video.read()
             if not ret:
-                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})", "", "", "", ""
+            processed_frames += 1
+            current_time = idx / fps
+            progress = (processed_frames / target_frames) * 100
+            yield f"Processing video... {progress:.1f}% complete (Frame {idx}/{total_frames})", "", "", "", ""
 
             # Run detection on this frame
-            results = model(frame, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], agnostic_nms=True)
+            results = model(frame, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"])
             
             current_detections = []
             for result in results:
@@ -347,7 +347,7 @@ def process_video(video_data):
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     
                     current_detections.append({
-                        "frame": frame_count,
+                        "frame": idx,
                         "violation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
@@ -355,16 +355,14 @@ def process_video(video_data):
                     })
                     detection_counts[label] += 1
 
-            logger.debug(f"Frame {frame_count}: Detected {len(current_detections)} violations: {[d['violation'] for d in current_detections]}")
+            logger.debug(f"Frame {idx}: Detected {len(current_detections)} violations: {[d['violation'] for d in current_detections]}")
 
-            # Process detections and associate with workers
             for detection in current_detections:
                 violation_type = detection.get("violation", None)
                 if violation_type is None:
                     logger.error(f"Invalid detection, missing 'violation' key: {detection}")
                     continue
 
-                # Helmet compliance check
                 if violation_type == "no_helmet":
                     matched_worker = None
                     max_iou = 0
@@ -376,7 +374,6 @@ def process_video(video_data):
                     
                     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
@@ -387,7 +384,6 @@ def process_video(video_data):
                             logger.debug(f"Worker {worker_id} has helmet, skipping no_helmet violation")
                             continue
                 
-                # Find or create worker
                 matched_worker = None
                 max_iou = 0
                 
@@ -409,11 +405,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]:
                     logger.debug(f"Violation {violation_type} already confirmed for worker {worker_id}, skipping")
                     continue
@@ -421,14 +415,10 @@ def process_video(video_data):
                 detection["worker_id"] = worker_id
                 violation_history[violation_type].append(detection)
 
-            # Clean up old workers
             workers = [w for w in workers if current_time - w["last_seen"] < 5.0]
 
-            frame_count += 1
-
         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:
                 logger.info(f"No detections for {violation_type}")
@@ -442,16 +432,13 @@ def process_video(video_data):
             
             for worker_id, worker_dets in worker_violations.items():
                 if len(worker_dets) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                    # Skip if already confirmed
                     if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
                         continue
                         
-                    # Skip No Helmet if worker is compliant
                     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
@@ -486,7 +473,6 @@ def process_video(video_data):
                         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")
@@ -537,7 +523,6 @@ def gradio_interface(video_file):
         with open(video_file, "rb") as f:
             video_data = f.read()
 
-        # Use generator to yield updates
         for status, score, snapshots_text, record_id, details_url in process_video(video_data):
             yield status, score, snapshots_text, record_id, details_url
     except Exception as e: