Update app.py

app.py

@@ -38,96 +38,155 @@ def check_ffmpeg():
 
 FFMPEG_AVAILABLE = check_ffmpeg()
 
-# ========================== # Improved ByteTrack Implementation # ==========================
+# ========================== # Optimized BYTETracker Implementation # ==========================
 class BYTETracker:
-    def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
+    def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.6, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
-        self.match_thresh = match_thresh  # Increased matching threshold
+        self.match_thresh = match_thresh  # Increased for stricter matching
         self.frame_rate = frame_rate
-        self.next_id = 1
+        self.next_id = 1  # Start IDs from 1
         self.tracks = {}
+        self.worker_history = {}
         self.last_positions = {}
-        self.worker_appearance = {}  # Track worker appearance patterns
+        self.recently_removed = {}
 
     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]
+        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:
+            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 0.5 seconds
+        to_remove = []
+        for track_id, info in self.recently_removed.items():
+            if current_time - info['last_seen'] > 0.5:
+                to_remove.append(track_id)
+        for track_id in to_remove:
+            del self.recently_removed[track_id]
+
+        # Precompute bounding box centers for efficiency
+        det_centers = [(det[0], det[1]) for det in dets]
         
-        for tid in stale_ids:
-            del self.tracks[tid]
-            if tid in self.last_positions:
-                del self.last_positions[tid]
-            if tid in self.worker_appearance:
-                del self.worker_appearance[tid]
-
-        for det, score, cl in zip(dets, scores, cls):
+        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])
-                
-                # Additional check for similar appearance
-                if tid in self.worker_appearance:
-                    appearance_similarity = self._appearance_similarity([x,y,w,h], self.worker_appearance[tid])
-                    iou = (iou + appearance_similarity) / 2  # Combine spatial and appearance similarity
-                
+
                 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]
-                self.worker_appearance[best_match] = [x, y, w, h]  # Update appearance
+                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:
-                # Check if this might be an existing worker based on movement pattern
-                existing_worker = self._find_existing_worker([x, y, w, h])
-                if existing_worker is not None:
-                    tid = existing_worker
+                # Try to re-identify with recently removed tracks
+                reidentified = False
+                min_distance = float('inf')
+                best_removed_id = None
+
+                for track_id, info in self.recently_removed.items():
+                    distance = self._calculate_distance([x, y], info['last_position'])
+                    if distance < CONFIG["MAX_WORKER_DISTANCE"] and distance < min_distance:
+                        min_distance = distance
+                        best_removed_id = track_id
+                        reidentified = True
+
+                if reidentified:
+                    self.tracks[best_removed_id] = {
+                        'bbox': [x, y, w, h],
+                        'score': score,
+                        'cls': cl,
+                        'last_seen': current_time
+                    }
+                    self.worker_history[best_removed_id] = self.worker_history.get(best_removed_id, []) + [[x, y]]
+                    self.last_positions[best_removed_id] = [x, y]
+                    tracks.append({
+                        'id': best_removed_id,
+                        'bbox': [x, y, w, h],
+                        'score': score,
+                        'cls': cl
+                    })
+                    del self.recently_removed[best_removed_id]
                 else:
