Update app.py

app.py

@@ -24,11 +24,9 @@ import shutil
 import tenacity
 
 # ========================== # Configuration and Setup # ==========================
-# Configure logging for better debugging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# Check for FFmpeg availability to ensure video processing works
 def check_ffmpeg():
     try:
         subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
@@ -42,21 +40,20 @@ FFMPEG_AVAILABLE = check_ffmpeg()
 
 # ========================== # ByteTrack 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.5, frame_rate=30):
         self.track_thresh = track_thresh
-        self.track_buffer = track_buffer
-        self.match_thresh = match_thresh
+        self.track_buffer = track_buffer  # Increased to 90 frames (3 seconds at 30 fps)
+        self.match_thresh = match_thresh  # Lowered to 0.5 for better matching
         self.frame_rate = frame_rate
         self.next_id = 1
-        self.tracks = {}  # Store active tracks
-        self.worker_history = {}  # Track worker positions over time
-        self.last_positions = {}  # Last known positions of workers
+        self.tracks = {}
+        self.worker_history = {}
+        self.last_positions = {}
 
     def update(self, dets, scores, cls):
         tracks = []
         current_time = time.time()
         
-        # Update existing tracks with new detections
         for i, (det, score, cl) in enumerate(zip(dets, scores, cls)):
             if score < self.track_thresh:
                 continue
@@ -80,15 +77,12 @@ class BYTETracker:
                     matched = True
             
             if matched:
-                # Update existing track
                 self.tracks[best_track_id].update({
                     'bbox': [x, y, w, h],
                     'score': score,
                     'cls': cl,
                     'last_seen': current_time
                 })
-                
-                # Update position history
                 if best_track_id not in self.worker_history:
                     self.worker_history[best_track_id] = []
                 self.worker_history[best_track_id].append([x, y])
@@ -101,7 +95,6 @@ class BYTETracker:
                     'cls': cl
                 })
             else:
-                # Create new track
                 same_worker = False
                 for worker_id, last_pos in self.last_positions.items():
                     if self._is_same_worker([x, y], last_pos):
@@ -137,7 +130,6 @@ class BYTETracker:
                     })
                     self.next_id += 1
         
-        # Clean up old tracks
         current_time = time.time()
         stale_ids = []
         for track_id, track_info in self.tracks.items():
@@ -152,31 +144,23 @@ class BYTETracker:
                 del self.last_positions[track_id]
             
         return tracks
-    
+
     def _calculate_iou(self, box1, box2):
-        """Calculate IOU between two boxes"""
         x1, y1, w1, h1 = box1
         x2, y2, w2, h2 = box2
-        
-        # Calculate intersection coordinates
         x_left = max(x1 - w1/2, x2 - w2/2)
         y_top = max(y1 - h1/2, y2 - h2/2)
         x_right = min(x1 + w1/2, x2 + w2/2)
         y_bottom = min(y1 + h1/2, y2 + h2/2)
-        
         if x_right < x_left or y_bottom < y_top:
             return 0.0
-            
         intersection_area = (x_right - x_left) * (y_bottom - y_top)
-        
         box1_area = w1 * h1
         box2_area = w2 * h2
-        
         iou = intersection_area / (box1_area + box2_area - intersection_area)
         return iou
     
-    def _is_same_worker(self, pos1, pos2, threshold=100):
-        """Check if two positions likely belong to the same worker"""
+    def _is_same_worker(self, pos1, pos2, threshold=300):  # Increased to 300 pixels
         x1, y1 = pos1
         x2, y2 = pos2
         distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
@@ -194,11 +178,11 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "no_helmet": (0, 0, 255),          # Red
-        "no_harness": (0, 165, 255),       # Orange
-        "unsafe_posture": (0, 255, 0),     # Green
-        "unsafe_zone": (255, 0, 0),        # Blue
-        "improper_tool_use": (255, 255, 0) # Cyan
+        "no_helmet": (0, 0, 255),
+        "no_harness": (0, 165, 255),
+        "unsafe_posture": (0, 255, 0),
+        "unsafe_zone": (255, 0, 0),
+        "improper_tool_use": (255, 255, 0)
     },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
