Update app.py

app.py

@@ -24,7 +24,6 @@ import pytz
 import shutil
 import tempfile
 from scipy.spatial import distance
-from scipy.optimize import linear_sum_assignment # <-- Added for improved tracking
 import asyncio
 from functools import partial
 from concurrent.futures import ThreadPoolExecutor
@@ -248,7 +247,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
         # 2. Upload ContentVersion to Salesforce
         title = f"Safety_Report_{datetime.now(IST).strftime('%Y%m%d_%H%M%S')}"
         b64_pdf = base64.b64encode(pdf_bytes).decode('utf-8')
-
+        
         logger.info(f"Uploading PDF '{title}.pdf' to Salesforce...")
         cv_result = sf.ContentVersion.create({
             'Title': title,
@@ -259,7 +258,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
         if not cv_result.get('success'):
             logger.error(f"Failed to create ContentVersion: {cv_result.get('errors')}")
             return None, None
-
+        
         content_version_id = cv_result['id']
         logger.info(f"Successfully created ContentVersion with ID: {content_version_id}")
 
@@ -278,7 +277,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
             'LinkedEntityId': record_id,
             'ShareType': 'V' # V = Viewer
         } for record_id in record_ids]
-
+        
         link_success_count = 0
         for payload in link_payloads:
             try:
@@ -289,7 +288,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
                     logger.warning(f"Failed to link to {payload['LinkedEntityId']}: {link_result.get('errors')}")
             except Exception as e:
                 logger.error(f"Error creating ContentDocumentLink for {payload['LinkedEntityId']}: {e}")
-
+        
         logger.info(f"Successfully created {link_success_count}/{len(record_ids)} links.")
 
         # 5. Construct URL and Update records
@@ -299,7 +298,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
 
         update_payloads = [{'Id': record_id, 'PDF_Report_URL__c': pdf_url} for record_id in record_ids]
         update_results = sf.bulk.Safety_Violation_Log__c.update(update_payloads)
-
+        
         successful_updates = sum(1 for res in update_results if res.get('success'))
         logger.info(f"Successfully updated {successful_updates}/{len(record_ids)} records with the PDF URL.")
 
@@ -307,7 +306,7 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
         with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', prefix='report_') as temp_pdf:
             temp_pdf.write(pdf_bytes)
             temp_pdf_path = temp_pdf.name
-
+            
         logger.info(f"Salesforce report URL: {pdf_url}")
         logger.info(f"Temporary local report for download: {temp_pdf_path}")
 
@@ -316,128 +315,103 @@ def generate_and_upload_report_to_salesforce(sf, violations, record_ids):
     except Exception as e:
         logger.error(f"Error in Salesforce PDF report generation/upload: {e}", exc_info=True)
         return None, None
-
-# --- Safety Violation Detector Class (ENHANCED FOR ACCURACY) ---
+        
+# --- Safety Violation Detector Class ---
 class SafetyViolationDetector:
     def __init__(self):
-        # Detection thresholds
+        # Detection thresholds (fine-tuned for better accuracy)
         self.helmet_threshold = 0.75
         self.person_threshold = 0.60
         self.unsafe_distance = 50  # pixels
+        self.violation_cooldown = 20  # seconds
 
         # Unauthorized zones (x1, y1, x2, y2)
         self.unauthorized_zones = [
-            [100, 100, 300, 300],
-            [400, 200, 600, 400]
+            [100, 100, 300, 300],  # Example zone 1
+            [400, 200, 600, 400]   # Example zone 2
         ]
-        
-        # Tracking parameters
-        self.person_tracker = {}
-        self.next_person_id = 1
-        self.max_age = 30  # Max frames to keep a track without a new detection
-        self.min_hits = 3  # Min consecutive detections to consider a track valid
-        self.iou_threshold = 0.3 # Min Intersection over Union for a match
 
-        # Violation tracking per session
+        self.active_violations = {}
         self.session_violations = {}
+        self.person_tracker = {}
+        self.person_positions_history = {}
+        self.next_person_id = 1
+        self.max_tracking_distance = 120
+        self.max_history_length = 20  # Increased for better tracking continuity
+        self.min_iou_for_match = 0.4  # Adjusted for stricter matching
+        self.min_score_for_match = 0.5  # Adjusted for stricter matching
 
     def reset_session(self):
-        """Resets tracking and violation history for a new video or stream."""
+        self.session_violations = {}
+        self.active_violations = {}
         self.person_tracker = {}
+        self.person_positions_history = {}
         self.next_person_id = 1
-        self.session_violations = {}
-        logger.info("Tracker and violation session reset.")
+        logger.info("Session violation tracking reset for new video")
 
     def has_reported_violation(self, person_id, violation_type):
-        """Checks if a specific violation has already been reported for a person in this session."""
-        return person_id in self.session_violations and violation_type in self.session_violations[person_id]
+        if person_id not in self.session_violations:
+            return False
+        return violation_type in self.session_violations[person_id]
 
     def mark_violation_reported(self, person_id, violation_type, timestamp):
-        """Marks a violation as reported for a person."""
         if person_id not in self.session_violations:
             self.session_violations[person_id] = {}