-                    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]
-                self.worker_appearance[tid] = [x, y, w, h]
-                tracks.append({
-                    'id': tid,
-                    'bbox': [x, y, w, h],
-                    'score': score,
-                    'cls': cl
-                })
-        
+                    # Only create new ID if no existing worker is close
+                    same_worker = False
+                    for track_id, last_pos in self.last_positions.items():
+                        if self._calculate_distance([x, y], last_pos) < CONFIG["MAX_WORKER_DISTANCE"]:
+                            self.tracks[track_id] = {
+                                'bbox': [x, y, w, h],
+                                'score': score,
+                                'cls': cl,
+                                'last_seen': current_time
+                            }
+                            self.worker_history[track_id].append([x, y])
+                            self.last_positions[track_id] = [x, y]
+                            tracks.append({
+                                'id': track_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):
@@ -142,23 +201,13 @@ 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)
-
-    def _appearance_similarity(self, box1, box2):
-        # Simple size similarity (can be enhanced with more sophisticated features)
-        _, _, w1, h1 = box1
-        _, _, w2, h2 = box2
-        size_similarity = 1 - abs(w1*h1 - w2*h2) / max(w1*h1, w2*h2)
-        return size_similarity
-
-    def _find_existing_worker(self, box):
-        x, y, w, h = box
-        for tid, last_pos in self.last_positions.items():
-            lx, ly = last_pos
-            distance = np.sqrt((x - lx)**2 + (y - ly)**2)
-            if distance < 50:  # If very close to last known position
-                return tid
-        return None
+        iou = intersection_area / (box1_area + box2_area - intersection_area)
+        return iou
+
+    def _calculate_distance(self, pos1, pos2):
+        x1, y1 = pos1
+        x2, y2 = pos2
+        return np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
 
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
@@ -203,15 +252,15 @@ CONFIG = {
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,  # Balanced processing speed and accuracy
-    "BATCH_SIZE": 4,
+    "FRAME_SKIP": 1,
+    "BATCH_SIZE": 8,  # Increased for better GPU utilization
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 30,
+    "TRACK_BUFFER": 150,
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.7,  # Increased for more strict matching
-    "SNAPSHOT_QUALITY": 90,
-    "MAX_WORKER_DISTANCE": 50,  # Reduced for more precise tracking
-    "TARGET_RESOLUTION": (384, 384)  # Balanced resolution
+    "MATCH_THRESH": 0.6,  # Increased for stricter matching
+    "SNAPSHOT_QUALITY": 95,
+    "MAX_WORKER_DISTANCE": 150,
+    "TARGET_RESOLUTION": (320, 320)  # Reduced for faster processing
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -228,7 +277,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()
@@ -240,34 +289,299 @@ def load_model():
 
 model = load_model()
 
-# [Rest of your helper functions (preprocess_frame, draw_detections, calculate_safety_score, 
-# generate_violation_pdf, connect_to_salesforce, push_report_to_salesforce, upload_pdf_to_salesforce)
-# remain exactly the same as in your original code]
+# ========================== # Helper Functions # ==========================
+def preprocess_frame(frame):
+    target_res = CONFIG["TARGET_RESOLUTION"]
+    frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_AREA)  # Faster interpolation
+    frame = cv2.convertScaleAbs(frame, alpha=1.1, beta=10)  # Reduced contrast adjustment
+    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 = 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)
+
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
+        text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
+        cv2.rectangle(result_frame, (x1, y1-text_size[1]-5), (x1+text_size[0]+5, y1), (0, 0, 0), -1)
+        cv2.putText(result_frame, display_text, (x1+3, y1-3), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+
+        conf_text = f"Conf: {confidence:.2f}"
+        cv2.putText(result_frame, conf_text, (x1+3, y2+15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
+
+    return result_frame
+
+def calculate_safety_score(violations):
+    penalties = {
+        "no_helmet": 25,
+        "no_harness": 30,
+        "unsafe_posture": 20,
+        "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
+
+def generate_violation_pdf(violations, score, output_dir):
+    try:
+        pdf_filename = f"violations_{int(time.time())}.pdf"
+        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}%")
+
+        y_position = 8.2 * inch
+        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
+
+        y_position -= 0.5 * inch
+        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)
+                    y_position = 10 * inch
+
+        c.save()
+        pdf_file.seek(0)
+
+        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
+    except Exception as e:
+        logger.error(f"Error generating PDF: {e}")
+        return "", "", None
+
+@retry(stop_max_attempt_number=3, wait_fixed=2000)
+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()
+
+        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 = {
+            "Compliance_Score__c": score,
+            "Violations_Found__c": len(violations),
+            "Violations_Details__c": violations_text,
+            "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']}")
+        except Exception as e:
+            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:
+            uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
+            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
+    except Exception as e:
+        logger.error(f"Salesforce record creation failed: {e}")
+        return "N/A", "Salesforce integration failed."
+
+@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
 
-# ========================== # Improved 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:
+        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"],
@@ -279,6 +593,7 @@ def process_video(video_data, temp_dir):
         worker_id_mapping = {}
         unique_violations = {}
         violation_frames = {}
+        worker_violation_count = {}
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
@@ -288,23 +603,23 @@ def process_video(video_data, temp_dir):
         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.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
@@ -320,7 +635,8 @@ def process_video(video_data, temp_dir):
                 if device.type == "cuda":
                     batch_frames_tensor = batch_frames_tensor.half()
 
-                results = model(batch_frames_tensor, device=device, conf=0.1, verbose=False)
+                with torch.no_grad():  # Disable gradient computation
+                    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)}")
@@ -330,27 +646,32 @@ def process_video(video_data, temp_dir):
                     torch.cuda.empty_cache()
 
             current_time = time.time()
