Update app.py

app.py

@@ -38,80 +38,147 @@ def check_ffmpeg():
 
 FFMPEG_AVAILABLE = check_ffmpeg()
 
-# ========================== # Optimized ByteTrack Implementation # ==========================
+# ========================== # ByteTrack Implementation # ==========================
 class BYTETracker:
-    def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.5, frame_rate=30):
+    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  # Reduced buffer for faster tracking
-        self.match_thresh = match_thresh
+        self.track_buffer = track_buffer
+        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
 
     def update(self, dets, scores, cls):
         tracks = []
         current_time = time.time()
-        
+       
         # Prune stale tracks
-        stale_ids = [tid for tid, track in self.tracks.items() 
-                    if current_time - track['last_seen'] > self.track_buffer / self.frame_rate]
-        
-        for tid in stale_ids:
-            del self.tracks[tid]
-            if tid in self.last_positions:
-                del self.last_positions[tid]
-
-        for det, score, cl in zip(dets, scores, cls):
+        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] = {
+                'bbox': self.tracks[track_id]['bbox'],
+                'last_seen': current_time,
+                'last_position': self.last_positions.get(track_id, [0, 0])
+            }
+            del self.tracks[track_id]
+            if track_id in self.worker_history:
+                del self.worker_history[track_id]
+            if track_id in self.last_positions:
+                del self.last_positions[track_id]
+
+        # Clean up recently_removed tracks older than 1 second
+        to_remove = []
+        for track_id, info in self.recently_removed.items():
+            if current_time - info['last_seen'] > 1.0:
+                to_remove.append(track_id)
+        for track_id in to_remove:
+            del self.recently_removed[track_id]
+
+        for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
                 continue
-                
+               
             x, y, w, h = det
             matched = False
-            
-            # Find best match among active tracks
-            best_match = None
             best_iou = 0
-            for tid, track in self.tracks.items():
-                tx, ty, tw, th = track['bbox']
+            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_match = tid
-            
-            if best_match is not None:
-                # Update existing track
-                self.tracks[best_match].update({
+                    best_track_id = track_id
+                    matched = True
+           
+            if matched:
+                self.tracks[best_track_id].update({
                     'bbox': [x, y, w, h],
                     'score': score,
                     'cls': cl,
                     'last_seen': current_time
                 })
-                self.last_positions[best_match] = [x, y]
+                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_match,
+                    'id': best_track_id,
                     'bbox': [x, y, w, h],
                     'score': score,
                     'cls': cl
                 })
             else:
-                # Create new track
-                tid = self.next_id
-                self.next_id += 1
-                self.tracks[tid] = {
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl,
-                    'last_seen': current_time
-                }
-                self.last_positions[tid] = [x, y]
-                tracks.append({
-                    'id': tid,
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl
-                })
-        
+                # Try to re-identify with recently removed tracks
+                reidentified = False
+                for track_id, info in self.recently_removed.items():
+                    if self._is_same_worker([x, y], info['last_position'], threshold=150):  # Increased threshold
+                        self.tracks[track_id] = {
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl,
+                            'last_seen': current_time
+                        }
+                        self.worker_history[track_id] = [[x, y]]
+                        self.last_positions[track_id] = [x, y]
+                        tracks.append({
+                            'id': track_id,
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl
+                        })
+                        reidentified = True
+                        del self.recently_removed[track_id]
+                        break
+               
+                if not reidentified:
+                    # Check if it matches an existing worker by position
+                    same_worker = False
+                    for worker_id, last_pos in self.last_positions.items():
+                        if self._is_same_worker([x, y], last_pos, threshold=150):  # Increased threshold
+                            self.tracks[worker_id] = {
+                                'bbox': [x, y, w, h],
+                                'score': score,
+                                'cls': cl,
+                                'last_seen': current_time
+                            }
+                            tracks.append({
+                                'id': worker_id,
+                                'bbox': [x, y, w, h],
+                                'score': score,
+                                'cls': cl
+                            })
+                            same_worker = True
+                            break
+                   
+                    if not same_worker:
+                        self.tracks[self.next_id] = {
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl,
+                            'last_seen': current_time
+                        }
+                        self.worker_history[self.next_id] = [[x, y]]
+                        self.last_positions[self.next_id] = [x, y]
+                        tracks.append({
+                            'id': self.next_id,
+                            'bbox': [x, y, w, h],
+                            'score': score,
+                            'cls': cl
+                        })
+                        self.next_id += 1
+       
         return tracks
 
     def _calculate_iou(self, box1, box2):
