Update app.py

app.py

@@ -43,13 +43,14 @@ class BYTETracker:
     def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.5, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
-        self.match_thresh = match_thresh  # Increased to 0.5 for better matching
+        self.match_thresh = match_thresh
         self.frame_rate = frame_rate
         self.next_id = 1
         self.tracks = {}
         self.worker_history = {}
         self.last_positions = {}
         self.recently_removed = {}  # Store recently removed tracks for re-identification
+        self.helmet_status = {}  # Track helmet status for each worker
 
     def update(self, dets, scores, cls):
         tracks = []
@@ -108,6 +109,13 @@ class BYTETracker:
                     'cls': cl,
                     'last_seen': current_time
                 })
+                
+                # Update helmet status if this is a helmet detection
+                if cl == 0:  # Helmet violation class
+                    # Higher confidence for helmet violations
+                    if score > 0.45:  # Increased threshold for helmet violations
+                        self.helmet_status[best_track_id] = True
+                
                 if best_track_id not in self.worker_history:
                     self.worker_history[best_track_id] = []
                 self.worker_history[best_track_id].append([x, y])
@@ -132,6 +140,13 @@ class BYTETracker:
                         }
                         self.worker_history[track_id] = [[x, y]]
                         self.last_positions[track_id] = [x, y]
