Update app.py

app.py

@@ -20,18 +20,14 @@ import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
 
-# ==========================
-# Configuration and Setup
-# ==========================
+# ========================== # Configuration and Setup # ==========================
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
 os.makedirs('/tmp/Ultralytics', exist_ok=True)
 
-logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# ==========================
-# ByteTrack Implementation
-# ==========================
+# ========================== # ByteTrack Implementation # ==========================
 class BYTETracker:
     def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
         self.track_thresh = track_thresh
@@ -39,27 +35,102 @@ class BYTETracker:
         self.match_thresh = match_thresh
         self.frame_rate = frame_rate
         self.next_id = 1
-        
+        self.tracks = {}  # Store active tracks
+
     def update(self, dets, scores, cls):
         tracks = []
+        
+        # Update existing tracks with new detections
         for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
                 logger.debug(f"Skipping detection with score {score} below threshold {self.track_thresh}")
                 continue
                 
             x, y, w, h = det
-            tracks.append({
-                'id': self.next_id,
-                'bbox': [x, y, w, h],
-                'score': score,
-                'cls': cl
-            })
-            self.next_id += 1
+            
+            # Try to match with existing tracks
+            matched = False
+            for track_id, track_info in self.tracks.items():
+                # Simple IOU-based matching
+                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 track_info['cls'] == cl:
+                    # Update existing track
+                    self.tracks[track_id] = {
+                        'bbox': [x, y, w, h],
+                        'score': score,
+                        'cls': cl,
+                        'last_seen': time.time()
+                    }
+                    tracks.append({
+                        'id': track_id,
+                        'bbox': [x, y, w, h],
+                        'score': score,
+                        'cls': cl
+                    })
+                    matched = True
+                    break
+            
+            if not matched:
+                # Create new track
+                self.tracks[self.next_id] = {
+                    'bbox': [x, y, w, h],
+                    'score': score,
+                    'cls': cl,
+                    'last_seen': time.time()
+                }
+                tracks.append({
+                    'id': self.next_id,
+                    'bbox': [x, y, w, h],
+                    'score': score,
+                    'cls': cl
+                })
+                self.next_id += 1
+        
+        # Remove stale tracks
+        current_time = time.time()
+        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:
+            del self.tracks[track_id]
+            
         return tracks
+    
+    def _calculate_iou(self, box1, box2):
+        """Calculate IOU between two boxes in format [x, y, w, h]"""
+        x1, y1, w1, h1 = box1
+        x2, y2, w2, h2 = box2
+        
+        # Convert to xmin, ymin, xmax, ymax
+        xmin1, ymin1 = x1 - w1/2, y1 - h1/2
+        xmax1, ymax1 = x1 + w1/2, y1 + h1/2
+        xmin2, ymin2 = x2 - w2/2, y2 - h2/2
+        xmax2, ymax2 = x2 + w2/2, y2 + h2/2
+        
+        # Calculate area of intersection
+        x_left = max(xmin1, xmin2)
+        y_top = max(ymin1, ymin2)
+        x_right = min(xmax1, xmax2)
+        y_bottom = min(ymax1, ymax2)
+        
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
+            
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+        
+        # Calculate area of both boxes
+        box1_area = w1 * h1
+        box2_area = w2 * h2
+        
+        # Calculate IOU
+        iou = intersection_area / (box1_area + box2_area - intersection_area)
+        return iou
 
-# ==========================
-# Optimized Configuration
-# ==========================
+# ========================== # Optimized Configuration # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
@@ -72,11 +143,11 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "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)
+        "no_helmet": (0, 0, 255),          # Red in BGR
+        "no_harness": (0, 165, 255),       # Orange in BGR
+        "unsafe_posture": (0, 255, 0),     # Green in BGR
+        "unsafe_zone": (255, 0, 0),        # Blue in BGR
+        "improper_tool_use": (255, 255, 0) # Cyan in BGR
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -93,21 +164,23 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Sadio2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.5,  # Lowered from 0.75
-        "no_harness": 0.3,  # Lowered from 0.4
-        "unsafe_posture": 0.3,  # Lowered from 0.4
-        "unsafe_zone": 0.3,  # Lowered from 0.4
-        "improper_tool_use": 0.3  # Lowered from 0.4
+        "no_helmet": 0.5,
+        "no_harness": 0.3,
+        "unsafe_posture": 0.3,
+        "unsafe_zone": 0.3,
+        "improper_tool_use": 0.3
     },