@@ -126,7 +193,14 @@ class BYTETracker:
         intersection_area = (x_right - x_left) * (y_bottom - y_top)
         box1_area = w1 * h1
         box2_area = w2 * h2
-        return intersection_area / (box1_area + box2_area - intersection_area)
+        iou = intersection_area / (box1_area + box2_area - intersection_area)
+        return iou
+   
+    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
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -169,17 +243,17 @@ CONFIG = {
     },
     "MIN_VIOLATION_FRAMES": 1,
     "VIOLATION_COOLDOWN": 30.0,
-    "WORKER_TRACKING_DURATION": 5.0,
+    "WORKER_TRACKING_DURATION": 10.0,  # Reverted to 5.0 seconds
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 3,  # Process every 3rd frame for speed
-    "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": 30,  # Reduced buffer for faster tracking
+    "TRACK_BUFFER": 150,  # 5.0 seconds at 30 fps
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.5,
-    "SNAPSHOT_QUALITY": 85,  # Reduced quality for faster processing
-    "MAX_WORKER_DISTANCE": 150,
-    "TARGET_RESOLUTION": (320, 320)  # Reduced resolution for faster processing
+    "MATCH_THRESH": 0.5,  # Increased to 0.5
+    "SNAPSHOT_QUALITY": 95,
+    "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")
@@ -199,7 +273,7 @@ def load_model():
                
         model = YOLO(model_path).to(device)
         if device.type == "cuda":
-            model.model.half()  # Use half precision for faster inference
+            model.model.half()
         logger.info(f"Model classes: {model.names}")
         return model
     except Exception as e:
@@ -208,25 +282,39 @@ def load_model():
 
 model = load_model()
 
-# ========================== # Optimized Helper Functions # ==========================
+# ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
-    frame = cv2.resize(frame, CONFIG["TARGET_RESOLUTION"], interpolation=cv2.INTER_LINEAR)
+    target_res = CONFIG["TARGET_RESOLUTION"]
+    frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
+    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)
         x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
         worker_id = det.get("worker_id", "Unknown")
 
-        x1, y1 = int(x - w/2), int(y - h/2)
-        x2, y2 = int(x + w/2), int(y + h/2)
+        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(result_frame, (x1, y1), (x2, y2), color, 2)
+       
+        cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
+       
         display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
-        cv2.putText(result_frame, display_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
+        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):
@@ -237,16 +325,23 @@ 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)
-    
-    return max(0, 100 - sum(penalties.get(v, 0) for violations_set in worker_violations.values() for v in violations_set))
+   
+    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
 
 def generate_violation_pdf(violations, score, output_dir):
     try:
@@ -330,11 +425,39 @@ def connect_to_salesforce():
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
         logger.info("Connected to Salesforce")
+        sf.describe()
         return sf
     except Exception as e:
         logger.error(f"Salesforce connection failed: {e}")
         raise
 
+def upload_pdf_to_salesforce(sf, pdf_file, report_id):
+    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())}",
+            "PathOnClient": f"safety_violation_{int(time.time())}.pdf",
+            "VersionData": encoded_pdf,
+            "FirstPublishLocationId": 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
+    except Exception as e:
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
+
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
@@ -361,6 +484,8 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             "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']}")
@@ -376,8 +501,11 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             if uploaded_url:
                 try:
                     sf.Safety_Video_Report__c.update(record_id, {"PDF_Report_URL__c": uploaded_url})
+                    logger.info(f"Updated record {record_id} with PDF URL: {uploaded_url}")
                 except Exception as e:
+                    logger.error(f"Failed to update Safety_Video_Report__c: {e}")
                     sf.Account.update(record_id, {"Description": uploaded_url})
+                    logger.info(f"Updated Account record {record_id} with PDF URL")
                 pdf_url = uploaded_url
 
         return record_id, pdf_url
@@ -385,47 +513,68 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.error(f"Salesforce record creation failed: {e}")
         return "N/A", "Salesforce integration failed."
 