-        self.session_violations[person_id][violation_type] = timestamp
+        self.session_violations[person_id][violation_type] = {
+            'first_detected': timestamp,
+            'count': self.session_violations[person_id].get(violation_type, {}).get('count', 0) + 1
+        }
 
-    def _iou(self, box1, box2):
-        """Calculates Intersection over Union between two bounding boxes."""
-        x1 = max(box1[0], box2[0])
-        y1 = max(box1[1], box2[1])
-        x2 = min(box1[2], box2[2])
-        y2 = min(box1[3], box2[3])
-        intersection = max(0, x2 - x1) * max(0, y2 - y1)
-        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
-        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
-        union = area1 + area2 - intersection
-        return intersection / (union + 1e-6)
-    
-    def _euclidean_distance(self, point1, point2):
-        return np.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)
+    def _get_stable_person_id(self, box, current_time):
+        center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
+        box_area = (box[2] - box[0]) * (box[3] - box[1])
 
-    def update_trackers(self, detections, current_time):
-        """
-        Updates person trackers with new detections using the Hungarian algorithm for matching.
-        This is the core of stable person tracking.
-        """
-        if not self.person_tracker:
-            for det in detections:
-                self._create_new_tracker(det, current_time)
-            return
-
-        tracked_ids = list(self.person_tracker.keys())
-        tracked_boxes = [self.person_tracker[tid]['box'] for tid in tracked_ids]
-        
-        iou_matrix = np.zeros((len(tracked_ids), len(detections)))
-        for t, trk_box in enumerate(tracked_boxes):
-            for d, det in enumerate(detections):
-                iou_matrix[t, d] = self._iou(trk_box, det['box'])
+        best_match_id = None
+        best_match_score = 0
+        min_distance = float('inf')
 
-        # Use Hungarian algorithm to find optimal matches (maximize IOU)
-        row_ind, col_ind = linear_sum_assignment(1 - iou_matrix)
-        
-        matched_indices = {r: c for r, c in zip(row_ind, col_ind) if iou_matrix[r, c] >= self.iou_threshold}
-        unmatched_trackers = set(range(len(tracked_ids))) - set(matched_indices.keys())
-        unmatched_detections = set(range(len(detections))) - set(matched_indices.values())
-
-        # Update matched trackers
-        for t_idx, d_idx in matched_indices.items():
-            track_id = tracked_ids[t_idx]
-            self.person_tracker[track_id].update({
-                'box': detections[d_idx]['box'],
-                'center': self._get_center(detections[d_idx]['box']),
-                'time_since_update': 0,
-                'hits': self.person_tracker[track_id]['hits'] + 1,
-                'hit_streak': self.person_tracker[track_id]['hit_streak'] + 1,
-                'last_seen': current_time
-            })
-
-        # Handle unmatched trackers (potential disappearance)
-        for t_idx in unmatched_trackers:
-            track_id = tracked_ids[t_idx]
-            self.person_tracker[track_id]['time_since_update'] += 1
-            self.person_tracker[track_id]['hit_streak'] = 0
-
-        # Create new trackers for unmatched detections
-        for d_idx in unmatched_detections:
-            self._create_new_tracker(detections[d_idx], current_time)
-
-        # Remove stale trackers
-        self.person_tracker = {tid: tdata for tid, tdata in self.person_tracker.items() if tdata['time_since_update'] <= self.max_age}
-
-    def _create_new_tracker(self, detection, current_time):
-        self.person_tracker[self.next_person_id] = {
-            'id': self.next_person_id,
-            'box': detection['box'],
-            'center': self._get_center(detection['box']),
-            'time_since_update': 0,
-            'hits': 1,
-            'hit_streak': 1,
-            'first_seen': current_time,
-            'last_seen': current_time,
-        }
-        self.next_person_id += 1
-        
-    def _get_center(self, box):
-        return ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
+        for person_id, history in self.person_positions_history.items():
+            if not history['positions']:
+                continue
+
+            last_position = history['positions'][-1]
+            last_box = history['boxes'][-1]
+
+            dist = np.sqrt((center[0] - last_position[0])**2 + (center[1] - last_position[1])**2)
+            iou = self._iou(box, last_box)
+
+            # Weighted scoring: more emphasis on IoU for stable tracking
+            score = (1.0 / (1.0 + dist/50)) * 0.3 + iou * 0.7
+
+            if score > best_match_score and score > self.min_score_for_match and iou > self.min_iou_for_match:
+                best_match_score = score
+                best_match_id = person_id
+                min_distance = dist
+
+        if best_match_id is not None:
+            person_id = best_match_id
+        else:
+            person_id = self.next_person_id
+            self.next_person_id += 1
+            self.person_positions_history[person_id] = {
+                'positions': [],
+                'boxes': [],
+                'first_seen': current_time,
+                'last_seen': current_time,
+                'features': []  # Store features for enhanced tracking
+            }
+
+        # Update tracking history
+        self.person_positions_history[person_id]['positions'].append(center)
+        self.person_positions_history[person_id]['boxes'].append(box)
+        self.person_positions_history[person_id]['last_seen'] = current_time
+
+        # Limit history length to prevent memory issues
+        if len(self.person_positions_history[person_id]['positions']) > self.max_history_length:
+            self.person_positions_history[person_id]['positions'].pop(0)
+            self.person_positions_history[person_id]['boxes'].pop(0)
+
+        return person_id
 
     def detect_violations(self, results, frame):
         current_time = time.time()