@@ -225,14 +209,14 @@ CONFIG = {
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,
-    "BATCH_SIZE": 8,  # Reduced batch size to lower memory usage
+    "FRAME_SKIP": 1,  # Reduced to 1 to process every frame
+    "BATCH_SIZE": 8,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
-    "TRACK_BUFFER": 30,
+    "TRACK_BUFFER": 90,  # Increased to 90
     "TRACK_THRESH": 0.3,
-    "MATCH_THRESH": 0.7,
+    "MATCH_THRESH": 0.5,  # Lowered to 0.5
     "SNAPSHOT_QUALITY": 95,
-    "MAX_WORKER_DISTANCE": 100
+    "MAX_WORKER_DISTANCE": 300  # Increased to 300
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -261,12 +245,10 @@ model = load_model()
 
 # ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
-    """Apply basic preprocessing to enhance detection"""
     frame = cv2.convertScaleAbs(frame, alpha=1.2, beta=20)
     return frame
 
 def draw_detections(frame, detections):
-    """Draw bounding boxes and labels on detection frame"""
     result_frame = frame.copy()
     
     for det in detections:
@@ -295,7 +277,6 @@ def draw_detections(frame, detections):
     return result_frame
 
 def calculate_safety_score(violations):
-    """Calculate safety score based on detected violations"""
     penalties = {
         "no_helmet": 25,
         "no_harness": 30,
@@ -322,7 +303,6 @@ def calculate_safety_score(violations):
     return score
 
 def generate_violation_pdf(violations, score, output_dir):
-    """Generate a PDF report for the detected violations"""
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
         pdf_path = os.path.join(output_dir, pdf_filename)
@@ -401,7 +381,6 @@ def generate_violation_pdf(violations, score, output_dir):
 
 @retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
-    """Connect to Salesforce with retry logic"""
     try:
         sf = Salesforce(**CONFIG["SF_CREDENTIALS"])
         logger.info("Connected to Salesforce")
@@ -412,7 +391,6 @@ def connect_to_salesforce():
         raise
 
 def upload_pdf_to_salesforce(sf, pdf_file, report_id):
-    """Upload PDF report to Salesforce"""
     try:
         if not pdf_file:
             logger.error("No PDF file provided for upload")
@@ -440,7 +418,6 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         return ""
 
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
-    """Push violation report to Salesforce"""
     try:
         sf = connect_to_salesforce()
         
@@ -502,7 +479,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     before_sleep=lambda retry_state: logger.info(f"Retrying file access (attempt {retry_state.attempt_number}/3)...")
 )
 def verify_and_open_video(video_path):
-    """Verify file existence and readability, then open with cv2.VideoCapture"""
     if not os.path.exists(video_path):
         raise FileNotFoundError(f"Temporary video file not found: {video_path}")
     
@@ -511,7 +487,7 @@ def verify_and_open_video(video_path):
         raise ValueError(f"Temporary video file is empty: {video_path}")
     
     with open(video_path, "rb") as f:
-        f.read(1)  # Test read access
+        f.read(1)
     
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
@@ -520,14 +496,12 @@ def verify_and_open_video(video_path):
     return cap
 
 def process_video(video_data, temp_dir):
-    """Process video to detect safety violations"""
     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  # Set YOLO config dir for this process
+    os.environ['YOLO_CONFIG_DIR'] = temp_dir
     
     try:
-        # Validate video data
         if not video_data:
             raise ValueError("Empty video data provided.")
         
@@ -535,14 +509,12 @@ def process_video(video_data, temp_dir):
         if len(video_data) == 0:
             raise ValueError("Video data is empty.")
 
-        # Save video to a temporary file
         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}")
 
-        # Verify the file exists and has content
         if not os.path.exists(video_path):
             raise FileNotFoundError(f"Temporary video file not found: {video_path}")
         file_size = os.path.getsize(video_path)
@@ -550,7 +522,6 @@ def process_video(video_data, temp_dir):
             raise ValueError(f"Temporary video file is empty: {video_path}")
         logger.info(f"Temporary video file size: {file_size} bytes")
 
-        # Open video with OpenCV (with retries)
         cap = verify_and_open_video(video_path)
         logger.info(f"Successfully opened video file: {video_path}")
 
