Update app.py

app.py

@@ -25,7 +25,7 @@ logger = logging.getLogger(__name__)
 
 # ========================== # Enhanced Tracker Implementation # ==========================
 class EnhancedTracker:
-    def __init__(self, track_thresh=0.4, track_buffer=60, match_thresh=0.9, frame_rate=30):
+    def __init__(self, track_thresh=0.4, track_buffer=60, match_thresh=0.8, frame_rate=30):
         self.track_thresh = track_thresh
         self.track_buffer = track_buffer
         self.match_thresh = match_thresh
@@ -34,8 +34,9 @@ class EnhancedTracker:
         self.tracks = {}
         self.worker_history = {}
         self.last_positions = {}
-        self.flagged_violations = {}
-        self.permanent_violations = {}
+        self.flagged_violations = {}  # {worker_id: {violation_type: (time, position, count)}}
+        self.permanent_violations = {}  # {worker_id: set(violation_types)}
+        self.violation_zones = {}
 
     def update(self, dets, scores, cls):
         tracks = []
@@ -96,7 +97,6 @@ class EnhancedTracker:
         for worker_id, last_pos in self.last_positions.items():
             if self._is_same_worker([x, y], last_pos):
                 self._update_existing_track(worker_id, x, y, w, h, score, cl, current_time)
-                logger.debug(f"Reused existing worker ID {worker_id} for new detection")
                 return worker_id
         
         self.tracks[self.next_id] = {
@@ -107,8 +107,7 @@ class EnhancedTracker:
         }
         self.worker_history[self.next_id] = [[x, y]]
         self.last_positions[self.next_id] = [x, y]
-        self.permanent_violations[self.next_id] = set()
-        logger.info(f"Created new worker ID {self.next_id}")
+        self.permanent_violations[self.next_id] = set()  # Initialize permanent violations set
         self.next_id += 1
         return self.next_id - 1
 
@@ -121,8 +120,6 @@ class EnhancedTracker:
             self.worker_history.pop(tid, None)
             self.last_positions.pop(tid, None)
             self.flagged_violations.pop(tid, None)
-            self.permanent_violations.pop(tid, None)
-            logger.debug(f"Cleaned up stale worker ID {tid}")
 
     def _calculate_iou(self, box1, box2):
         x1, y1, w1, h1 = box1
@@ -141,37 +138,32 @@ class EnhancedTracker:
         area2 = w2 * h2
         return intersection / (area1 + area2 - intersection)
     
-    def _is_same_worker(self, pos1, pos2, threshold=50):
+    def _is_same_worker(self, pos1, pos2, threshold=80):
         x1, y1 = pos1
         x2, y2 = pos2
-        distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
-        logger.debug(f"Distance between positions: {distance:.2f}")
-        return distance < threshold
+        return np.sqrt((x1 - x2)**2 + (y1 - y2)**2) < threshold
 
     def should_alert_violation(self, worker_id, violation_type, position, cooldown=45.0, distance_thresh=120):
         if worker_id not in self.permanent_violations:
             self.permanent_violations[worker_id] = set()
-            logger.debug(f"Initialized permanent_violations for worker {worker_id}")
-        
+            
+        # For no_helmet, only report once per worker
         if violation_type == "no_helmet":
             if "no_helmet" in self.permanent_violations[worker_id]:
-                logger.debug(f"Skipped no_helmet violation for worker {worker_id} (already reported)")
                 return False
             self.permanent_violations[worker_id].add("no_helmet")
-            logger.info(f"Recorded no_helmet violation for worker {worker_id}")
             return True
             
+        # For other violations, use cooldown logic
         if worker_id not in self.flagged_violations:
             self.flagged_violations[worker_id] = {}
             
         if violation_type in self.flagged_violations[worker_id]:
             last_time, last_pos, _ = self.flagged_violations[worker_id][violation_type]
             if time.time() - last_time < cooldown and self._is_same_worker(position, last_pos, distance_thresh):
-                logger.debug(f"Skipped {violation_type} for worker {worker_id} due to cooldown")
                 return False
                 
         self.flagged_violations[worker_id][violation_type] = (time.time(), position, 1)
-        logger.info(f"Recorded {violation_type} violation for worker {worker_id}")
         return True
 
 # ========================== # Optimized Configuration # ==========================