@@ -447,27 +421,49 @@ class SafetyViolationDetector:
         class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
         class_names = results[0].names
 
-        # Step 1: Extract all person and helmet detections from the current frame
-        person_detections = []
+        persons = []
         helmets = []
+
         for box, conf, cls_id in zip(boxes, confidences, class_ids):
             class_name = class_names[cls_id]
             if class_name == "person" and conf >= self.person_threshold:
-                person_detections.append({'box': box, 'confidence': conf})
+                person_id = self._get_stable_person_id(box, current_time)
+                persons.append({
+                    'box': box,
+                    'confidence': conf,
+                    'center': ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2),
+                    'id': person_id
+                })
             elif class_name == "hard hat" and conf >= self.helmet_threshold:
-                helmets.append({'box': box, 'confidence': conf})
+                helmets.append({
+                    'box': box,
+                    'confidence': conf,
+                    'area': (box[2] - box[0]) * (box[3] - box[1])
+                })
+
+        current_person_ids = set()
+        for person in persons:
+            person_id = person['id']
+            current_person_ids.add(person_id)
+
+            if person_id not in self.person_tracker:
+                self.person_tracker[person_id] = {
+                    'first_seen': current_time,
+                    'last_seen': current_time,
+                    'positions': [person['center']],
+                    'helmet_status': False,
+                    'violations': {},
+                    'last_violation_frame': {}
+                }
+            else:
+                self.person_tracker[person_id]['last_seen'] = current_time
+                self.person_tracker[person_id]['positions'].append(person['center'])
+                if len(self.person_tracker[person_id]['positions']) > self.max_history_length:
+                    self.person_tracker[person_id]['positions'].pop(0)
 
-        # Step 2: Update the tracker with the new person detections
-        self.update_trackers(person_detections, current_time)
-        
-        # Step 3: Get a list of *confirmed* persons from the tracker
-        persons_to_check = [
-            data for data in self.person_tracker.values()
-            if data['hit_streak'] >= self.min_hits and data['time_since_update'] == 0
-        ]
+        for person in persons:
+            person_id = person['id']
 
-        # Step 4: Check for violations for each confirmed person
-        for person in persons_to_check:
             helmet_violation = self._check_helmet_violation(person, helmets, frame, current_time)
             if helmet_violation:
                 violations.append(helmet_violation)
@@ -476,34 +472,77 @@ class SafetyViolationDetector:
             if unauthorized_violation:
                 violations.append(unauthorized_violation)
 
-        # Step 5: Check for pairwise violations like unsafe distance
-        distance_violations = self._check_distance_violations(persons_to_check, frame, current_time)
+        distance_violations = self._check_distance_violations(persons, frame, current_time)
         violations.extend(distance_violations)
 
+        self._cleanup_violations(current_time)
+        self._cleanup_inactive_persons(current_person_ids, current_time)
+
         return violations
 
     def _check_helmet_violation(self, person, helmets, frame, current_time):
         person_id = person['id']
+        person_box = person['box']
         violation_type = 'no_helmet'
 
         if self.has_reported_violation(person_id, violation_type):
             return None
 
-        person_box = person['box']
         head_region = [
-            person_box[0], person_box[1], person_box[2], person_box[1] + (person_box[3] - person_box[1]) * 0.4
+            person_box[0],
+            max(person_box[1], person_box[1] + (person_box[3] - person_box[1]) * 0.3),
+            person_box[2],
+            person_box[1] + (person_box[3] - person_box[1]) * 0.3
         ]
 
-        has_helmet = any(self._iou(helmet['box'], head_region) > 0.1 for helmet in helmets)
-        
+        has_helmet = False
+        for helmet in helmets:
+            if self._iou(helmet['box'], head_region) > 0.1:
+                has_helmet = True
+                break
+
+        self.person_tracker[person_id]['helmet_status'] = has_helmet
+
         if not has_helmet:
-            self.mark_violation_reported(person_id, violation_type, current_time)
-            self._annotate_frame(frame, person_box, person_id, "NO HELMET", (0, 0, 255))
-            logger.info(f"NEW VIOLATION: No helmet detected for person {person_id}")
-            return {
-                'type': violation_type, 'severity': 'Critical', 'person_id': person_id,
-                'timestamp': current_time
-            }
+            violation_key = f"no_helmet_{person_id}"
+
+            # Ensure persistent violation check across frames
+            if violation_key not in self.active_violations:
+                self.active_violations[violation_key] = {
+                    'type': 'no_helmet',
+                    'person_id': person_id,
+                    'first_detected': current_time,
+                    'last_detected': current_time,
+                    'count': 0,
+                    'confirmed': False
+                }
+
+            self.active_violations[violation_key]['count'] += 1
+            self.active_violations[violation_key]['last_detected'] = current_time
+
+            # Require consistent detection over multiple frames for confirmation
+            if self.active_violations[violation_key]['count'] >= 5 and not self.active_violations[violation_key]['confirmed']:
+                self.mark_violation_reported(person_id, violation_type, current_time)
+                self.active_violations[violation_key]['confirmed'] = True
+
+                if 'no_helmet' not in self.person_tracker[person_id]['violations']:
+                    self.person_tracker[person_id]['violations']['no_helmet'] = {
+                        'count': 0,
+                        'last_time': 0
+                    }
+                self.person_tracker[person_id]['violations']['no_helmet']['count'] += 1
+                self.person_tracker[person_id]['violations']['no_helmet']['last_time'] = current_time
+
+                self._annotate_frame(frame, person_box, person_id, "NO HELMET", (0, 0, 255))
+                logger.info(f"CONFIRMED VIOLATION: No helmet detected for person {person_id}")
+
+                return {
+                    'type': 'no_helmet',
+                    'severity': 'Critical',
+                    'person': person,
+                    'person_id': person_id,
+                    'timestamp': current_time
+                }
         return None
 
     def _check_unauthorized_area(self, person, frame, current_time):