@@ -571,7 +542,9 @@ def process_video(video_data, temp_dir):
             frame_rate=fps
         )
 
-        unique_violations = {}
+        # Modified: Use a single worker ID since we know there's only one worker
+        worker_id_mapping = {}  # Map tracker IDs to a single worker ID
+        unique_violations = {}  # Keyed by violation type to ensure uniqueness
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
@@ -650,19 +623,28 @@ def process_video(video_data, temp_dir):
                 )
 
                 for obj in tracked_objects:
-                    worker_id = obj['id']
+                    tracker_id = obj['id']
                     label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
                     conf = obj['score']
                     bbox = obj['bbox']
                     
                     if label is None:
                         continue
-                        
-                    if worker_id not in unique_violations:
-                        unique_violations[worker_id] = {}
                     
-                    if label not in unique_violations[worker_id]:
-                        unique_violations[worker_id][label] = current_time
+                    # Map all tracker IDs to a single worker ID (since we know there's only one worker)
+                    if not worker_id_mapping:
+                        worker_id_mapping[tracker_id] = 1  # First worker ID is 1
+                    else:
+                        # Map all subsequent tracker IDs to the same worker ID
+                        worker_id_mapping[tracker_id] = worker_id_mapping[list(worker_id_mapping.keys())[0]]
+                    
+                    worker_id = worker_id_mapping[tracker_id]
+                    
+                    # Use violation type as key to ensure uniqueness per worker
+                    violation_key = (worker_id, label)
+                    
+                    if violation_key not in unique_violations:
+                        unique_violations[violation_key] = current_time
                         
                         detection = {
                             "worker_id": worker_id,
@@ -712,19 +694,18 @@ def process_video(video_data, temp_dir):
         logger.info(f"Snapshots: {snapshots}")
 
         violations = []
-        for worker_id, worker_violations in unique_violations.items():
-            for label, detection_time in worker_violations.items():
-                confidence = next(
-                    (float(s["confidence"]) for s in snapshots if s["worker_id"] == worker_id and s["violation"] == label),
-                    0.0
-                )
-                violation = {
-                    "worker_id": worker_id,
-                    "violation": label,
-                    "timestamp": detection_time,
-                    "confidence": confidence
-                }
-                violations.append(violation)
+        for (worker_id, label), detection_time in unique_violations.items():
+            confidence = next(
+                (float(s["confidence"]) for s in snapshots if s["worker_id"] == worker_id and s["violation"] == label),
+                0.0
+            )
+            violation = {
+                "worker_id": worker_id,
+                "violation": label,
+                "timestamp": detection_time,
+                "confidence": confidence
+            }
+            violations.append(violation)
 
         logger.info(f"Violations: {violations}")
 
@@ -785,18 +766,15 @@ def process_video(video_data, temp_dir):
                 logger.error(f"Failed to clean up temporary video file {video_path}: {e}")
 
 def gradio_interface(video_file):
-    """Gradio interface for the video processing"""
     temp_dir = None
     local_video_path = None
     try:
         if not video_file:
             return "No file uploaded.", "", "No file uploaded.", "", ""
         
-        # Create a unique temporary directory for this video
         temp_dir = tempfile.mkdtemp(prefix="Ultralytics_")
         logger.info(f"Created temporary directory for video processing: {temp_dir}")
 
-        # Copy Gradio's video file to a local temporary file
         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")
@@ -804,7 +782,6 @@ def gradio_interface(video_file):
         if len(video_data) == 0:
             return "Uploaded video file is empty.", "", "", "", ""
 
-        # Save to a local temporary file to avoid Gradio file deletion
         with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
             temp_file.write(video_data)
             temp_file.flush()
@@ -821,7 +798,6 @@ def gradio_interface(video_file):
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
     finally:
-        # Clean up the local temporary video file
         if local_video_path and os.path.exists(local_video_path):
             try:
                 os.remove(local_video_path)
@@ -829,7 +805,6 @@ def gradio_interface(video_file):
             except Exception as e:
                 logger.error(f"Failed to clean up local temporary video file {local_video_path}: {e}")
         
-        # Clean up the temporary directory
         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}")