Update app.py

app.py

@@ -40,20 +40,48 @@ FFMPEG_AVAILABLE = check_ffmpeg()
 
 # ========================== # ByteTrack Implementation # ==========================
 class BYTETracker:
-    def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.4, 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
-        self.match_thresh = match_thresh  # Lowered to 0.4 to improve tracking sensitivity
+        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 = []
+        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
@@ -63,10 +91,8 @@ class BYTETracker:
             best_iou = 0
             best_track_id = None
             
+            # Try to match with active tracks
             for track_id, track_info in self.tracks.items():
-                if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
-                    continue
-                    
                 tx, ty, tw, th = track_info['bbox']
                 iou = self._calculate_iou([x, y, w, h], [tx, ty, tw, th])
                 
@@ -94,54 +120,65 @@ class BYTETracker:
                     'cls': cl
                 })
             else:
-                same_worker = False
-                for worker_id, last_pos in self.last_positions.items():
-                    if self._is_same_worker([x, y], last_pos):
-                        self.tracks[worker_id] = {
+                # 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': worker_id,
+                            'id': track_id,
                             'bbox': [x, y, w, h],
                             'score': score,
                             'cls': cl
                         })
-                        same_worker = True
+                        reidentified = True
+                        del self.recently_removed[track_id]
                         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
-        
-        current_time = time.time()
-        stale_ids = []
-        for track_id, track_info in self.tracks.items():
-            if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate:
-                stale_ids.append(track_id)
+                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
         
-        for track_id in stale_ids:
-            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]
-            
         return tracks
 
     def _calculate_iou(self, box1, box2):
@@ -159,7 +196,7 @@ class BYTETracker:
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
     
-    def _is_same_worker(self, pos1, pos2, threshold=100):  # Reduced threshold for 384x384 frames
+    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)
@@ -197,7 +234,7 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "CONFIDENCE_THRESHOLDS": {  # Lowered thresholds to improve detection
+    "CONFIDENCE_THRESHOLDS": {
         "no_helmet": 0.4,
         "no_harness": 0.25,
         "unsafe_posture": 0.25,
@@ -206,16 +243,16 @@ CONFIG = {
     },
     "MIN_VIOLATION_FRAMES": 1,
     "VIOLATION_COOLDOWN": 30.0,
-    "WORKER_TRACKING_DURATION": 30.0,
+    "WORKER_TRACKING_DURATION": 5.0,  # Reverted to 5.0 seconds
     "MAX_PROCESSING_TIME": 60,
     "FRAME_SKIP": 1,
     "BATCH_SIZE": 4,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 90,
+    "TRACK_BUFFER": 150,  # 5.0 seconds at 30 fps
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.4,
+    "MATCH_THRESH": 0.5,  # Increased to 0.5
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 100,  # Adjusted to match BYTETracker threshold
+    "MAX_WORKER_DISTANCE": 150,  # Increased to match _is_same_worker threshold
     "TARGET_RESOLUTION": (384, 384)
 }
 
@@ -553,7 +590,7 @@ def process_video(video_data, temp_dir):
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
         last_yield_time = start_time
-        worker_counter = 1  # For assigning unique worker IDs
+        worker_counter = 1
 
         while processed_frames < total_frames:
             batch_frames = []