@@ -513,73 +552,190 @@ class SafetyViolationDetector:
         if self.has_reported_violation(person_id, violation_type):
             return None
 
+        x1, y1, x2, y2 = person['box']
+        person_center = ((x1 + x2) / 2, (y1 + y2) / 2)
+
         for zone in self.unauthorized_zones:
             zx1, zy1, zx2, zy2 = zone
-            if (zx1 <= person['center'][0] <= zx2 and zy1 <= person['center'][1] <= zy2):
-                self.mark_violation_reported(person_id, violation_type, current_time)
-                cv2.rectangle(frame, (zx1, zy1), (zx2, zy2), (255, 0, 255), 2)
-                self._annotate_frame(frame, person['box'], person_id, "UNAUTHORIZED", (255, 0, 255))
-                logger.info(f"NEW VIOLATION: Unauthorized area detected for person {person_id}")
-                return {
-                    'type': violation_type, 'severity': 'High', 'person_id': person_id,
-                    'zone': zone, 'timestamp': current_time
-                }
+            if (zx1 <= person_center[0] <= zx2 and zy1 <= person_center[1] <= zy2):
+                violation_key = f"unauthorized_area_{person_id}_{zx1}_{zy1}"
+
+                if violation_key not in self.active_violations:
+                    self.active_violations[violation_key] = {
+                        'type': 'unauthorized_area',
+                        'person_id': person_id,
+                        'zone': zone,
+                        'first_detected': current_time,
+                        'last_detected': current_time,
+                        'count': 0,
+                        'confirmed': False
+                    }
+
+                self.active_violations[violation_key]['count'] += 1
+                self.active_violations[violation_key]['last_detected'] = current_time
+
+                if self.active_violations[violation_key]['count'] >= 5 and not self.active_violations[violation_key]['confirmed']:
+                    self.mark_violation_reported(person_id, violation_type, current_time)
+                    self.active_violations[violation_key]['confirmed'] = True
+
+                    if 'unauthorized_area' not in self.person_tracker[person_id]['violations']:
+                        self.person_tracker[person_id]['violations']['unauthorized_area'] = {
+                            'count': 0,
+                            'last_time': 0
+                        }
+                    self.person_tracker[person_id]['violations']['unauthorized_area']['count'] += 1
+                    self.person_tracker[person_id]['violations']['unauthorized_area']['last_time'] = current_time
+
+                    cv2.rectangle(frame, (zx1, zy1), (zx2, zy2), (255, 0, 255), 2)
+                    self._annotate_frame(frame, person['box'], person_id, "UNAUTHORIZED", (255, 0, 255))
+                    logger.info(f"CONFIRMED VIOLATION: Unauthorized area detected for person {person_id}")
+
+                    return {
+                        'type': 'unauthorized_area',
+                        'severity': 'High',
+                        'person': person,
+                        'person_id': person_id,
+                        'zone': zone,
+                        'timestamp': current_time
+                    }
         return None
 
     def _check_distance_violations(self, persons, frame, current_time):
         violations = []
-        violation_type = 'unsafe_distance'
+        if len(persons) < 2:
+            return violations
 
         for i in range(len(persons)):
-            for j in range(i + 1, len(persons)):
-                p1 = persons[i]
-                p2 = persons[j]
-                dist = self._euclidean_distance(p1['center'], p2['center'])
-
+            for j in range(i+1, len(persons)):
+                dist = self._euclidean_distance(persons[i]['center'], persons[j]['center'])
                 if dist < self.unsafe_distance:
-                    # Report once per pair of individuals
-                    if self.has_reported_violation(p1['id'], f"{violation_type}_{p2['id']}") or \
-                       self.has_reported_violation(p2['id'], f"{violation_type}_{p1['id']}"):
+                    person1_id = persons[i]['id']
+                    person2_id = persons[j]['id']
+                    violation_type = 'unsafe_distance'
+
+                    if (self.has_reported_violation(person1_id, violation_type) or
+                            self.has_reported_violation(person2_id, violation_type)):
                         continue