@@ -217,13 +209,13 @@ CONFIG = {
     "VIOLATION_COOLDOWN": 45.0,
     "WORKER_TRACKING_DURATION": 20.0,
     "MAX_PROCESSING_TIME": 45,
-    "FRAME_SKIP": 5,
-    "BATCH_SIZE": 16,
+    "FRAME_SKIP": 5,  # Increased for faster processing
+    "BATCH_SIZE": 16,  # Increased for better throughput
     "TRACK_BUFFER": 60,
     "TRACK_THRESH": 0.4,
-    "MATCH_THRESH": 0.9,
-    "SNAPSHOT_QUALITY": 85,
-    "MAX_WORKER_DISTANCE": 50,
+    "MATCH_THRESH": 0.8,
+    "SNAPSHOT_QUALITY": 85,  # Slightly reduced for faster saving
+    "MAX_WORKER_DISTANCE": 80,
     "VIOLATION_DISTANCE_THRESH": 120
 }
 
@@ -232,18 +224,10 @@ logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        model_path = CONFIG["MODEL_PATH"]
+        model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
         if not os.path.isfile(model_path):
-            logger.error(f"Custom model {model_path} not found. Please provide a trained safety model.")
-            raise FileNotFoundError(f"Custom model {model_path} not found")
-        
+            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)
-        expected_classes = set(CONFIG["VIOLATION_LABELS"].values())
-        model_classes = set(model.names.values())
-        if not expected_classes.issubset(model_classes):
-            logger.error(f"Model classes {model.names} do not contain required safety classes: {expected_classes}")
-            raise ValueError("Loaded model does not support required safety violation classes")
-        
         logger.info(f"Model loaded with classes: {model.names}")
         return model
     except Exception as e:
@@ -252,13 +236,10 @@ def load_model():
 
 model = load_model()
 
-# Global cache for no_helmet violations
-NO_HELMET_CACHE = set()
-
 # ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
     frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=25)
-    frame = cv2.GaussianBlur(frame, (5, 5), 0)
+    frame = cv2.GaussianBlur(frame, (5, 5), 0)  # Noise reduction
     return frame
 
 def draw_detections(frame, detections):
@@ -501,20 +482,6 @@ def process_video(video_file):
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
                     if label and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.3):
-                        worker_id = None
-                        if label == "no_helmet":
-                            for tid, track in tracker.tracks.items():
-                                if tid in NO_HELMET_CACHE:
-                                    continue
-                                tx, ty, tw, th = track['bbox']
-                                iou = tracker._calculate_iou(box.xywh.cpu().numpy()[0], [tx, ty, tw, th])
-                                if iou > CONFIG["MATCH_THRESH"]:
-                                    worker_id = tid
-                                    break
-                            if worker_id and worker_id in NO_HELMET_CACHE:
-                                logger.debug(f"Skipped no_helmet detection for worker {worker_id} (already in cache)")
-                                continue
-                        
                         track_inputs.append({
                             "bbox": box.xywh.cpu().numpy()[0],
                             "conf": conf,
@@ -535,19 +502,11 @@ def process_video(video_file):
                         bbox = obj['bbox']
                         position = (bbox[0], bbox[1])
                         
-                        if label == "no_helmet" and worker_id in NO_HELMET_CACHE:
-                            logger.debug(f"Skipped no_helmet violation for worker {worker_id} (already reported)")
-                            continue
-                        
                         if label and tracker.should_alert_violation(
                             worker_id, label, position,
                             CONFIG["VIOLATION_COOLDOWN"],
                             CONFIG["VIOLATION_DISTANCE_THRESH"]
                         ):
-                            if label == "no_helmet":
-                                NO_HELMET_CACHE.add(worker_id)
-                                logger.info(f"Added worker {worker_id} to no_helmet cache")
-                            
                             if worker_id not in unique_violations:
                                 unique_violations[worker_id] = {}
                             
@@ -625,9 +584,6 @@ def gradio_interface(video_file):
     if not video_file:
         return "No file uploaded", "", "", "", ""
     
-    NO_HELMET_CACHE.clear()
-    logger.info("Cleared NO_HELMET_CACHE for new video processing")
-    
     for result in process_video(video_file):
         yield result