-def upload_pdf_to_salesforce(sf, pdf_file, report_id):
-    try:
-        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
-        content_version_data = {
-            "Title": f"Safety_Violation_Report_{int(time.time())}",
-            "PathOnClient": f"safety_violation_{int(time.time())}.pdf",
-            "VersionData": encoded_pdf,
-            "FirstPublishLocationId": 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']:
-            return ""
-           
-        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
-        return file_url
-    except Exception as e:
-        logger.error(f"Error uploading PDF to Salesforce: {e}")
-        return ""
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_fixed(1),
+    retry=tenacity.retry_if_exception_type((IOError, OSError)),
+    before_sleep=lambda retry_state: logger.info(f"Retrying file access (attempt {retry_state.attempt_number}/3)...")
+)
+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
 
-# ========================== # Fast Video Processing # ==========================
 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:
-        # Write video to temp file
+        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.")
+
         with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
             temp_file.write(video_data)
+            temp_file.flush()
             video_path = temp_file.name
+        logger.info(f"Video saved to temporary file: {video_path}")
 
-        cap = cv2.VideoCapture(video_path)
-        if not cap.isOpened():
-            raise ValueError("Could not open video file")
+        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}")
+        logger.info(f"Temporary video file size: {file_size} bytes")
+
+        cap = verify_and_open_video(video_path)
+        logger.info(f"Successfully opened video file: {video_path}")
 
-        fps = cap.get(cv2.CAP_PROP_FPS) or 30
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30
+        duration = total_frames / fps
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
+
         if total_frames <= 0:
-            raise ValueError("Video has no frames")
+            raise ValueError("Video has no frames.")
 
         tracker = BYTETracker(
             track_thresh=CONFIG["TRACK_THRESH"],
@@ -436,167 +585,326 @@ def process_video(video_data, temp_dir):
 
         worker_id_mapping = {}
         unique_violations = {}
-        worker_counter = 1
+        violation_frames = {}
+        worker_violation_count = {}  # Track violation count per worker
+        start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
-        start_time = time.time()
+        last_yield_time = start_time
+        worker_counter = 1
 
-        while True:
-            # Skip frames for faster processing
-            for _ in range(frame_skip):
-                if not cap.grab():
+        while processed_frames < total_frames:
+            batch_frames = []
+            batch_indices = []
+           
+            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.retrieve()
-            if not ret:
-                break
+               
+                ret, frame = cap.read()
+                if not ret:
+                    logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
+                    break
+               
+                frame = preprocess_frame(frame)
+               
+                for _ in range(frame_skip - 1):
+                    if not cap.grab():
+                        break
+               
+                batch_frames.append(frame)
+                batch_indices.append(frame_idx)
+                processed_frames += 1
 
-            frame = preprocess_frame(frame)
-            frame_tensor = torch.from_numpy(frame).permute(2, 0, 1).float().to(device) / 255.0
-            if device.type == "cuda":
-                frame_tensor = frame_tensor.half()
+            if not batch_frames:
+                logger.info("No more frames to process.")
+                break
 
-            results = model(frame_tensor.unsqueeze(0), verbose=False)[0]
-            
-            current_time = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000
-            boxes = results.boxes
-            detections = []
-            
-            for box in boxes:
-                cls = int(box.cls)
-                conf = float(box.conf)
-                label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+            try:
+                batch_frames_np = np.array(batch_frames)
+                batch_frames_tensor = torch.from_numpy(batch_frames_np).permute(0, 3, 1, 2).float() / 255.0
+                batch_frames_tensor = batch_frames_tensor.to(device)
+                if device.type == "cuda":
+                    batch_frames_tensor = batch_frames_tensor.half()
+
+                results = model(batch_frames_tensor, device=device, conf=0.1, verbose=False)
+            except Exception as e:
+                logger.error(f"Model inference failed: {e}")
+                raise ValueError(f"Failed to process video frames with YOLO model: {str(e)}")
+            finally:
+                batch_frames = []
+                if device.type == "cuda":
+                    torch.cuda.empty_cache()
+
+            current_time = time.time()
+            if current_time - last_yield_time > 0.1:
+                progress = (processed_frames / total_frames) * 100
+                elapsed_time = current_time - start_time
+                fps_processed = processed_frames / elapsed_time if elapsed_time > 0 else 0
+                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)):
+                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
+                       
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
+                        continue
+
+                    bbox = box.xywh.cpu().numpy()[0]
+                    track_inputs.append({
+                        "bbox": bbox,
+                        "conf": conf,
+                        "cls": cls
+                    })
+
+                if not track_inputs:
                     continue
-                
-                detections.append({
-                    "bbox": box.xywh.cpu().numpy()[0],
-                    "conf": conf,
-                    "cls": cls
-                })
-
-            if detections:
-                tracked_objs = tracker.update(
-                    np.array([d["bbox"] for d in detections]),
-                    np.array([d["conf"] for d in detections]),
-                    np.array([d["cls"] for d in detections])
+                   
+                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])
                 )
-                
-                for obj in tracked_objs:
-                    tid = obj['id']
-                    if tid not in worker_id_mapping:
-                        worker_id_mapping[tid] = worker_counter
+                logger.info(f"Frame {frame_idx}: Detected {len(tracked_objects)} workers")
+
+                for obj in tracked_objects:
+                    tracker_id = obj['id']
+                    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[tid]
-                    label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), "Unknown")
+                   
+                    worker_id = worker_id_mapping[tracker_id]
+                   
                     violation_key = (worker_id, label)
-                    
+                   
                     if violation_key not in unique_violations:
-                        unique_violations[violation_key] = {
-                            "timestamp": current_time,
-                            "frame_idx": processed_frames,
-                            "confidence": obj['score']
-                        }
-
-            processed_frames += 1
-            if time.time() - start_time > 0.5:  # Update progress every 0.5s
-                progress = (processed_frames * frame_skip / total_frames) * 100
-                yield f"Processing... {min(100, progress):.1f}%", "", "", "", ""
-                start_time = time.time()
+                        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()
-
-        # Prepare violations list
-        violations = [{
-            "worker_id": worker_id,
-            "violation": label,
-            "timestamp": info["timestamp"],
-            "confidence": info["confidence"],
-            "frame_idx": info["frame_idx"]
-        } for (worker_id, label), info in unique_violations.items()]
+        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():
+            violations.append({
+                "worker_id": worker_id,
+                "violation": label,
+                "timestamp": detection_time,
+                "confidence": 0.0,
+                "frame_idx": violation_frames[(worker_id, label)]
+            })
 
         if not violations:
-            yield "No violations detected.", "Safety Score: 100%", "", "N/A", "N/A"
+            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
 
-        # Generate snapshots (only for first violation of each type per worker)
         snapshots = []
         cap = cv2.VideoCapture(video_path)
-        snapshots_to_capture = min(10, len(violations))  # Limit to 10 snapshots
-        for violation in violations[:snapshots_to_capture]:
-            cap.set(cv2.CAP_PROP_POS_FRAMES, violation["frame_idx"])
+        for violation in violations:
+            frame_idx = violation["frame_idx"]
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
             ret, frame = cap.read()