-                    
-                    self.mark_violation_reported(p1['id'], f"{violation_type}_{p2['id']}", current_time)
-                    self.mark_violation_reported(p2['id'], f"{violation_type}_{p1['id']}", current_time)
-
-                    self._annotate_distance(frame, p1['box'], p2['box'], p1['id'], p2['id'], dist)
-                    logger.info(f"NEW VIOLATION: Unsafe distance between persons {p1['id']} and {p2['id']}")
-                    violations.append({
-                        'type': violation_type, 'severity': 'Moderate', 'distance': dist,
-                        'person1_id': p1['id'], 'person2_id': p2['id'], 'timestamp': current_time
-                    })
+
+                    pair_key = f"{min(person1_id, person2_id)}_{max(person1_id, person2_id)}"
+                    violation_key = f"unsafe_distance_{pair_key}"
+
+                    if violation_key not in self.active_violations:
+                        self.active_violations[violation_key] = {
+                            'type': 'unsafe_distance',
+                            'person1_id': person1_id,
+                            'person2_id': person2_id,
+                            'first_detected': current_time,
+                            'last_detected': current_time,
+                            'count': 0,
+                            'confirmed': False
+                        }
+
+                    self.active_violations[violation_key]['count'] += 1
+                    self.active_violations[violation_key]['last_detected'] = current_time
+
+                    if self.active_violations[violation_key]['count'] >= 5 and not self.active_violations[violation_key]['confirmed']:
+                        self.mark_violation_reported(person1_id, violation_type, current_time)
+                        self.mark_violation_reported(person2_id, violation_type, current_time)
+                        self.active_violations[violation_key]['confirmed'] = True
+
+                        for pid in [person1_id, person2_id]:
+                            if 'unsafe_distance' not in self.person_tracker[pid]['violations']:
+                                self.person_tracker[pid]['violations']['unsafe_distance'] = {
+                                    'count': 0,
+                                    'last_time': 0
+                                }
+                            self.person_tracker[pid]['violations']['unsafe_distance']['count'] += 1
+                            self.person_tracker[pid]['violations']['unsafe_distance']['last_time'] = current_time
+
+                        self._annotate_distance(frame, persons[i]['box'], persons[j]['box'],
+                                                person1_id, person2_id, dist)
+                        logger.info(f"CONFIRMED VIOLATION: Unsafe distance detected between persons {person1_id} and {person2_id}")
+
+                        violations.append({
+                            'type': 'unsafe_distance',
+                            'severity': 'Moderate',
+                            'person1': persons[i],
+                            'person2': persons[j],
+                            'distance': dist,
+                            'person1_id': person1_id,
+                            'person2_id': person2_id,
+                            'timestamp': current_time
+                        })
         return violations
 
-    def _annotate_frame(self, frame, box, person_id, label_text, color):
+    def _cleanup_violations(self, current_time):
+        expired_violations = [
+            k for k, v in self.active_violations.items()
+            if current_time - v['last_detected'] > self.violation_cooldown
+        ]
+        for key in expired_violations:
+            del self.active_violations[key]
+
+    def _cleanup_inactive_persons(self, current_person_ids, current_time):
+        inactive_timeout = 60
+        expired_persons = [
+            pid for pid, data in self.person_tracker.items()
+            if pid not in current_person_ids and
+            current_time - data['last_seen'] > inactive_timeout
+        ]
+        for pid in expired_persons:
+            del self.person_tracker[pid]
+            if pid in self.person_positions_history:
+                del self.person_positions_history[pid]
+
+    def _annotate_frame(self, frame, box, person_id, violation_type, color):
         x1, y1, x2, y2 = map(int, box)
         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-        label = f"ID:{person_id:02d} {label_text}"