-            if current_time - last_yield_time > 0.5:  # Update progress every 0.5s
+            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
 
+            # Early stopping if enough violations are detected
+            if len(unique_violations) >= 10 and processed_frames > total_frames * 0.5:
+                logger.info("Early stopping: Sufficient violations detected.")
+                break
+
             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
 
@@ -363,7 +684,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]),
@@ -376,32 +697,30 @@ 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
-                   
-                    # More conservative worker ID assignment
+
                     if tracker_id not in worker_id_mapping:
-                        # Check if this is likely the same worker as before
-                        if len(worker_id_mapping) > 0:  # If we already have a worker
-                            existing_worker_id = next(iter(worker_id_mapping.values()))
-                            worker_id_mapping[tracker_id] = existing_worker_id
-                        else:
-                            worker_id_mapping[tracker_id] = worker_counter
-                            worker_counter += 1
-                   
+                        worker_id_mapping[tracker_id] = worker_counter
+                        worker_counter += 1
+
                     worker_id = worker_id_mapping[tracker_id]
-                   
+
                     violation_key = (worker_id, label)
-                   
+
                     if violation_key not in unique_violations:
                         unique_violations[violation_key] = current_time
                         violation_frames[violation_key] = frame_idx
+                        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():
@@ -418,20 +737,14 @@ def process_video(video_data, temp_dir):
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
 
-        # Generate snapshots (only for the first worker)
         snapshots = []
         cap = cv2.VideoCapture(video_path)
-        worker_ids = set(v["worker_id"] for v in violations)
-        
-        # Only capture snapshots for the first worker (assuming single worker)
-        first_worker_id = min(worker_ids) if worker_ids else 1
-        worker_violations = [v for v in violations if v["worker_id"] == first_worker_id][:5]  # Limit to 5 violations
-        
-        for violation in worker_violations:
+        for violation in violations:
             frame_idx = violation["frame_idx"]
             cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
             ret, frame = cap.read()
             if not ret:
+                logger.warning(f"Failed to read frame {frame_idx} for snapshot.")
                 continue
 
             frame = preprocess_frame(frame)
@@ -446,13 +759,13 @@ def process_video(video_data, temp_dir):
             for box in boxes:
                 cls = int(box.cls)
                 conf = float(box.conf)
-                box_label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                if box_label == violation["violation"]:
+                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": box_label,
+                        "violation": label,
                         "confidence": violation["confidence"],
                         "bounding_box": bbox,
                         "timestamp": violation["timestamp"]
@@ -468,7 +781,7 @@ def process_video(video_data, temp_dir):
                         (255, 255, 255),
                         2
                     )
-                    snapshot_filename = f"violation_{box_label}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                    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,
@@ -476,55 +789,59 @@ def process_video(video_data, temp_dir):
                         [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
                     )
                     snapshots.append({
-                        "violation": box_label,
+                        "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['factor']:.2f}s")
                     break
 
         cap.release()
 
         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)
 
-        # Generate output
-        violation_table = "## Safety Violation Report\n"
-        
-        # Worker summary
         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"])
+            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"])
 
+        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_table += "\n## Detailed Violations\n"
-        violation_table += "| Violation | Time (s) | Confidence |\n"
-        violation_table += "|-----------|----------|------------|\n"
-        
-        for v in sorted(violations, key=lambda x: x.get("timestamp", 0.0)):
+            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")
             timestamp = v.get("timestamp", 0.0)
             confidence = v.get("confidence", 0.0)
-            violation_table += f"| {display_name} | {timestamp:.2f} | {confidence:.2f} |\n"
+            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} at {timestamp:.2f}s\n\n"
+            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:
@@ -535,37 +852,36 @@ def process_video(video_data, temp_dir):
             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: {str(e)}", exc_info=True)
-        yield f"Error: {str(e)}", "", "", "", ""
+        yield f"Error processing video: {str(e)}", "", "", "", ""
     finally:
         if video_path and os.path.exists(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()
 
-# [Rest of your code (gradio_interface function and interface setup) remains the same]
-
 def gradio_interface(video_file):
     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()
         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.", "", "", "", ""
 
@@ -580,7 +896,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.", "", ""
@@ -591,7 +907,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}")