-            if ret:
-                frame = draw_detections(frame, [violation])
-                snapshot_path = os.path.join(output_dir, 
-                    f"violation_{violation['worker_id']}_{violation['violation']}_{int(violation['timestamp']*100)}.jpg")
-                cv2.imwrite(snapshot_path, frame, [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
-                snapshots.append({
-                    "path": snapshot_path,
-                    "url": f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(snapshot_path)}",
-                    "worker_id": violation["worker_id"],
-                    "violation": violation["violation"]
-                })
+            if not ret:
+                logger.warning(f"Failed to read frame {frame_idx} for snapshot.")
+                continue
+
+            frame = preprocess_frame(frame)
+            frame_tensor = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
+            frame_tensor = frame_tensor.unsqueeze(0).to(device)
+            if device.type == "cuda":
+                frame_tensor = frame_tensor.half()
+
+            result = model(frame_tensor, device=device, conf=0.1, verbose=False)[0]
+            boxes = result.boxes
+
+            for box in boxes:
+                cls = int(box.cls)
+                conf = float(box.conf)
+                label = CONFIG["VIOLATION_LABELS"].get(cls, None)
+                if label == violation["violation"]:
+                    violation["confidence"] = round(conf, 2)
+                    bbox = box.xywh.cpu().numpy()[0]
+                    detection = {
+                        "worker_id": violation["worker_id"],
+                        "violation": label,
+                        "confidence": violation["confidence"],
+                        "bounding_box": bbox,
+                        "timestamp": violation["timestamp"]
+                    }
+                    snapshot_frame = frame.copy()
+                    snapshot_frame = draw_detections(snapshot_frame, [detection])
+                    cv2.putText(
+                        snapshot_frame,
+                        f"Time: {violation['timestamp']:.2f}s",
+                        (10, 30),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.7,
+                        (255, 255, 255),
+                        2
+                    )
+                    snapshot_filename = f"violation_{label}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                    snapshot_path = os.path.join(output_dir, snapshot_filename)
+                    cv2.imwrite(
+                        snapshot_path,
+                        snapshot_frame,
+                        [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
+                    )
+                    snapshots.append({
+                        "violation": label,
+                        "worker_id": violation["worker_id"],
+                        "timestamp": violation["timestamp"],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
+                        "confidence": violation["confidence"]
+                    })
+                    logger.info(f"Captured snapshot for {label} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
+                    break
+
         cap.release()
 
-        # Generate report
         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)
-        
-        # Worker summary
+
+        # Generate summary of workers and their violations
         worker_summary = {}
         for v in violations:
-            if v["worker_id"] not in worker_summary:
-                worker_summary[v["worker_id"]] = {"count": 0, "types": set()}
-            worker_summary[v["worker_id"]]["count"] += 1
-            worker_summary[v["worker_id"]]["types"].add(v["violation"])
-
-        # Generate output
-        violation_table = "## Worker Safety Summary\n"
-        violation_table += "| Worker ID | Violations | Types |\n|-----------|------------|-------|\n"
+            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():
-            types = ", ".join([CONFIG["DISPLAY_NAMES"].get(t, t) for t in info["types"]])
-            violation_table += f"| {worker_id} | {info['count']} | {types} |\n"
+            violation_types = ", ".join([CONFIG["DISPLAY_NAMES"].get(v, v) for v in info["violations"]])
+            violation_table += f"| {worker_id} | {info['count']} | {violation_types} |\n"
         
-        snapshots_text = "\n".join(
-            f"### Worker {s['worker_id']} - {CONFIG['DISPLAY_NAMES'].get(s['violation'], s['violation'])}\n\n" +
-            f"![Violation]({s['url']})\n"
-            for s in snapshots[:5]  # Limit to 5 snapshots in output
-        ) if snapshots else "No snapshots captured"
+        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")
+            timestamp = v.get("timestamp", 0.0)
+            confidence = v.get("confidence", 0.0)
+            violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
+
+        snapshots_text = ""
+        for s in 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,
             f"Safety Score: {score}%",
             snapshots_text,
             f"Salesforce Record ID: {record_id}",
-            final_pdf_url if final_pdf_url else pdf_url
+            final_pdf_url
         )
 
     except Exception as e:
-        logger.error(f"Error processing video: {e}")
-        yield f"Error: {str(e)}", "", "", "", ""
+        logger.error(f"Error processing video: {str(e)}", exc_info=True)
+        yield f"Error processing video: {str(e)}", "", "", "", ""
     finally:
         if video_path and os.path.exists(video_path):
-            os.remove(video_path)
+            try:
+                os.remove(video_path)
+                logger.info(f"Cleaned up temporary video file: {video_path}")
+            except Exception as e:
+                logger.error(f"Failed to clean up temporary video file {video_path}: {e}")
         if device.type == "cuda":
             torch.cuda.empty_cache()
 
-# ========================== # Gradio Interface # ==========================
 def gradio_interface(video_file):
-    temp_dir = tempfile.mkdtemp()
+    temp_dir = None
+    local_video_path = None
     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()
-        
-        for output in process_video(video_data, temp_dir):
-            yield output
+        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.", "", "", "", ""
+
+        with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
+            temp_file.write(video_data)
+            temp_file.flush()
+            local_video_path = temp_file.name
+        logger.info(f"Copied Gradio video to local temporary file: {local_video_path}")
+
+        if not FFMPEG_AVAILABLE:
+            return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", ""
+
+        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:
-        yield f"Error: {str(e)}", "", "", "", ""
+        logger.error(f"Error in Gradio interface: {e}", exc_info=True)
+        yield f"Error: {str(e)}", "", "Error in processing.", "", ""
     finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
+        if local_video_path and os.path.exists(local_video_path):
+            try:
+                os.remove(local_video_path)
+                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}")
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
 
+# ========================== # Gradio Interface # ==========================
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Video(label="Upload Site Video"),
@@ -608,10 +916,10 @@ interface = gr.Interface(
         gr.Textbox(label="Violation Details URL")
     ],
     title="Worksite Safety Violation Analyzer",
-    description="Fast safety violation detection for worksite videos. Upload a video to analyze worker safety compliance.",
+    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"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Fast Safety Analyzer App...")
+    logger.info("Launching Enhanced Safety Analyzer App...")
     interface.launch()