-        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+        label = f"ID:{person_id:03d} {violation_type}"
+        cv2.putText(frame, label, (x1, y1 - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
     def _annotate_distance(self, frame, box1, box2, id1, id2, dist):
-        for box, pid in [(box1, id1), (box2, id2)]:
-            x1, y1, x2, y2 = map(int, box)
-            color = (0, 165, 255) # Orange
-            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-            cv2.putText(frame, f"ID:{pid:02d}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
-        
-        center1 = self._get_center(box1)
-        center2 = self._get_center(box2)
-        cv2.line(frame, (int(center1[0]), int(center1[1])), (int(center2[0]), int(center2[1])), (0, 165, 255), 2)
+        x1, y1, x2, y2 = map(int, box1)
+        x3, y3, x4, y4 = map(int, box2)
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 165, 255), 2)
+        cv2.rectangle(frame, (x3, y3), (x4, y4), (0, 165, 255), 2)
+        center1 = ((x1 + x2) // 2, (y1 + y2) // 2)
+        center2 = ((x3 + x4) // 2, (y3 + y4) // 2)
+        cv2.line(frame, center1, center2, (0, 165, 255), 2)
+        mid_point = ((center1[0] + center2[0]) // 2, (center1[1] + center2[1]) // 2)
+        cv2.putText(frame, f"{dist:.1f}px", mid_point,
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 165, 255), 2)
+
+    def _iou(self, box1, box2):
+        x1 = max(box1[0], box2[0])
+        y1 = max(box1[1], box2[1])
+        x2 = min(box1[2], box2[2])
+        y2 = min(box1[3], box2[3])
+        intersection = max(0, x2 - x1) * max(0, y2 - y1)
+        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
+        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
+        return intersection / (area1 + area2 - intersection + 1e-6)
+
+    def _euclidean_distance(self, point1, point2):
+        return np.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)
 
     def get_session_summary(self):
         summary = {
-            'total_persons_tracked': self.next_person_id - 1,
+            'total_persons': len(self.session_violations),
             'violations_by_type': {},
-            'persons_with_violations': list(self.session_violations.keys())
+            'persons_with_violations': []
         }
+
         for person_id, violations in self.session_violations.items():
-            for violation_type in violations:
-                vt = violation_type.split('_')[0] # Group distance violations
-                summary['violations_by_type'][vt] = summary['violations_by_type'].get(vt, 0) + 1
+            person_info = {
+                'person_id': person_id,
+                'violations': list(violations.keys()),
+                'violation_count': len(violations)
+            }
+            summary['persons_with_violations'].append(person_info)
+
+            for violation_type in violations.keys():
+                if violation_type not in summary['violations_by_type']:
+                    summary['violations_by_type'][violation_type] = 0
+                summary['violations_by_type'][violation_type] += 1
+
         return summary
 
 # --- Frame Processing Functions ---
@@ -642,7 +798,7 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
         frame_count = 0
         processed_frames = 0
         violation_count = 0
-
+        
         # Get Salesforce connection once at the beginning
         sf = None
         if SALESFORCE_ENABLED:
@@ -650,7 +806,7 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
                 sf = get_salesforce_connection()
             except Exception as e:
                 logger.error(f"Could not connect to Salesforce at start: {e}")
-
+                
         fps = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         duration = total_frames / fps if fps > 0 else 0
@@ -666,53 +822,50 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
             if frame_count % frame_skip != 0:
                 continue
 
-            original_frame_copy = frame.copy()
             processed_frames += 1
             timestamp = datetime.now(IST).isoformat()
 
             progress_percent = min(100, (frame_count / total_frames) * 100)
             progress(progress_percent / 100, desc=f"Processing frame {frame_count}/{total_frames}")
 
-            # Note: YOLO model expects a standard image, not pre-resized
-            results = yolo_model(original_frame_copy)
+            processed_frame = preprocess_frame(frame)
+            results = yolo_model.predict(processed_frame)
 
-            # The detect_violations method now annotates the frame directly
-            violations = tracker.detect_violations(results, original_frame_copy)
+            violations = tracker.detect_violations(results, frame)
 
-            if violations:
-                frames.append(original_frame_copy)
-
-                for violation in violations:
-                    violation_count += 1
-                    snapshot_url = save_snapshot(original_frame_copy)
-                    worker_id = f"WORKER{violation.get('person_id', 'UNKNOWN')}"
-                    if violation['type'] == 'unsafe_distance':
-                        worker_id = f"WORKER{violation['person1_id']} & WORKER{violation['person2_id']}"
-                    
-                    violation_data = {
-                        'violation_type': violation['type'].replace('_', ' ').title(),
-                        'severity': violation['severity'],
-                        'timestamp': timestamp,
-                        'snapshot_url': snapshot_url,
-                        'site_id': 'SITE001',
-                        'camera_id': 'CAM001',
-                        'worker_id': worker_id,
-                        'frame_number': frame_count
-                    }
+            for violation in violations:
+                violation_count += 1
+                snapshot_url = save_snapshot(frame)
+                worker_id = f"WORKER{violation.get('person_id', 'UNKNOWN')}"
+                if violation['type'] == 'unsafe_distance':
+                    worker_id = f"WORKER{violation['person1_id']} & WORKER{violation['person2_id']}"
+                violation_data = {
+                    'violation_type': violation['type'].replace('_', ' ').title(),
+                    'severity': violation['severity'],
+                    'timestamp': timestamp,
+                    'snapshot_url': snapshot_url,
+                    'site_id': 'SITE001',
+                    'camera_id': 'CAM001',
+                    'worker_id': worker_id,
+                    'frame_number': frame_count
+                }
 
-                    if violation['type'] == 'unsafe_distance':
-                        violation_data['distance'] = f"{violation['distance']:.1f}px"
+                if violation['type'] == 'unsafe_distance':
+                    violation_data['distance'] = f"{violation['distance']:.1f}px"
+
+                current_run_violations.append(violation_data)
+                log_violation(violation_data)
+                send_alert(violation_data)
 
-                    current_run_violations.append(violation_data)
-                    log_violation(violation_data)
-                    send_alert(violation_data)
+                if sf:
+                    record_id, message = create_salesforce_violation_record(sf, violation_data)
+                    if record_id:
+                        new_sf_record_ids.append(record_id)
+                    else:
+                        logger.error(f"Salesforce push failed for violation: {message}")
 
-                    if sf:
-                        record_id, message = create_salesforce_violation_record(sf, violation_data)
-                        if record_id:
-                            new_sf_record_ids.append(record_id)
-                        else:
-                            logger.error(f"Salesforce push failed for violation: {message}")
+            if violations:
+                frames.append(frame)
 
         cap.release()
         processing_time = time.time() - start_time
@@ -727,12 +880,12 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
             logger.info(f"Generating PDF report and uploading to Salesforce for {len(new_sf_record_ids)} violations...")
             pdf_temp_path, pdf_sf_url = generate_and_upload_report_to_salesforce(sf, current_run_violations, new_sf_record_ids)
             if not pdf_temp_path:
-                   logger.error("Failed to generate and upload Salesforce report.")
+                 logger.error("Failed to generate and upload Salesforce report.")
         elif not current_run_violations:
             logger.info("No violations detected, skipping report generation.")
         else:
             logger.warning("Salesforce not configured or no violations recorded. Skipping Salesforce report upload.")
-
+            
         session_summary = tracker.get_session_summary()
         logger.info(f"Video analysis complete. Session summary: {session_summary}")
 
@@ -741,16 +894,13 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
             frame_count,
             processed_frames,
             duration,
-            len(current_run_violations), # Use length of unique violations
+            violation_count,
             processing_time,
             actual_fps,
             session_summary
         )
 
-        # Convert processed frames to RGB for Gradio display
-        display_frames = [cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in frames]
-
-        return display_frames, status_message, pdf_temp_path, format_violations_as_text(current_run_violations)
+        return frames, status_message, pdf_temp_path, format_violations_as_text(current_run_violations)
     except Exception as e:
         logger.error(f"Video processing error: {e}", exc_info=True)
         error_message = f"Video processing failed: {str(e)}"
@@ -760,7 +910,7 @@ async def process_video(video_path, frame_skip=1, progress=gr.Progress()):
         cv2.destroyAllWindows()
 
 # --- RTSP Processing ---
-async def process_rtsp_stream(rtsp_url, max_frames=300, frame_skip=5, progress=gr.Progress()):
+async def process_rtsp_stream(rtsp_url, max_frames=None, frame_skip=1, progress=gr.Progress()):
     global processing_active
     processing_active = True
 
@@ -782,39 +932,33 @@ async def process_rtsp_stream(rtsp_url, max_frames=300, frame_skip=5, progress=g
                 sf = get_salesforce_connection()
             except Exception as e:
                 logger.error(f"Could not connect to Salesforce at start: {e}")
-
-        processed_frames = []
+        
+        frames = []
         violation_count = 0
 
         progress(0, desc="Connecting to RTSP stream...")
-        frame_gen = capture_rtsp_frames(rtsp_url, max_frames)
-        
-        fc = 0
-        for frame, timestamp, _, _ in frame_gen:
-            fc += 1
+
+        for frame, timestamp, fc, _ in capture_rtsp_frames(rtsp_url, max_frames):
             if not processing_active:
                 break
 
             if fc % frame_skip != 0:
                 continue
-            
-            original_frame_copy = frame.copy()
 
-            progress(fc / max_frames, desc=f"Processing frame {fc}/{max_frames}")
+            progress_percent = min(100, (fc / (max_frames if max_frames else 100)) * 100)
+            progress(progress_percent / 100, desc=f"Processing frame {fc}")
 
-            results = yolo_model(original_frame_copy)
-            violations = tracker.detect_violations(results, original_frame_copy)
+            processed_frame = preprocess_frame(frame)
+            results = yolo_model.predict(processed_frame)
+
+            violations = tracker.detect_violations(results, frame)
 
-            if violations:
-                processed_frames.append(original_frame_copy)
-            
             for violation in violations:
                 violation_count += 1
-                snapshot_url = save_snapshot(original_frame_copy)
+                snapshot_url = save_snapshot(frame)
                 worker_id = f"WORKER{violation.get('person_id', 'UNKNOWN')}"
                 if violation['type'] == 'unsafe_distance':
                     worker_id = f"WORKER{violation['person1_id']} & WORKER{violation['person2_id']}"
-                
                 violation_data = {
                     'violation_type': violation['type'].replace('_', ' ').title(),
                     'severity': violation['severity'],
@@ -840,6 +984,9 @@ async def process_rtsp_stream(rtsp_url, max_frames=300, frame_skip=5, progress=g
                     else:
                         logger.error(f"Salesforce push failed for violation: {message}")
 
+            if violations:
+                frames.append(frame)
+
         if not processing_active:
             logger.info("Processing cancelled by user.")
 
@@ -849,20 +996,21 @@ async def process_rtsp_stream(rtsp_url, max_frames=300, frame_skip=5, progress=g
             logger.info(f"Generating PDF report and uploading to Salesforce for {len(new_sf_record_ids)} violations...")
             pdf_temp_path, pdf_sf_url = generate_and_upload_report_to_salesforce(sf, current_run_violations, new_sf_record_ids)
             if not pdf_temp_path:
-                logger.error("Failed to generate and upload Salesforce report.")
+                 logger.error("Failed to generate and upload Salesforce report.")
         elif not current_run_violations:
             logger.info("No violations detected, skipping report generation.")
         else:
             logger.warning("Salesforce not configured or no violations recorded. Skipping Salesforce report upload.")
 
+        if not processing_active:
+            return "Processing cancelled.", frames, format_violations_as_text(current_run_violations), generate_heatmap(current_run_violations), pdf_temp_path
+
         session_summary = tracker.get_session_summary()
         logger.info(f"RTSP analysis complete. Session summary: {session_summary}")
 
-        status_message = f"Processed {fc} frames with {violation_count} unique violations. Persons tracked: {session_summary['total_persons_tracked']}"
-        
-        display_frames = [cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in processed_frames]
+        status_message = f"Processed {len(frames)} frames with {violation_count} unique violations. Persons tracked: {session_summary['total_persons']}"
 
-        return status_message, display_frames, format_violations_as_text(current_run_violations), generate_heatmap(current_run_violations), pdf_temp_path
+        return status_message, frames, format_violations_as_text(current_run_violations), generate_heatmap(current_run_violations), pdf_temp_path
     except Exception as e:
         logger.error(f"RTSP processing error: {e}", exc_info=True)
         error_message = f"RTSP processing failed: {str(e)}"
@@ -884,23 +1032,23 @@ def generate_status_message(has_violations, total_frames, processed_frames, dura
 
     if session_summary:
         base_message += f"""
-👥 TOTAL PERSONS TRACKED: {session_summary['total_persons_tracked']}
+👥 UNIQUE PERSONS TRACKED: {session_summary['total_persons']}
 🔍 VIOLATION TYPES: {', '.join(session_summary['violations_by_type'].keys()) if session_summary['violations_by_type'] else 'None'}"""
 
     if has_violations:
         return f"""{base_message}
-🚨 UNIQUE VIOLATIONS LOGGED: {violation_count}
+🚨 UNIQUE VIOLATIONS: {violation_count}
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-Each violation type is logged only once per person."""
+Each violation reported only once per person"""
     else:
         return f"""{base_message}
 ✅ NO VIOLATIONS DETECTED
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-All safety protocols followed."""
+All safety protocols followed"""
 
 def save_snapshot(frame):
     try:
-        filename = f"snapshot_{int(time.time())}_{np.random.randint(100,999)}.jpg"
+        filename = f"snapshot_{int(time.time())}.jpg"
         snapshot_dir = "./snapshots"
         os.makedirs(snapshot_dir, exist_ok=True)
         snapshot_path = os.path.join(snapshot_dir, filename)
@@ -960,7 +1108,7 @@ Note: Each violation type reported only once per person
         text += f"""
 ┌─ ALERT #{i:02d} ─ {severity_emoji} {violation['violation_type'].upper()}
 │
-├─ 🕐 Time: {datetime.fromisoformat(violation['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}
+├─ 🕐 Time: {violation['timestamp']}
 ├─ ⚠️  Severity: {violation['severity']}
 ├─ 📍 Location: Site {violation['site_id']} | Camera {violation['camera_id']}
 ├─ 👷 Worker: {violation.get('worker_id', 'UNKNOWN')}
@@ -974,15 +1122,14 @@ Note: Each violation type reported only once per person
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 • Total Violations: {len(violations)}
 • Critical: {sum(1 for v in violations if v['severity'] == 'Critical')}
-• High: {sum(1 for v in violations if v['severity'] == 'High')}
 • Moderate: {sum(1 for v in violations if v['severity'] == 'Moderate')}
-• Last Alert: {datetime.fromisoformat(violations[-1]['timestamp']).strftime('%H:%M:%S') if violations else 'N/A'}
+• Last Alert: {violations[-1]['timestamp'] if violations else 'N/A'}
 
 🔄 System Status: ACTIVELY MONITORING
 ⚡ Response Time: <5 seconds
 🎯 Detection Accuracy: >90% confidence"""
     return text
-
+    
 def generate_heatmap(violations):
     if not violations:
         return None
@@ -1508,7 +1655,7 @@ with gr.Blocks(
             animateStars();
         </script>
     """)
-
+    
     # Professional Header
     gr.HTML("""
         <div class="main-header">
@@ -1516,7 +1663,7 @@ with gr.Blocks(
             <p class="header-subtitle">Enhanced Multi-Person Tracking - Each violation type detected only once per person per video</p>
         </div>
     """)
-
+    
     # Smart Media Analysis Section
     gr.HTML('<div class="section-header">📷 Smart Media Analysis</div>')
     with gr.Row():
@@ -1535,7 +1682,7 @@ with gr.Blocks(
                         elem_classes=["btn-primary"],
                         size="lg"
                     )
-
+    
     # Analysis Results Section
     gr.HTML('<div class="section-header">📊 Analysis Results & Violation Details</div>')
     with gr.Row():
@@ -1560,7 +1707,7 @@ with gr.Blocks(
                     label="📥 Download Professional Report",
                     elem_classes=["file-component"]
                 )
-
+    
     # Violation Details Section
     gr.HTML('<div class="section-header">🚨 Real-time Violation Monitoring</div>')
     with gr.Group(elem_classes=["professional-card", "alert-panel"]):
@@ -1572,7 +1719,7 @@ with gr.Blocks(
             value=format_violations_as_text(recent_violations),
             interactive=False
         )
-
+    
     # Live Stream Processing Section
     gr.HTML('<div class="section-header">📹 Live Stream Monitoring</div>')
     with gr.Row():
@@ -1615,7 +1762,7 @@ with gr.Blocks(
                     rows=2,
                     object_fit="cover"
                 )
-
+    
     # Live Violation Log Section
     gr.HTML('<div class="section-header">📊 Live Violation Analytics</div>')
     with gr.Row():
@@ -1639,8 +1786,8 @@ with gr.Blocks(
                     label="📥 Download RTSP Professional Report",
                     elem_classes=["file-component"]
                 )
-
-    # Professional Footer
+    
+     # Professional Footer
     gr.HTML(f"""
         <div class="footer-info">
             <h3>🛡️ Dynamic Safety Violation Detection using CCTV + AI</h3>