+                        
+                        # Update helmet status if this is a helmet detection
+                        if cl == 0:  # Helmet violation class
+                            # Higher confidence for helmet violations
+                            if score > 0.45:  # Increased threshold for helmet violations
+                                self.helmet_status[track_id] = True
+                        
                         tracks.append({
                             'id': track_id,
                             'bbox': [x, y, w, h],
@@ -153,6 +168,13 @@ class BYTETracker:
                                 'cls': cl,
                                 'last_seen': current_time
                             }
+                            
+                            # Update helmet status if this is a helmet detection
+                            if cl == 0:  # Helmet violation class
+                                # Higher confidence for helmet violations
+                                if score > 0.45:  # Increased threshold for helmet violations
+                                    self.helmet_status[worker_id] = True
+                            
                             tracks.append({
                                 'id': worker_id,
                                 'bbox': [x, y, w, h],
@@ -171,6 +193,13 @@ class BYTETracker:
                         }
                         self.worker_history[self.next_id] = [[x, y]]
                         self.last_positions[self.next_id] = [x, y]
+                        
+                        # Update helmet status if this is a helmet detection
+                        if cl == 0:  # Helmet violation class
+                            # Higher confidence for helmet violations
+                            if score > 0.45:  # Increased threshold for helmet violations
+                                self.helmet_status[self.next_id] = True
+                        
                         tracks.append({
                             'id': self.next_id,
                             'bbox': [x, y, w, h],
@@ -196,12 +225,17 @@ class BYTETracker:
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
     
-    def _is_same_worker(self, pos1, pos2, threshold=150):  # Increased threshold to 150
+    def _is_same_worker(self, pos1, pos2, threshold=150):
         x1, y1 = pos1
         x2, y2 = pos2
         distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
         return distance < threshold
 
+    # Function to validate if a helmet violation is consistent across frames
+    def validate_helmet_violation(self, worker_id, current_confidence):
+        # If we have consistent high confidence or multiple detections, it's a valid violation
+        return worker_id in self.helmet_status and self.helmet_status[worker_id]
+
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
@@ -235,25 +269,26 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.4,
+        "no_helmet": 0.45,  # Increased threshold for helmet violations
         "no_harness": 0.25,
         "unsafe_posture": 0.25,
         "unsafe_zone": 0.25,
         "improper_tool_use": 0.25
     },
-    "MIN_VIOLATION_FRAMES": 1,
+    "MIN_VIOLATION_FRAMES": 2,  # Increased to require multiple frames for confirmation
     "VIOLATION_COOLDOWN": 30.0,
-    "WORKER_TRACKING_DURATION": 10.0,  # Reverted to 5.0 seconds
+    "WORKER_TRACKING_DURATION": 10.0,
     "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 1,
     "BATCH_SIZE": 4,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 150,  # 5.0 seconds at 30 fps
+    "TRACK_BUFFER": 150,
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.5,  # Increased to 0.5
+    "MATCH_THRESH": 0.5,
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 150,  # Increased to match _is_same_worker threshold
-    "TARGET_RESOLUTION": (384, 384)
+    "MAX_WORKER_DISTANCE": 150,
+    "TARGET_RESOLUTION": (384, 384),
+    "HELMET_VALIDATION_FRAMES": 3  # Number of frames to validate helmet violations
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -285,8 +320,18 @@ model = load_model()
 # ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
     target_res = CONFIG["TARGET_RESOLUTION"]
+    # Enhanced preprocessing for better helmet detection
     frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
-    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
+    # Increase contrast to better differentiate helmets from other head coverings
+    frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=20)  # Increased contrast
+    
+    # Additional preprocessing to enhance head/helmet features
+    # Apply slight sharpening to make edges more distinct
+    kernel = np.array([[-1,-1,-1], 
+                       [-1, 9,-1],
+                       [-1,-1,-1]])
+    frame = cv2.filter2D(frame, -1, kernel)
+    
     return frame
 
 def draw_detections(frame, detections):
@@ -305,7 +350,10 @@ def draw_detections(frame, detections):
         
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         
-        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+        # Make no_helmet violations more prominent
+        line_thickness = 4 if label == "no_helmet" else 3
+        
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, line_thickness)
         
         display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
         text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
@@ -536,6 +584,83 @@ def verify_and_open_video(video_path):
     
     return cap
 
+# Helper for helmet validation
+def validate_helmet_detection(frame, bbox, confidence_threshold=0.45):
+    """
+    Additional validation for helmet detection to reduce false positives.
+    This function performs additional checks on the region to confirm it's a true helmet violation.
+    """
+    x, y, w, h = bbox
+    x1 = int(max(0, x - w/2))
+    y1 = int(max(0, y - h/2))
+    x2 = int(min(frame.shape[1], x + w/2))
+    y2 = int(min(frame.shape[0], y + h/2))
+    
+    # Extract head region
+    head_region = frame[y1:y2, x1:x2]
+    if head_region.size == 0:
+        return False
+    
+    # Check if this is truly a helmet violation by analyzing the region
+    # 1. Check color distribution - helmets often have more uniform color
+    hsv = cv2.cvtColor(head_region, cv2.COLOR_BGR2HSV)
+    
+    # Check for typical helmet colors (many construction helmets are yellow, white, orange, blue)
+    # This helps differentiate from cloth head coverings
+    yellow_lower = np.array([20, 100, 100])
+    yellow_upper = np.array([30, 255, 255])
+    yellow_mask = cv2.inRange(hsv, yellow_lower, yellow_upper)
+    
+    white_lower = np.array([0, 0, 200])
+    white_upper = np.array([180, 30, 255])
+    white_mask = cv2.inRange(hsv, white_lower, white_upper)
+    
+    orange_lower = np.array([5, 100, 100])
+    orange_upper = np.array([15, 255, 255])
+    orange_mask = cv2.inRange(hsv, orange_lower, orange_upper)
+    
+    blue_lower = np.array([100, 100, 100])
+    blue_upper = np.array([130, 255, 255])
+    blue_mask = cv2.inRange(hsv, blue_lower, blue_upper)
+    
+    helmet_mask = cv2.bitwise_or(yellow_mask, white_mask)
+    helmet_mask = cv2.bitwise_or(helmet_mask, orange_mask)
+    helmet_mask = cv2.bitwise_or(helmet_mask, blue_mask)
+    
+    # If there's a significant amount of helmet-colored pixels, this might be a helmet
+    helmet_percentage = np.sum(helmet_mask > 0) / (head_region.shape[0] * head_region.shape[1])
+    
+    # If the region has a significant amount of helmet-like colors, it's probably a helmet
+    # so we should NOT flag it as a violation (return False)
+    if helmet_percentage > 0.25:
+        return False
+    
+    # Check texture uniformity - helmets have more uniform texture compared to head coverings
+    gray = cv2.cvtColor(head_region, cv2.COLOR_BGR2GRAY)
+    texture_score = np.std(gray)
+    
+    # If texture is very uniform (low standard deviation), it might be a helmet or bare head
+    # Very uniform texture (like a hard helmet) would have low texture_score
+    if texture_score < 15:  # Low texture suggests uniform surface like a helmet
+        return False
+    
+    # Additional check for cloth-like textures
+    edges = cv2.Canny(gray, 50, 150)
+    edge_density = np.sum(edges > 0) / (head_region.shape[0] * head_region.shape[1])
+    
+    # If there are many edges (cloth wrinkles), this might be a kurchief
+    if edge_density > 0.15:
+        # This is likely a cloth head covering, not a helmet violation
+        # But also not a proper helmet, so we should still detect as violation
+        return True
+    
+    # If confidence is very high, trust the model
+    if confidence_threshold >= 0.6:
+        return True
+        
+    # Default to the original detection
+    return True
+
 def process_video(video_data, temp_dir):
     video_path = None
     output_dir = os.path.join(temp_dir, "output")
@@ -586,6 +711,8 @@ def process_video(video_data, temp_dir):
         worker_id_mapping = {}
         unique_violations = {}
         violation_frames = {}
+        # Track helmet detections across frames for each worker
+        helmet_detections = {}  
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
@@ -595,6 +722,7 @@ def process_video(video_data, temp_dir):
         while processed_frames < total_frames:
             batch_frames = []
             batch_indices = []
+            batch_originals = []  # Store original frames for helmet validation
             
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
@@ -606,6 +734,9 @@ def process_video(video_data, temp_dir):
                     logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
                     break
                 
+                # Store original frame for validation
+                original_frame = frame.copy()
+                
                 frame = preprocess_frame(frame)
                 
                 for _ in range(frame_skip - 1):
@@ -614,6 +745,7 @@ def process_video(video_data, temp_dir):
                 
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
+                batch_originals.append(original_frame)
                 processed_frames += 1
 
             if not batch_frames:
@@ -644,7 +776,7 @@ def process_video(video_data, temp_dir):
                 yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames}, {fps_processed:.1f} FPS)", "", "", "", ""
                 last_yield_time = current_time
 