-    "MIN_VIOLATION_FRAMES": 1,  # Lowered from 3
+    "MIN_VIOLATION_FRAMES": 1,
+    "VIOLATION_COOLDOWN": 5.0,        # Time in seconds before same violation type can be detected again for the same worker
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 1,
-    "BATCH_SIZE": 16,  # Reduced from 32 to prevent memory issues
+    "BATCH_SIZE": 16,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 30,
-    "TRACK_THRESH": 0.3,  # Lowered from 0.4
-    "MATCH_THRESH": 0.7  # Lowered from 0.8
+    "TRACK_THRESH": 0.3,
+    "MATCH_THRESH": 0.7,
+    "SNAPSHOT_QUALITY": 90           # JPEG quality for snapshots (0-100)
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -124,6 +197,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)
         logger.info(f"Model classes: {model.names}")
         return model
@@ -133,33 +207,47 @@ def load_model():
 
 model = load_model()
 
-# ==========================
-# Helper Functions
-# ==========================
+# ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
     """Apply basic preprocessing to enhance detection"""
     frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)  # Increase contrast
     return frame
 
 def draw_detections(frame, detections):
+    """Draw bounding boxes and labels on detection frame with improved visibility"""
+    result_frame = frame.copy()
+    
     for det in detections:
         label = det.get("violation", "Unknown")
         confidence = det.get("confidence", 0.0)
         x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
-        
+
         x1 = int(x - w/2)
         y1 = int(y - h/2)
         x2 = int(x + w/2)
         y2 = int(y + h/2)
         
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
-        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
         
+        # Draw thicker rectangle with border for better visibility
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+        
+        # Add a black background behind text for better readability
         display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)}: {confidence:.2f}"
