Update app.py

app.py

@@ -40,10 +40,10 @@ FFMPEG_AVAILABLE = check_ffmpeg()
 
 # ========================== # ByteTrack Implementation # ==========================
 class BYTETracker:
-    def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.5, frame_rate=30):
+    def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.4, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
-        self.match_thresh = match_thresh
+        self.match_thresh = match_thresh  # Lowered to 0.4 to improve tracking sensitivity
         self.frame_rate = frame_rate
         self.next_id = 1
         self.tracks = {}
@@ -159,7 +159,7 @@ class BYTETracker:
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
     
-    def _is_same_worker(self, pos1, pos2, threshold=300):
+    def _is_same_worker(self, pos1, pos2, threshold=100):  # Reduced threshold for 384x384 frames
         x1, y1 = pos1
         x2, y2 = pos2
         distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
@@ -197,12 +197,12 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.5,
-        "no_harness": 0.3,
-        "unsafe_posture": 0.3,
-        "unsafe_zone": 0.3,
-        "improper_tool_use": 0.3
+    "CONFIDENCE_THRESHOLDS": {  # Lowered thresholds to improve detection
+        "no_helmet": 0.4,
+        "no_harness": 0.25,
+        "unsafe_posture": 0.25,
+        "unsafe_zone": 0.25,
+        "improper_tool_use": 0.25
     },
     "MIN_VIOLATION_FRAMES": 1,
     "VIOLATION_COOLDOWN": 30.0,
@@ -213,10 +213,10 @@ CONFIG = {
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 90,
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.5,
+    "MATCH_THRESH": 0.4,
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 300,
-    "TARGET_RESOLUTION": (384, 384)  # Changed to 384x384 (divisible by 32)
+    "MAX_WORKER_DISTANCE": 100,  # Adjusted to match BYTETracker threshold
+    "TARGET_RESOLUTION": (384, 384)
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -548,11 +548,12 @@ def process_video(video_data, temp_dir):
 
         worker_id_mapping = {}
         unique_violations = {}
-        violation_frames = {}  # Store frame indices for violations
+        violation_frames = {}
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
         last_yield_time = start_time
+        worker_counter = 1  # For assigning unique worker IDs
 
         while processed_frames < total_frames:
             batch_frames = []
@@ -583,7 +584,7 @@ def process_video(video_data, temp_dir):
                 break
 
             try:
-                batch_frames_np = np.array(batch_frames)  # Shape: (batch, height, width, channels)
+                batch_frames_np = np.array(batch_frames)
                 batch_frames_tensor = torch.from_numpy(batch_frames_np).permute(0, 3, 1, 2).float() / 255.0
                 batch_frames_tensor = batch_frames_tensor.to(device)
                 if device.type == "cuda":
@@ -638,6 +639,7 @@ def process_video(video_data, temp_dir):
                     np.array([t["conf"] for t in track_inputs]),
                     np.array([t["cls"] for t in track_inputs])
                 )
+                logger.info(f"Frame {frame_idx}: Detected {len(tracked_objects)} workers")
 
                 for obj in tracked_objects:
                     tracker_id = obj['id']
@@ -648,10 +650,9 @@ def process_video(video_data, temp_dir):
                     if label is None:
                         continue
                     
-                    if not worker_id_mapping:
-                        worker_id_mapping[tracker_id] = 1
-                    else:
-                        worker_id_mapping[tracker_id] = worker_id_mapping[list(worker_id_mapping.keys())[0]]
+                    if tracker_id not in worker_id_mapping:
+                        worker_id_mapping[tracker_id] = worker_counter
+                        worker_counter += 1
                     
                     worker_id = worker_id_mapping[tracker_id]
                     
@@ -664,6 +665,7 @@ def process_video(video_data, temp_dir):
         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()))}")
 
         violations = []
         for (worker_id, label), detection_time in unique_violations.items():
@@ -671,7 +673,7 @@ def process_video(video_data, temp_dir):
                 "worker_id": worker_id,
                 "violation": label,
                 "timestamp": detection_time,
-                "confidence": 0.0,  # Will be updated after reprocessing frames
+                "confidence": 0.0,
                 "frame_idx": violation_frames[(worker_id, label)]
             })
 
@@ -680,7 +682,6 @@ 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
 
-        # Reopen video to capture snapshots for violations
         snapshots = []
         cap = cv2.VideoCapture(video_path)
         for violation in violations: