Update app.py

app.py

@@ -43,25 +43,24 @@ 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
+        self.match_thresh = match_thresh  # Increased to 0.5 for better matching
         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 = []
         current_time = time.time()
-        
+       
         # Prune stale tracks
         stale_ids = []
         for track_id, track_info in self.tracks.items():
             if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
                 stale_ids.append(track_id)
-        
+       
         for track_id in stale_ids:
             # Store recently removed tracks for re-identification (for 1 second)
             self.recently_removed[track_id] = {
@@ -86,22 +85,22 @@ class BYTETracker:
         for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
                 continue
-                
+               
             x, y, w, h = det
             matched = False
             best_iou = 0
             best_track_id = None
-            
+           
             # Try to match with active tracks
             for track_id, track_info in self.tracks.items():
                 tx, ty, tw, th = track_info['bbox']
                 iou = self._calculate_iou([x, y, w, h], [tx, ty, tw, th])
-                
+               
                 if iou > self.match_thresh and iou > best_iou:
                     best_iou = iou
                     best_track_id = track_id
                     matched = True
-            
+           
             if matched:
                 self.tracks[best_track_id].update({
                     'bbox': [x, y, w, h],
@@ -109,18 +108,11 @@ 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])
                 self.last_positions[best_track_id] = [x, y]
-                
+               
                 tracks.append({
                     'id': best_track_id,
                     'bbox': [x, y, w, h],
@@ -140,13 +132,6 @@ 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],
@@ -156,7 +141,7 @@ class BYTETracker:
                         reidentified = True
                         del self.recently_removed[track_id]
                         break
-                
+               
                 if not reidentified:
                     # Check if it matches an existing worker by position
                     same_worker = False
@@ -168,13 +153,6 @@ 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],
@@ -183,7 +161,7 @@ class BYTETracker:
                             })
                             same_worker = True
                             break
-                    
+                   
                     if not same_worker:
                         self.tracks[self.next_id] = {
                             'bbox': [x, y, w, h],
@@ -193,13 +171,6 @@ 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],
@@ -207,7 +178,7 @@ class BYTETracker:
                             'cls': cl
                         })
                         self.next_id += 1
-        
+       
         return tracks
 
     def _calculate_iou(self, box1, box2):
@@ -224,18 +195,13 @@ class BYTETracker:
         box2_area = w2 * h2
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
-    
-    def _is_same_worker(self, pos1, pos2, threshold=150):
+   
+    def _is_same_worker(self, pos1, pos2, threshold=150):  # Increased threshold to 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",
@@ -269,26 +235,25 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.45,  # Increased threshold for helmet violations
+        "no_helmet": 0.4,
         "no_harness": 0.25,
         "unsafe_posture": 0.25,
         "unsafe_zone": 0.25,
         "improper_tool_use": 0.25
     },
-    "MIN_VIOLATION_FRAMES": 2,  # Increased to require multiple frames for confirmation
+    "MIN_VIOLATION_FRAMES": 1,
     "VIOLATION_COOLDOWN": 30.0,
-    "WORKER_TRACKING_DURATION": 10.0,
+    "WORKER_TRACKING_DURATION": 10.0,  # Reverted to 5.0 seconds
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,  # Increased frame skip for faster processing
-    "BATCH_SIZE": 8,  # Increased batch size for better GPU utilization
+    "FRAME_SKIP": 1,
+    "BATCH_SIZE": 4,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 150,
+    "TRACK_BUFFER": 150,  # 5.0 seconds at 30 fps
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.5,
+    "MATCH_THRESH": 0.5,  # Increased to 0.5
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 150,
-    "TARGET_RESOLUTION": (384, 384),
-    "HELMET_VALIDATION_FRAMES": 3  # Number of frames to validate helmet violations
+    "MAX_WORKER_DISTANCE": 150,  # Increased to match _is_same_worker threshold
+    "TARGET_RESOLUTION": (384, 384)
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -305,7 +270,7 @@ def load_model():
             if not os.path.isfile(model_path):
                 logger.info(f"Downloading fallback model: {model_path}")
                 torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