-        cv2.putText(frame, display_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-    return frame
+        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)
+        
+        # Draw worker ID
+        worker_id = det.get("worker_id", "Unknown")
+        worker_text = f"Worker: {worker_id}"
+        cv2.putText(result_frame, worker_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
+        cv2.putText(result_frame, worker_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+        
+    return result_frame
 
 def calculate_safety_score(violations):
+    """Calculate safety score based on detected violations"""
     penalties = {
         "no_helmet": 25,
         "no_harness": 30,
@@ -167,41 +255,78 @@ def calculate_safety_score(violations):
         "unsafe_zone": 35,
         "improper_tool_use": 25
     }
-    total_penalty = sum(penalties.get(v.get("violation", "Unknown"), 0) for v in violations)
+    
+    # Count unique violation types
+    unique_violations = set()
+    for v in violations:
+        unique_violations.add(v.get("violation", "Unknown"))
+    
+    # Calculate penalty based on unique violation types
+    total_penalty = sum(penalties.get(v, 0) for v in unique_violations)
     score = 100 - total_penalty
     return max(score, 0)
 
 def generate_violation_pdf(violations, score):
+    """Generate a PDF report for the detected violations"""
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
         pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         pdf_file = BytesIO()
         c = canvas.Canvas(pdf_file, pagesize=letter)
-        c.setFont("Helvetica", 12)
+        c.setFont("Helvetica-Bold", 16)
         c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
-        c.setFont("Helvetica", 10)
+        
+        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}%")
 
-        y_position = 9.5 * inch
-        report_data = {
-            "Compliance Score": f"{score}%",
-            "Violations Found": len(violations),
-            "Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
+        y_position = 8.2 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Summary:")
+        y_position -= 0.3 * inch
+        
+        c.setFont("Helvetica", 10)
+        summary_data = {
+            "Total Violations Found": len(violations),
+            "Unique Workers with Violations": len(set(v.get("worker_id", "Unknown") for v in violations)),
+            "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
-        for key, value in report_data.items():
+        
+        for key, value in summary_data.items():
             c.drawString(1 * inch, y_position, f"{key}: {value}")
-            y_position -= 0.3 * inch
+            y_position -= 0.25 * inch
 
-        y_position -= 0.3 * inch
-        c.drawString(1 * inch, y_position, "Violation Details:")
-        y_position -= 0.3 * inch
         if not violations:
+            y_position -= 0.3 * inch
             c.drawString(1 * inch, y_position, "No violations detected.")
         else:
-            for v in violations:
+            y_position -= 0.5 * inch
+            c.setFont("Helvetica-Bold", 12)
+            c.drawString(1 * inch, y_position, "Violation Details:")
+            y_position -= 0.3 * inch
+            
+            c.setFont("Helvetica", 10)
+            # Sort violations by worker ID and type for better organization
+            sorted_violations = sorted(violations, key=lambda v: (v.get("worker_id", "Unknown"), v.get("violation", "Unknown")))
+            
+            for v in sorted_violations:
+                worker_id = v.get("worker_id", "Unknown")
                 display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-                text = f"{display_name} from {v.get('start_timestamp', 0.0):.2f}s to {v.get('end_timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f}, Worker ID: {v.get('worker_id', 'N/A')})"
+                start_time = v.get('start_timestamp', 0.0)
+                end_time = v.get('end_timestamp', 0.0)
+                confidence = v.get('confidence', 0.0)
+                
+                text = f"Worker ID: {worker_id} - {display_name}"
                 c.drawString(1 * inch, y_position, text)
+                y_position -= 0.2 * inch
+                
+                details = f"  Time: {start_time:.2f}s to {end_time:.2f}s (Confidence: {confidence:.2f})"
+                c.drawString(1.2 * inch, y_position, details)
                 y_position -= 0.3 * inch
+                
                 if y_position < 1 * inch:
                     c.showPage()
                     c.setFont("Helvetica", 10)
@@ -222,6 +347,7 @@ def generate_violation_pdf(violations, score):
 
 @retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
+    """Connect to Salesforce with retry logic"""
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
         logger.info("Connected to Salesforce")
@@ -232,10 +358,12 @@ def connect_to_salesforce():
         raise
 
 def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    """Upload PDF report to Salesforce"""
     try:
         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())}",
@@ -245,9 +373,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
@@ -256,12 +386,24 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         return ""
 
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
+    """Push violation report to Salesforce"""
     try:
         sf = connect_to_salesforce()
-        violations_text = "\n".join(
-            f"{CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')} from {v.get('start_timestamp', 0.0):.2f}s to {v.get('end_timestamp', 0.0):.2f}s (Confidence: {v.get('confidence', 0.0):.2f}, Worker ID: {v.get('worker_id', 'N/A')})"
-            for v in violations
-        ) or "No violations detected."
+        
+        # Format violations for Salesforce
+        violations_text = ""
+        for v in violations:
+            display_name = CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')
+            worker_id = v.get('worker_id', 'N/A')
+            start_time = v.get('start_timestamp', 0.0)
+            end_time = v.get('end_timestamp', 0.0)
+            confidence = v.get('confidence', 0.0)
+            
+            violations_text += f"Worker {worker_id}: {display_name} ({start_time:.2f}s-{end_time:.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 = {
@@ -271,7 +413,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']}")
@@ -279,6 +423,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:
@@ -299,6 +444,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         return None, ""
 
 def process_video(video_data):
+    """Process video to detect safety violations"""
     try:
         os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
         logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
@@ -327,12 +473,14 @@ def process_video(video_data):
             frame_rate=fps
         )
 
-        violation_tracker = {}
+        # Track unique violations by worker ID
+        unique_violations = {}  # {worker_id: {violation_type: {first_detection, last_detection, best_confidence, best_frame, cooldown}}}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
+        processed_frames = 0
 
-        while True:
+        while processed_frames < total_frames:
             batch_frames = []
             batch_indices = []
             
@@ -347,6 +495,7 @@ def process_video(video_data):
                 
                 frame = preprocess_frame(frame)
                 
+                # Skip frames if needed
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
                         break
@@ -357,11 +506,14 @@ def process_video(video_data):
             if not batch_frames:
                 break
 
+            # Process batch with YOLO model
             results = model(batch_frames, device=device, conf=0.1, verbose=False)
             
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
+                processed_frames += 1
                 current_time = frame_idx / fps
                 
+                # Update progress every second
                 if time.time() - start_time > 1.0:
                     progress = (frame_idx / total_frames) * 100
                     yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{total_frames})", "", "", "", ""
@@ -369,6 +521,7 @@ def process_video(video_data):
 
                 boxes = result.boxes
                 track_inputs = []
+                
                 for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
@@ -377,6 +530,7 @@ def process_video(video_data):
                     if label is None:
                         logger.debug(f"Unknown class ID {cls} detected, skipping")
                         continue
+                        
                     if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
                         logger.debug(f"Detection for {label} with confidence {conf} below threshold {CONFIG['CONFIDENCE_THRESHOLDS'].get(label, 0.25)}")
                         continue