-            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
+            for i, (result, frame_idx, original_frame) in enumerate(zip(results, batch_indices, batch_originals)):
                 current_time = frame_idx / fps
                 
                 boxes = result.boxes
@@ -658,8 +790,20 @@ def process_video(video_data, temp_dir):
                     if label is None:
                         continue
                         
-                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
-                        continue
+                    # Enhanced confidence threshold handling, especially for helmet detection
+                    if label == "no_helmet":
+                        if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.45):
+                            continue
+                        
+                        # Additional validation for helmet detection
+                        bbox = box.xywh.cpu().numpy()[0]
+                        if not validate_helmet_detection(original_frame, bbox, conf):
+                            logger.info(f"Frame {frame_idx}: Helmet false positive filtered at {conf:.2f} confidence")
+                            continue
+                    else:
+                        # Use regular thresholds for other violations
+                        if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                            continue
 
                     bbox = box.xywh.cpu().numpy()[0]
                     track_inputs.append({
@@ -693,11 +837,37 @@ def process_video(video_data, temp_dir):
                     
                     worker_id = worker_id_mapping[tracker_id]
                     
-                    violation_key = (worker_id, label)
-                    
-                    if violation_key not in unique_violations:
-                        unique_violations[violation_key] = current_time
-                        violation_frames[violation_key] = frame_idx
+                    # Special handling for helmet violations to ensure consistency
+                    if label == "no_helmet":
+                        # Track helmet violations for this worker
+                        if worker_id not in helmet_detections:
+                            helmet_detections[worker_id] = []
+                        
+                        # Store this detection with frame index and confidence
+                        helmet_detections[worker_id].append({
+                            "frame_idx": frame_idx,
+                            "confidence": conf,
+                            "bbox": bbox
+                        })
+                        
+                        # Only record a helmet violation if we have multiple consistent detections
+                        if len(helmet_detections[worker_id]) >= CONFIG["HELMET_VALIDATION_FRAMES"]:
+                            # Calculate average confidence
+                            avg_conf = sum(d["confidence"] for d in helmet_detections[worker_id]) / len(helmet_detections[worker_id])
+                            
+                            # If confidence is consistently high across multiple frames, record the violation
+                            if avg_conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                                violation_key = (worker_id, label)
+                                if violation_key not in unique_violations:
+                                    unique_violations[violation_key] = current_time
+                                    violation_frames[violation_key] = frame_idx
+                                    logger.info(f"Frame {frame_idx}: Valid helmet violation for worker {worker_id} with avg conf {avg_conf:.2f}")
+                    else:
+                        # Regular handling for other violations
+                        violation_key = (worker_id, label)
+                        if violation_key not in unique_violations:
+                            unique_violations[violation_key] = current_time
+                            violation_frames[violation_key] = frame_idx
 
         cap.release()
         processing_time = time.time() - start_time