-                
+               
         model = YOLO(model_path).to(device)
         if device.type == "cuda":
             model.model.half()
@@ -320,23 +285,13 @@ 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)
-    # 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)
-    
+    frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
     return frame
 
 def draw_detections(frame, detections):
     result_frame = frame.copy()
-    
+   
     for det in detections:
         label = det.get("violation", "Unknown")
         confidence = det.get("confidence", 0.0)
@@ -347,22 +302,19 @@ def draw_detections(frame, detections):
         y1 = int(y - h/2)
         x2 = int(x + w/2)
         y2 = int(y + h/2)
-        
+       
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
-        
-        # 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)
-        
+       
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+       
         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]
         cv2.rectangle(result_frame, (x1, y1-text_size[1]-10), (x1+text_size[0]+10, y1), (0, 0, 0), -1)
         cv2.putText(result_frame, display_text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-        
+       
         conf_text = f"Conf: {confidence:.2f}"
         cv2.putText(result_frame, conf_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-        
+       
     return result_frame
 
 def calculate_safety_score(violations):
@@ -373,21 +325,21 @@ def calculate_safety_score(violations):
         "unsafe_zone": 35,
         "improper_tool_use": 25
     }
-    
+   
     worker_violations = {}
     for v in violations:
         worker_id = v.get("worker_id", "Unknown")
         violation_type = v.get("violation", "Unknown")
-        
+       
         if worker_id not in worker_violations:
             worker_violations[worker_id] = set()
         worker_violations[worker_id].add(violation_type)
-    
+   
     total_penalty = 0
     for worker_violations_set in worker_violations.values():
         worker_penalty = sum(penalties.get(v, 0) for v in worker_violations_set)
         total_penalty += worker_penalty
-    
+   
     score = max(0, 100 - total_penalty)
     return score
 
@@ -397,14 +349,14 @@ def generate_violation_pdf(violations, score, output_dir):
         pdf_path = os.path.join(output_dir, pdf_filename)
         pdf_file = BytesIO()
         c = canvas.Canvas(pdf_file, pagesize=letter)
-        
+       
         c.setFont("Helvetica-Bold", 16)
         c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
-        
+       
         c.setFont("Helvetica", 12)
         c.drawString(1 * inch, 9.5 * inch, f"Date: {time.strftime('%Y-%m-%d')}")
         c.drawString(1 * inch, 9.2 * inch, f"Time: {time.strftime('%H:%M:%S')}")
-        
+       
         c.setFont("Helvetica-Bold", 14)
         c.drawString(1 * inch, 8.7 * inch, f"Safety Compliance Score: {score}%")
 
@@ -412,21 +364,21 @@ def generate_violation_pdf(violations, score, output_dir):
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Summary:")
         y_position -= 0.3 * inch
-        
+       
         worker_violations = {}
         for v in violations:
             worker_id = v.get("worker_id", "Unknown")
             if worker_id not in worker_violations:
                 worker_violations[worker_id] = []
             worker_violations[worker_id].append(v)
-        
+       
         c.setFont("Helvetica", 10)
         summary_data = {
             "Total Workers with Violations": len(worker_violations),
             "Total Violations Found": len(violations),
             "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
-        
+       
         for key, value in summary_data.items():
             c.drawString(1 * inch, y_position, f"{key}: {value}")
             y_position -= 0.25 * inch
@@ -435,21 +387,21 @@ def generate_violation_pdf(violations, score, output_dir):
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Violations by Worker:")
         y_position -= 0.3 * inch
-        
+       
         c.setFont("Helvetica", 10)
         for worker_id, worker_vios in worker_violations.items():
             c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
             y_position -= 0.2 * inch
-            
+           
             for v in worker_vios:
                 display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
                 time_str = f"{v.get('timestamp', 0.0):.2f}s"
                 conf_str = f"{v.get('confidence', 0.0):.2f}"
-                
+               
                 violation_text = f"  - {display_name} at {time_str} (Confidence: {conf_str})"
                 c.drawString(1.2 * inch, y_position, violation_text)
                 y_position -= 0.2 * inch
-                
+               
                 if y_position < 1 * inch:
                     c.showPage()
                     c.setFont("Helvetica", 10)
@@ -460,7 +412,7 @@ def generate_violation_pdf(violations, score, output_dir):
 
         with open(pdf_path, "wb") as f:
             f.write(pdf_file.getvalue())
-            
+           
         public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
         logger.info(f"PDF generated: {public_url}")
         return pdf_path, public_url, pdf_file
@@ -484,7 +436,7 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         if not pdf_file:
             logger.error("No PDF file provided for upload")
             return ""
-            
+           
         encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
         content_version_data = {
             "Title": f"Safety_Violation_Report_{int(time.time())}",
@@ -494,11 +446,11 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         }
         content_version = sf.ContentVersion.create(content_version_data)
         result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
-        
+       
         if not result['records']:
             logger.error("Failed to retrieve ContentVersion")
             return ""
-            
+           
         file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
         logger.info(f"PDF uploaded to Salesforce: {file_url}")
         return file_url
@@ -509,19 +461,19 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
-        
+       
         violations_text = ""
         for v in violations:
             display_name = CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')
             worker_id = v.get('worker_id', 'Unknown')
             timestamp = v.get('timestamp', 0.0)
             confidence = v.get('confidence', 0.0)
-            
+           
             violations_text += f"Worker {worker_id}: {display_name} at {timestamp:.2f}s (Conf: {confidence:.2f})\n"
-        
+       
         if not violations_text:
             violations_text = "No violations detected."
-            
+           
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
 
         record_data = {
@@ -531,9 +483,9 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             "Status__c": "Pending",
             "PDF_Report_URL__c": pdf_url
         }
-        
+       
         logger.info(f"Creating Salesforce record with data: {record_data}")
-        
+       
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
             logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
@@ -541,7 +493,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
             logger.warning(f"Fell back to Account record: {record['id']}")
-            
+           
         record_id = record["id"]
 
         if pdf_file:
@@ -570,107 +522,30 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 def verify_and_open_video(video_path):
     if not os.path.exists(video_path):
         raise FileNotFoundError(f"Temporary video file not found: {video_path}")
-    
+   
     file_size = os.path.getsize(video_path)
     if file_size == 0:
         raise ValueError(f"Temporary video file is empty: {video_path}")
-    
+   
     with open(video_path, "rb") as f:
         f.read(1)
-    
+   
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         raise ValueError("Could not open video file. Ensure the video format is supported (e.g., MP4) and FFmpeg is installed.")
-    
+   
     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")
     os.makedirs(output_dir, exist_ok=True)
     os.environ['YOLO_CONFIG_DIR'] = temp_dir
-    
+   
     try:
         if not video_data:
             raise ValueError("Empty video data provided.")
-        
+       
         logger.info(f"Received video data size: {len(video_data)} bytes")
         if len(video_data) == 0:
             raise ValueError("Video data is empty.")
@@ -711,8 +586,7 @@ def process_video(video_data, temp_dir):
         worker_id_mapping = {}
         unique_violations = {}
         violation_frames = {}
-        # Track helmet detections across frames for each worker
-        helmet_detections = {}  
+        worker_violation_count = {}  # Track violation count per worker
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
@@ -722,30 +596,25 @@ 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))
                 if frame_idx >= total_frames:
                     break