@@ -388,89 +542,178 @@ def process_video(video_data):
                         "cls": cls
                     })
 
+                # Skip tracking if no detections
+                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]),
                     np.array([t["cls"] for t in track_inputs])
                 )
+                
                 logger.debug(f"Frame {frame_idx}: {len(tracked_objects)} objects tracked")
 
-                for obj, track_input in zip(tracked_objects, track_inputs):
+                # Process tracked objects for violations
+                for obj in tracked_objects:
                     worker_id = obj['id']
                     label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
-                    bbox = track_input["bbox"]
-                    conf = track_input["conf"]
+                    conf = obj['score']
+                    bbox = obj['bbox']
                     
-                    detection = {
-                        "frame": frame_idx,
-                        "violation": label,
-                        "confidence": round(conf, 2),
-                        "bounding_box": [round(x, 2) for x in bbox],
-                        "timestamp": current_time,
-                        "worker_id": worker_id
-                    }
-                    logger.debug(f"Detection: {detection}")
-
-                    if worker_id not in violation_tracker:
-                        violation_tracker[worker_id] = {}
-                    if label not in violation_tracker[worker_id]:
-                        violation_tracker[worker_id][label] = []
-                    violation_tracker[worker_id][label].append(detection)
+                    if label is None:
+                        continue
+                        
+                    # Initialize worker if not seen before
+                    if worker_id not in unique_violations:
+                        unique_violations[worker_id] = {}
+                        
+                    # Check if this is a new violation type for this worker or if cooldown has passed
+                    is_new_violation = False
+                    if label not in unique_violations[worker_id]:
+                        # New violation type for this worker
+                        unique_violations[worker_id][label] = {
+                            'first_detection': current_time,
+                            'last_detection': current_time,
+                            'best_confidence': conf,
+                            'best_frame': frame_idx,
+                            'best_bbox': bbox,
+                            'cooldown': current_time + CONFIG["VIOLATION_COOLDOWN"]
+                        }
+                        is_new_violation = True
+                    elif current_time > unique_violations[worker_id][label]['cooldown']:
+                        # Cooldown period has passed, treat as a new violation
+                        unique_violations[worker_id][label] = {
+                            'first_detection': current_time,
+                            'last_detection': current_time,
+                            'best_confidence': conf,
+                            'best_frame': frame_idx,
+                            'best_bbox': bbox,
+                            'cooldown': current_time + CONFIG["VIOLATION_COOLDOWN"]
+                        }
+                        is_new_violation = True
+                    else:
+                        # Update existing violation
+                        violation_info = unique_violations[worker_id][label]
+                        violation_info['last_detection'] = current_time
+                        
+                        # Update if this is a better detection (higher confidence)
+                        if conf > violation_info['best_confidence']:
+                            violation_info['best_confidence'] = conf
+                            violation_info['best_frame'] = frame_idx
+                            violation_info['best_bbox'] = bbox
+                    
+                    # If this is a new violation, capture a snapshot
+                    if is_new_violation:
+                        # Create a detection object for the snapshot
+                        detection = {
+                            "frame": frame_idx,
+                            "violation": label,
+                            "confidence": round(conf, 2),
+                            "bounding_box": bbox,
+                            "timestamp": current_time,
+                            "worker_id": worker_id
+                        }
+                        
+                        # Take a snapshot for the new violation
+                        snapshot_frame = batch_frames[i].copy()
+                        snapshot_frame = draw_detections(snapshot_frame, [detection])
+                        
+                        # Add timestamp to the image
+                        cv2.putText(
+                            snapshot_frame, 
+                            f"Time: {current_time:.2f}s", 
+                            (10, 30), 
+                            cv2.FONT_HERSHEY_SIMPLEX, 
+                            0.7, 
+                            (255, 255, 255), 
+                            2
+                        )
+                        
+                        # Save snapshot with high quality
+                        snapshot_filename = f"{label}_worker{worker_id}_{int(current_time)}_{frame_idx}.jpg"
+                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                        
+                        # Use higher quality for JPEG to ensure better visibility
+                        cv2.imwrite(
+                            snapshot_path, 
+                            snapshot_frame, 
+                            [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
+                        )
+                        
+                        snapshots.append({
+                            "violation": label,
+                            "worker_id": worker_id,
+                            "frame": frame_idx,
+                            "timestamp": current_time,
+                            "snapshot_path": snapshot_path,
+                            "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                        })
+                        
+                        logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at frame {frame_idx}")
 
         cap.release()
         if os.path.exists(video_path):
             os.remove(video_path)
+            
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
+        # Convert tracked violations to final violation list
         violations = []
-        for worker_id, worker_violations in violation_tracker.items():
-            for label, detections in worker_violations.items():
-                if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
-                    best_detection = max(detections, key=lambda x: x["confidence"])
-                    best_detection["start_timestamp"] = min(d["timestamp"] for d in detections)
-                    best_detection["end_timestamp"] = max(d["timestamp"] for d in detections)
-                    violations.append(best_detection)
-
-                    cap = cv2.VideoCapture(video_path)
-                    if not cap.isOpened():
-                        logger.warning(f"Could not reopen video for snapshot at frame {best_detection['frame']}")
-                        continue
-                    cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
-                    ret, snapshot_frame = cap.read()
-                    if ret:
-                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
-                        snapshot_filename = f"{label}_{best_detection['frame']}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        cv2.imwrite(snapshot_path, snapshot_frame)
-                        snapshots.append({
-                            "violation": label,
-                            "frame": best_detection["frame"],
-                            "snapshot_path": snapshot_path,
-                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                        })
-                    cap.release()
+        for worker_id, worker_violations in unique_violations.items():
+            for label, violation_info in worker_violations.items():
+                violation = {
+                    "worker_id": worker_id,
+                    "violation": label,
+                    "confidence": violation_info['best_confidence'],
+                    "start_timestamp": violation_info['first_detection'],
+                    "end_timestamp": violation_info['last_detection'],
+                    "frame": violation_info['best_frame'],
+                    "bounding_box": violation_info['best_bbox']
+                }
+                violations.append(violation)
 
         if not violations:
             logger.info("No violations detected after processing")
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
 
+        # Calculate safety score
         score = calculate_safety_score(violations)
+        
+        # Generate PDF report
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        
+        # Push report to Salesforce
         report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        violation_table = "| Violation              | Time Range (s) | Confidence | Worker ID |\n"
-        violation_table += "|------------------------|----------------|------------|-----------|\n"
-        for v in sorted(violations, key=lambda x: x["start_timestamp"]):
+        # Format violations table for display
+        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("start_timestamp", 0.0))):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
-            row = f"| {display_name:<22} | {v.get('start_timestamp', 0.0):.2f}-{v.get('end_timestamp', 0.0):.2f} | {v.get('confidence', 0.0):.2f} | {v.get('worker_id', 'N/A')} |\n"
+            worker_id = v.get("worker_id", "Unknown")
+            start_time = v.get('start_timestamp', 0.0)
+            end_time = v.get('end_timestamp', 0.0)
+            confidence = v.get('confidence', 0.0)
+            
+            row = f"| {display_name} | {worker_id} | {start_time:.2f}-{end_time:.2f} | {confidence:.2f} |\n"
             violation_table += row
 
-        snapshots_text = "\n".join(
-            f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
-            for s in snapshots
-        ) if snapshots else "No snapshots captured."
+        # Format snapshots for display
+        snapshots_text = ""
+        for i, s in enumerate(snapshots):
+            display_name = CONFIG["DISPLAY_NAMES"].get(s['violation'], "Unknown")
+            worker_id = s.get("worker_id", "Unknown")
+            timestamp = s.get("timestamp", 0.0)
+            
+            snapshots_text += f"### {display_name} - Worker {worker_id} at {timestamp:.2f}s\n\n"
+            snapshots_text += f"![Violation]({s['snapshot_url']})\n\n"
+
+        if not snapshots_text:
+            snapshots_text = "No snapshots captured."
 
         yield (
             violation_table,
@@ -487,21 +730,22 @@ def process_video(video_data):
         yield f"Error processing video: {e}", "", "", "", ""
 
 def gradio_interface(video_file):
+    """Gradio interface for the video processing"""
     if not video_file:
         return "No file uploaded.", "", "No file uploaded.", "", ""
+        
     try:
         with open(video_file, "rb") as f:
             video_data = f.read()
 
         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:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
 
-# ==========================
-# Gradio Interface
-# ==========================
+# ========================== # Gradio Interface # ==========================
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
@@ -513,7 +757,7 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Non-violations are ignored.",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use). Each unique violation is detected only once per worker.",
     allow_flagging="never"
 )