-                
+               
                 ret, frame = cap.read()
                 if not ret:
                     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):
                     if not cap.grab():
                         break
-                
+               
                 batch_frames.append(frame)
                 batch_indices.append(frame_idx)
-                batch_originals.append(original_frame)
                 processed_frames += 1
 
             if not batch_frames:
@@ -776,34 +645,22 @@ 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, original_frame) in enumerate(zip(results, batch_indices, batch_originals)):
+            for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
-                
+               
                 boxes = result.boxes
                 track_inputs = []
-                
+               
                 for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                    
+                   
                     if label is None:
                         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
+                       
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        continue
 
                     bbox = box.xywh.cpu().numpy()[0]
                     track_inputs.append({
@@ -814,7 +671,7 @@ def process_video(video_data, temp_dir):
 
                 if not track_inputs:
                     continue
-                    
+                   
                 tracked_objects = tracker.update(
                     np.array([t["bbox"] for t in track_inputs]),
                     np.array([t["conf"] for t in track_inputs]),
@@ -827,52 +684,31 @@ def process_video(video_data, temp_dir):
                     label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
                     conf = obj['score']
                     bbox = obj['bbox']
-                    
+                   
                     if label is None:
                         continue
-                    
+                   
                     if tracker_id not in worker_id_mapping:
                         worker_id_mapping[tracker_id] = worker_counter
                         worker_counter += 1
-                    
+                   
                     worker_id = worker_id_mapping[tracker_id]
-                    
-                    # 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
+                   
+                    violation_key = (worker_id, label)
+                   
+                    if violation_key not in unique_violations:
+                        unique_violations[violation_key] = current_time
+                        violation_frames[violation_key] = frame_idx
+                        # Update violation count for this worker
+                        if worker_id not in worker_violation_count:
+                            worker_violation_count[worker_id] = 0
+                        worker_violation_count[worker_id] += 1
 
         cap.release()
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
         logger.info(f"Total unique workers detected: {len(set(worker_id_mapping.values()))}")
+        logger.info(f"Violations per worker: {worker_violation_count}")
 
         violations = []
         for (worker_id, label), detection_time in unique_violations.items():
@@ -955,12 +791,34 @@ def process_video(video_data, temp_dir):
 
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score, output_dir)
-        
+       
         record_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
-        violation_table += "|-----------|-----------|----------|------------|\n"
+        # Generate summary of workers and their violations
+        worker_summary = {}
+        for v in violations:
+            worker_id = v["worker_id"]
+            if worker_id not in worker_summary:
+                worker_summary[worker_id] = {
+                    "count": 0,
+                    "violations": set()
+                }
+            worker_summary[worker_id]["count"] += 1
+            worker_summary[worker_id]["violations"].add(v["violation"])
+
+        # Create violation table with worker summary
+        violation_table = "## Worker Safety Violation Summary\n\n"
+        violation_table += "| Worker ID | Total Violations | Violation Types |\n"
+        violation_table += "|-----------|------------------|-----------------|\n"
         
+        for worker_id, info in worker_summary.items():
+            violation_types = ", ".join([CONFIG["DISPLAY_NAMES"].get(v, v) for v in info["violations"]])
+            violation_table += f"| {worker_id} | {info['count']} | {violation_types} |\n"
+        
+        violation_table += "\n## Detailed Violation Log\n\n"
+        violation_table += "| Violation | Worker ID | Time (s) | Confidence |\n"
+        violation_table += "|-----------|-----------|----------|------------|\n"
+       
         for v in sorted(violations, key=lambda x: (x.get("worker_id", "Unknown"), x.get("timestamp", 0.0))):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
             worker_id = v.get("worker_id", "Unknown")
@@ -1006,14 +864,14 @@ def gradio_interface(video_file):
     try:
         if not video_file:
             return "No file uploaded.", "", "No file uploaded.", "", ""
-        
+       
         temp_dir = tempfile.mkdtemp(prefix="Ultralytics_")
         logger.info(f"Created temporary directory for video processing: {temp_dir}")
 
         with open(video_file, "rb") as f:
             video_data = f.read()
         logger.info(f"Read Gradio video file: {video_file}, size: {len(video_data)} bytes")
-        
+       
         if len(video_data) == 0:
             return "Uploaded video file is empty.", "", "", "", ""
 
@@ -1028,7 +886,7 @@ def gradio_interface(video_file):
 
         for status, score, snapshots_text, record_id, details_url in process_video(video_data, temp_dir):
             yield status, score, snapshots_text, record_id, details_url
-            
+           
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
@@ -1039,7 +897,7 @@ def gradio_interface(video_file):
                 logger.info(f"Cleaned up local temporary video file: {local_video_path}")
             except Exception as e:
                 logger.error(f"Failed to clean up local temporary video file {local_video_path}: {e}")
-        
+       
         if temp_dir and os.path.exists(temp_dir):
             shutil.rmtree(temp_dir, ignore_errors=True)
             logger.info(f"Cleaned up temporary directory: {temp_dir}")