Update app.py

app.py

@@ -7,7 +7,8 @@ import cv2
 import gradio as gr
 import torch
 import numpy as np
-from ultralytics import YOLO
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from PIL import Image
 import time
 from simple_salesforce import Salesforce
 from reportlab.lib.pagesizes import letter
@@ -22,6 +23,7 @@ from functools import partial
 import tempfile
 import shutil
 import tenacity
+from scipy.spatial import distance
 
 # ========================== # Configuration and Setup # ==========================
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -38,7 +40,7 @@ def check_ffmpeg():
 
 FFMPEG_AVAILABLE = check_ffmpeg()
 
-# ========================== # ByteTrack Implementation # ==========================
+# ========================== # BYTETracker Implementation # ==========================
 class BYTETracker:
     def __init__(self, track_thresh=0.3, track_buffer=90, match_thresh=0.5, frame_rate=30):
         self.track_thresh = track_thresh
@@ -49,37 +51,30 @@ class BYTETracker:
         self.tracks = {}
         self.worker_history = {}
         self.last_positions = {}
-        self.recently_removed = {}  # Store recently removed tracks for re-identification
-        self.helmet_status = {}  # Track helmet status for each worker
+        self.recently_removed = {}
+        self.helmet_status = {}
+        self.harness_status = {}
 
     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)
-        
+        stale_ids = [track_id for track_id, track_info in self.tracks.items()
+                     if current_time - track_info['last_seen'] > self.track_buffer / self.frame_rate]
         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]
+            self.worker_history.pop(track_id, None)
+            self.last_positions.pop(track_id, None)
 
         # 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)
+        to_remove = [track_id for track_id, info in self.recently_removed.items()
+                     if current_time - info['last_seen'] > 1.0]
         for track_id in to_remove:
             del self.recently_removed[track_id]
 
@@ -92,7 +87,6 @@ class BYTETracker:
             best_iou = 0
             best_track_id = None
             
-            # Try to match with active tracks
             for track_id, track_info in self.tracks.items():
                 tx, ty, tw, th = track_info['bbox']
                 iou = self._calculate_iou([x, y, w, h], [tx, ty, tw, th])
@@ -110,15 +104,12 @@ class BYTETracker:
                     'last_seen': current_time
                 })
                 
-                # Update helmet status if this is a helmet detection
-                if cl == 0:  # Helmet violation class
-                    # Higher confidence for helmet violations
-                    if score > 0.45:  # Increased threshold for helmet violations
-                        self.helmet_status[best_track_id] = True
+                if cl == "no_helmet" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                    self.helmet_status[best_track_id] = True
+                elif cl == "no_harness" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_harness"]:
+                    self.harness_status[best_track_id] = True
                 
-                if best_track_id not in self.worker_history:
-                    self.worker_history[best_track_id] = []
-                self.worker_history[best_track_id].append([x, y])
+                self.worker_history[best_track_id] = self.worker_history.get(best_track_id, []) + [[x, y]]
                 self.last_positions[best_track_id] = [x, y]
                 
                 tracks.append({
@@ -128,10 +119,9 @@ class BYTETracker:
                     'cls': cl
                 })
             else:
-                # 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
+                for track_id, info in list(self.recently_removed.items()):
+                    if self._is_same_worker([x, y], info['last_position'], threshold=CONFIG["MAX_WORKER_DISTANCE"]):
                         self.tracks[track_id] = {
                             'bbox': [x, y, w, h],
                             'score': score,
@@ -141,11 +131,10 @@ class BYTETracker:
                         self.worker_history[track_id] = [[x, y]]
                         self.last_positions[track_id] = [x, y]
                         
-                        # Update helmet status if this is a helmet detection
-                        if cl == 0:  # Helmet violation class
-                            # Higher confidence for helmet violations
-                            if score > 0.45:  # Increased threshold for helmet violations
-                                self.helmet_status[track_id] = True
+                        if cl == "no_helmet" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                            self.helmet_status[track_id] = True
+                        elif cl == "no_harness" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_harness"]:
+                            self.harness_status[track_id] = True
                         
                         tracks.append({
                             'id': track_id,
@@ -158,10 +147,9 @@ class BYTETracker:
                         break
                 
                 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
+                        if self._is_same_worker([x, y], last_pos, threshold=CONFIG["MAX_WORKER_DISTANCE"]):
                             self.tracks[worker_id] = {
                                 'bbox': [x, y, w, h],
                                 'score': score,
@@ -169,11 +157,10 @@ class BYTETracker:
                                 'last_seen': current_time
                             }
                             
-                            # Update helmet status if this is a helmet detection
-                            if cl == 0:  # Helmet violation class
-                                # Higher confidence for helmet violations
-                                if score > 0.45:  # Increased threshold for helmet violations
-                                    self.helmet_status[worker_id] = True
+                            if cl == "no_helmet" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                                self.helmet_status[worker_id] = True
+                            elif cl == "no_harness" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_harness"]:
+                                self.harness_status[worker_id] = True
                             
                             tracks.append({
                                 'id': worker_id,
@@ -194,11 +181,10 @@ class BYTETracker:
                         self.worker_history[self.next_id] = [[x, y]]
                         self.last_positions[self.next_id] = [x, y]
                         
-                        # Update helmet status if this is a helmet detection
-                        if cl == 0:  # Helmet violation class
-                            # Higher confidence for helmet violations
-                            if score > 0.45:  # Increased threshold for helmet violations
-                                self.helmet_status[self.next_id] = True
+                        if cl == "no_helmet" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
+                            self.helmet_status[self.next_id] = True
+                        elif cl == "no_harness" and score > CONFIG["CONFIDENCE_THRESHOLDS"]["no_harness"]:
+                            self.harness_status[self.next_id] = True
                         
                         tracks.append({
                             'id': self.next_id,
@@ -228,24 +214,23 @@ class BYTETracker:
     def _is_same_worker(self, pos1, pos2, threshold=150):
         x1, y1 = pos1
         x2, y2 = pos2
-        distance = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
-        return distance < threshold
+        return np.sqrt((x1 - x2)**2 + (y1 - y2)**2) < threshold
 
-    # Function to validate if a helmet violation is consistent across frames
     def validate_helmet_violation(self, worker_id, current_confidence):
-        # If we have consistent high confidence or multiple detections, it's a valid violation
         return worker_id in self.helmet_status and self.helmet_status[worker_id]
 
+    def validate_harness_violation(self, worker_id, current_confidence):
+        return worker_id in self.harness_status and self.harness_status[worker_id]
+
 # ========================== # Optimized Configuration # ==========================
 CONFIG = {
-    "MODEL_PATH": "yolov8_safety.pt",
-    "FALLBACK_MODEL": "yolov8n.pt",
+    "MODEL_NAME": "facebook/detr-resnet-50",  # Fine-tune with your dataset, e.g., "your-username/detr-resnet-50-finetuned-safety"
     "VIOLATION_LABELS": {
-        0: "no_helmet",
-        1: "no_harness",
-        2: "unsafe_posture",
-        3: "unsafe_zone",
-        4: "improper_tool_use"
+        "no_helmet": "No Helmet",
+        "no_harness": "No Harness",
+        "unsafe_posture": "Unsafe Posture",
+        "unsafe_zone": "Unsafe Zone",
+        "improper_tool_use": "Improper Tool Use"
     },
     "CLASS_COLORS": {
         "no_helmet": (0, 0, 255),
@@ -269,18 +254,18 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.45,  # Increased threshold for helmet violations
+        "no_helmet": 0.45,
         "no_harness": 0.25,
         "unsafe_posture": 0.25,
         "unsafe_zone": 0.25,
         "improper_tool_use": 0.25
     },
-    "MIN_VIOLATION_FRAMES": 2,  # Increased to require multiple frames for confirmation
+    "MIN_VIOLATION_FRAMES": 2,
     "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 10.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,  # Increased frame skip for faster processing
-    "BATCH_SIZE": 8,  # Increased batch size for better GPU utilization
+    "FRAME_SKIP": 2,
+    "BATCH_SIZE": 8,
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 150,
     "TRACK_THRESH": 0.3,
@@ -288,7 +273,7 @@ CONFIG = {
     "SNAPSHOT_QUALITY": 95,
     "MAX_WORKER_DISTANCE": 150,
     "TARGET_RESOLUTION": (384, 384),
-    "HELMET_VALIDATION_FRAMES": 3  # Number of frames to validate helmet violations
+    "HELMET_VALIDATION_FRAMES": 3
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -296,73 +281,72 @@ logger.info(f"Using device: {device}")
 
 def load_model():
     try:
-        if os.path.isfile(CONFIG["MODEL_PATH"]):
-            model_path = CONFIG["MODEL_PATH"]
-            logger.info(f"Model loaded: {model_path}")
-        else:
-            model_path = CONFIG["FALLBACK_MODEL"]
-            logger.warning("Using fallback model. Train yolov8_safety.pt for best results.")
-            if not os.path.isfile(model_path):
-                logger.info(f"Downloading fallback model: {model_path}")
-                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)
+        processor = DetrImageProcessor.from_pretrained(CONFIG["MODEL_NAME"])
+        model = DetrForObjectDetection.from_pretrained(CONFIG["MODEL_NAME"]).to(device)
         if device.type == "cuda":
-            model.model.half()
-        logger.info(f"Model classes: {model.names}")
-        return model
+            model = model.half()
+        logger.info(f"Loaded DETR model: {CONFIG['MODEL_NAME']}")
+        logger.info(f"Model classes: {model.config.id2label}")
+        return processor, model
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
         raise
 
-model = load_model()
+processor, model = load_model()
 
 # ========================== # Helper Functions # ==========================
 def preprocess_frame(frame):
     target_res = CONFIG["TARGET_RESOLUTION"]
-    # Enhanced preprocessing for better helmet detection
     frame = cv2.resize(frame, target_res, interpolation=cv2.INTER_LINEAR)
-    # Increase contrast to better differentiate helmets from other head coverings
-    frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=20)  # Increased contrast
-    
-    # Additional preprocessing to enhance head/helmet features
-    # Apply slight sharpening to make edges more distinct
-    kernel = np.array([[-1,-1,-1], 
-                       [-1, 9,-1],
-                       [-1,-1,-1]])
+    frame = cv2.convertScaleAbs(frame, alpha=1.3, beta=20)
+    kernel = np.array([[-1,-1,-1], [-1, 9,-1], [-1,-1,-1]])
     frame = cv2.filter2D(frame, -1, kernel)
-    
     return frame
 
+def is_unsafe_posture(box, frame_shape):
+    """Placeholder for unsafe posture detection. Replace with pose estimation (e.g., MediaPipe)."""
+    x1, y1, x2, y2 = box
+    height = y2 - y1
+    width = x2 - x1
+    aspect_ratio = height / max(width, 1)
+    return aspect_ratio > 2.0  # Tall, narrow box suggests bending/unsafe posture
+
+def is_improper_tool_use(person_box, tool_box):
+    """Placeholder for improper tool use. Fine-tune DETR for specific tools."""
+    person_center = ((person_box[0] + person_box[2]) / 2, (person_box[1] + person_box[3]) / 2)
+    tool_center = ((tool_box[0] + tool_box[2]) / 2, (tool_box[1] + tool_box[3]) / 2)
+    dist = distance.euclidean(person_center, tool_center)
+    return dist > 100  # Tool too far from person
+
+def is_unsafe_zone(person_box, frame_shape):
+    """Check if person is in restricted area (e.g., top-left quadrant)."""
+    px, py, pw, ph = person_box
+    frame_h, frame_w = frame_shape
+    person_center = (px + pw / 2, py + ph / 2)
+    unsafe_zone = (0, 0, 0.5, 0.5)  # Top-left quadrant
+    return (unsafe_zone[0] * frame_w < person_center[0] < unsafe_zone[2] * frame_w and
+            unsafe_zone[1] * frame_h < person_center[1] < unsafe_zone[3] * frame_h)
+
 def draw_detections(frame, detections):
     result_frame = frame.copy()
-    
     for det in detections:
         label = det.get("violation", "Unknown")
         confidence = det.get("confidence", 0.0)
         x, y, w, h = det.get("bounding_box", [0, 0, 0, 0])
         worker_id = det.get("worker_id", "Unknown")
-
         x1 = int(x - w/2)
         y1 = int(y - h/2)
         x2 = int(x + w/2)
         y2 = int(y + h/2)
-        
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
-        
-        # Make no_helmet violations more prominent
         line_thickness = 4 if label == "no_helmet" else 3
-        
         cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, line_thickness)
-        
         display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
         text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
         cv2.rectangle(result_frame, (x1, y1-text_size[1]-10), (x1+text_size[0]+10, y1), (0, 0, 0), -1)
         cv2.putText(result_frame, display_text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-        
         conf_text = f"Conf: {confidence:.2f}"
         cv2.putText(result_frame, conf_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
-        
     return result_frame
 
 def calculate_safety_score(violations):
@@ -373,23 +357,15 @@ def calculate_safety_score(violations):
         "unsafe_zone": 35,
         "improper_tool_use": 25
     }
-    
     worker_violations = {}
     for v in violations:
         worker_id = v.get("worker_id", "Unknown")
         violation_type = v.get("violation", "Unknown")
-        
         if worker_id not in worker_violations:
             worker_violations[worker_id] = set()
         worker_violations[worker_id].add(violation_type)
-    
-    total_penalty = 0
-    for worker_violations_set in worker_violations.values():
-        worker_penalty = sum(penalties.get(v, 0) for v in worker_violations_set)
-        total_penalty += worker_penalty
-    
-    score = max(0, 100 - total_penalty)
-    return score
+    total_penalty = sum(sum(penalties.get(v, 0) for v in worker_violations[wid]) for wid in worker_violations)
+    return max(0, 100 - total_penalty)
 
 def generate_violation_pdf(violations, score, output_dir):
     try:
@@ -397,70 +373,55 @@ def generate_violation_pdf(violations, score, output_dir):
         pdf_path = os.path.join(output_dir, pdf_filename)
         pdf_file = BytesIO()
         c = canvas.Canvas(pdf_file, pagesize=letter)
-        
         c.setFont("Helvetica-Bold", 16)
         c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
-        
         c.setFont("Helvetica", 12)
         c.drawString(1 * inch, 9.5 * inch, f"Date: {time.strftime('%Y-%m-%d')}")
         c.drawString(1 * inch, 9.2 * inch, f"Time: {time.strftime('%H:%M:%S')}")
-        
         c.setFont("Helvetica-Bold", 14)
         c.drawString(1 * inch, 8.7 * inch, f"Safety Compliance Score: {score}%")
-
         y_position = 8.2 * inch
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Summary:")
         y_position -= 0.3 * inch
-        
         worker_violations = {}
         for v in violations:
             worker_id = v.get("worker_id", "Unknown")
             if worker_id not in worker_violations:
                 worker_violations[worker_id] = []
             worker_violations[worker_id].append(v)
-        
         c.setFont("Helvetica", 10)
         summary_data = {
             "Total Workers with Violations": len(worker_violations),
             "Total Violations Found": len(violations),
             "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
         }
-        
         for key, value in summary_data.items():
             c.drawString(1 * inch, y_position, f"{key}: {value}")
             y_position -= 0.25 * inch
-
         y_position -= 0.5 * inch
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Violations by Worker:")
         y_position -= 0.3 * inch
-        
         c.setFont("Helvetica", 10)
         for worker_id, worker_vios in worker_violations.items():
             c.drawString(1 * inch, y_position, f"Worker {worker_id}:")
             y_position -= 0.2 * inch
-            
             for v in worker_vios:
                 display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
                 time_str = f"{v.get('timestamp', 0.0):.2f}s"
                 conf_str = f"{v.get('confidence', 0.0):.2f}"
-                
                 violation_text = f"  - {display_name} at {time_str} (Confidence: {conf_str})"
                 c.drawString(1.2 * inch, y_position, violation_text)
                 y_position -= 0.2 * inch
-                
                 if y_position < 1 * inch:
                     c.showPage()
                     c.setFont("Helvetica", 10)
                     y_position = 10 * inch
-
         c.save()
         pdf_file.seek(0)
-
         with open(pdf_path, "wb") as f:
             f.write(pdf_file.getvalue())
-            
         public_url = f"{CONFIG['PUBLIC_URL_BASE']}{pdf_filename}"
         logger.info(f"PDF generated: {public_url}")
         return pdf_path, public_url, pdf_file
@@ -484,7 +445,6 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         if not pdf_file:
             logger.error("No PDF file provided for upload")
             return ""
-            
         encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
         content_version_data = {
             "Title": f"Safety_Violation_Report_{int(time.time())}",
@@ -494,11 +454,9 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
         }
         content_version = sf.ContentVersion.create(content_version_data)
         result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
-        
         if not result['records']:
             logger.error("Failed to retrieve ContentVersion")
             return ""
-            
         file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
         logger.info(f"PDF uploaded to Salesforce: {file_url}")
         return file_url
@@ -509,21 +467,16 @@ def upload_pdf_to_salesforce(sf, pdf_file, report_id):
 def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
-        
         violations_text = ""
         for v in violations:
             display_name = CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')
             worker_id = v.get('worker_id', 'Unknown')
             timestamp = v.get('timestamp', 0.0)
             confidence = v.get('confidence', 0.0)
-            
             violations_text += f"Worker {worker_id}: {display_name} at {timestamp:.2f}s (Conf: {confidence:.2f})\n"
-        
         if not violations_text:
             violations_text = "No violations detected."
-            
         pdf_url = f"{CONFIG['PUBLIC_URL_BASE']}{os.path.basename(pdf_path)}" if pdf_path else ""
-
         record_data = {
             "Compliance_Score__c": score,
             "Violations_Found__c": len(violations),
@@ -531,9 +484,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             "Status__c": "Pending",
             "PDF_Report_URL__c": pdf_url
         }
-        
         logger.info(f"Creating Salesforce record with data: {record_data}")
-        
         try:
             record = sf.Safety_Video_Report__c.create(record_data)
             logger.info(f"Created Safety_Video_Report__c record: {record['id']}")
@@ -541,9 +492,7 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
             logger.error(f"Failed to create Safety_Video_Report__c: {e}")
             record = sf.Account.create({"Name": f"Safety_Report_{int(time.time())}"})
             logger.warning(f"Fell back to Account record: {record['id']}")
-            
         record_id = record["id"]
-
         if pdf_file:
             uploaded_url = upload_pdf_to_salesforce(sf, pdf_file, record_id)
             if uploaded_url:
@@ -555,7 +504,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
                     sf.Account.update(record_id, {"Description": uploaded_url})
                     logger.info(f"Updated Account record {record_id} with PDF URL")
                 pdf_url = uploaded_url
-
         return record_id, pdf_url
     except Exception as e:
         logger.error(f"Salesforce record creation failed: {e}")
@@ -570,102 +518,60 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 def verify_and_open_video(video_path):
     if not os.path.exists(video_path):
         raise FileNotFoundError(f"Temporary video file not found: {video_path}")
-    
     file_size = os.path.getsize(video_path)
     if file_size == 0:
         raise ValueError(f"Temporary video file is empty: {video_path}")
-    
     with open(video_path, "rb") as f:
         f.read(1)
-    
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
         raise ValueError("Could not open video file. Ensure the video format is supported (e.g., MP4) and FFmpeg is installed.")
-    
     return cap
 
-# Helper for helmet validation
 def validate_helmet_detection(frame, bbox, confidence_threshold=0.45):
-    """
-    Additional validation for helmet detection to reduce false positives.
-    This function performs additional checks on the region to confirm it's a true helmet violation.
-    """
     x, y, w, h = bbox
     x1 = int(max(0, x - w/2))
     y1 = int(max(0, y - h/2))
     x2 = int(min(frame.shape[1], x + w/2))
     y2 = int(min(frame.shape[0], y + h/2))
-    
-    # Extract head region
     head_region = frame[y1:y2, x1:x2]
     if head_region.size == 0:
         return False
-    
-    # Check if this is truly a helmet violation by analyzing the region
-    # 1. Check color distribution - helmets often have more uniform color
     hsv = cv2.cvtColor(head_region, cv2.COLOR_BGR2HSV)
-    
-    # Check for typical helmet colors (many construction helmets are yellow, white, orange, blue)
-    # This helps differentiate from cloth head coverings
     yellow_lower = np.array([20, 100, 100])
     yellow_upper = np.array([30, 255, 255])
     yellow_mask = cv2.inRange(hsv, yellow_lower, yellow_upper)
-    
     white_lower = np.array([0, 0, 200])
     white_upper = np.array([180, 30, 255])
     white_mask = cv2.inRange(hsv, white_lower, white_upper)
-    
     orange_lower = np.array([5, 100, 100])
     orange_upper = np.array([15, 255, 255])
     orange_mask = cv2.inRange(hsv, orange_lower, orange_upper)
-    
     blue_lower = np.array([100, 100, 100])
     blue_upper = np.array([130, 255, 255])
     blue_mask = cv2.inRange(hsv, blue_lower, blue_upper)
-    
     helmet_mask = cv2.bitwise_or(yellow_mask, white_mask)
     helmet_mask = cv2.bitwise_or(helmet_mask, orange_mask)
     helmet_mask = cv2.bitwise_or(helmet_mask, blue_mask)
-    
-    # If there's a significant amount of helmet-colored pixels, this might be a helmet
     helmet_percentage = np.sum(helmet_mask > 0) / (head_region.shape[0] * head_region.shape[1])
-    
-    # If the region has a significant amount of helmet-like colors, it's probably a helmet
-    # so we should NOT flag it as a violation (return False)
     if helmet_percentage > 0.25:
         return False
-    
-    # Check texture uniformity - helmets have more uniform texture compared to head coverings
     gray = cv2.cvtColor(head_region, cv2.COLOR_BGR2GRAY)
     texture_score = np.std(gray)
-    
-    # If texture is very uniform (low standard deviation), it might be a helmet or bare head
-    # Very uniform texture (like a hard helmet) would have low texture_score
-    if texture_score < 15:  # Low texture suggests uniform surface like a helmet
+    if texture_score < 15:
         return False
-    
-    # Additional check for cloth-like textures
     edges = cv2.Canny(gray, 50, 150)
     edge_density = np.sum(edges > 0) / (head_region.shape[0] * head_region.shape[1])
-    
-    # If there are many edges (cloth wrinkles), this might be a kurchief
     if edge_density > 0.15:
-        # This is likely a cloth head covering, not a helmet violation
-        # But also not a proper helmet, so we should still detect as violation
         return True
-    
-    # If confidence is very high, trust the model
     if confidence_threshold >= 0.6:
         return True
-        
-    # Default to the original detection
     return True
 
 def process_video(video_data, temp_dir):
     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
     
     try:
         if not video_data:
@@ -681,16 +587,7 @@ def process_video(video_data, temp_dir):
             video_path = temp_file.name
         logger.info(f"Video saved to temporary file: {video_path}")
 
-        if not os.path.exists(video_path):
-            raise FileNotFoundError(f"Temporary video file not found: {video_path}")
-        file_size = os.path.getsize(video_path)
-        if file_size == 0:
-            raise ValueError(f"Temporary video file is empty: {video_path}")
-        logger.info(f"Temporary video file size: {file_size} bytes")
-
         cap = verify_and_open_video(video_path)
-        logger.info(f"Successfully opened video file: {video_path}")
-
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
         duration = total_frames / fps
@@ -711,8 +608,7 @@ def process_video(video_data, temp_dir):
         worker_id_mapping = {}
         unique_violations = {}
         violation_frames = {}
-        # Track helmet detections across frames for each worker
-        helmet_detections = {}  
+        helmet_detections = {}
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
         processed_frames = 0
@@ -722,28 +618,22 @@ def process_video(video_data, temp_dir):
         while processed_frames < total_frames:
             batch_frames = []
             batch_indices = []
-            batch_originals = []  # Store original frames for helmet validation
+            batch_originals = []
             
             for _ in range(CONFIG["BATCH_SIZE"]):
                 frame_idx = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
                 if frame_idx >= total_frames:
                     break
-                
                 ret, frame = cap.read()
                 if not ret:
                     logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
                     break
-                
-                # Store original frame for validation
                 original_frame = frame.copy()
-                
                 frame = preprocess_frame(frame)
-                
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
                         break
-                
-                batch_frames.append(frame)
+                batch_frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
                 batch_indices.append(frame_idx)
                 batch_originals.append(original_frame)
                 processed_frames += 1
@@ -753,16 +643,16 @@ def process_video(video_data, temp_dir):
                 break
 
             try:
-                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)
+                inputs = processor(images=batch_frames, return_tensors="pt").to(device)
                 if device.type == "cuda":
-                    batch_frames_tensor = batch_frames_tensor.half()
-
-                results = model(batch_frames_tensor, device=device, conf=0.1, verbose=False)
+                    inputs = {k: v.half() for k, v in inputs.items()}
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                target_sizes = torch.tensor([frame.size[::-1] for frame in batch_frames]).to(device)
+                results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.1)
             except Exception as e:
                 logger.error(f"Model inference failed: {e}")
-                raise ValueError(f"Failed to process video frames with YOLO model: {str(e)}")
+                raise ValueError(f"Failed to process video frames with DETR model: {str(e)}")
             finally:
                 batch_frames = []
                 if device.type == "cuda":
@@ -778,39 +668,37 @@ def process_video(video_data, temp_dir):
 
             for i, (result, frame_idx, original_frame) in enumerate(zip(results, batch_indices, batch_originals)):
                 current_time = frame_idx / fps
-                
-                boxes = result.boxes
                 track_inputs = []
-                
-                for box in boxes:
-                    cls = int(box.cls)
-                    conf = float(box.conf)
-                    label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                    
-                    if label is None:
-                        continue
-                        
-                    # Enhanced confidence threshold handling, especially for helmet detection
-                    if label == "no_helmet":
-                        if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.45):
-                            continue
-                        
-                        # Additional validation for helmet detection
-                        bbox = box.xywh.cpu().numpy()[0]
-                        if not validate_helmet_detection(original_frame, bbox, conf):
+                person_boxes = []
+                tool_boxes = []
+
+                for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
+                    label_name = model.config.id2label[label.item()]
+                    conf = float(score)
+                    bbox = box.cpu().numpy()
+                    x, y, x2, y2 = bbox
+                    w, h = x2 - x, y2 - y
+                    bbox_xywh = [x + w/2, y + h/2, w, h]
+
+                    if label_name in ["no_helmet", "no_harness"] and conf >= CONFIG["CONFIDENCE_THRESHOLDS"].get(label_name, 0.25):
+                        if label_name == "no_helmet" and not validate_helmet_detection(original_frame, bbox_xywh, conf):
                             logger.info(f"Frame {frame_idx}: Helmet false positive filtered at {conf:.2f} confidence")
                             continue
-                    else:
-                        # Use regular thresholds for other violations
-                        if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
-                            continue
-
-                    bbox = box.xywh.cpu().numpy()[0]
-                    track_inputs.append({
-                        "bbox": bbox,
-                        "conf": conf,
-                        "cls": cls
-                    })
+                        track_inputs.append({"bbox": bbox_xywh, "conf": conf, "cls": label_name})
+                    elif label_name == "person":
+                        person_boxes.append(bbox_xywh)
+                    elif label_name in ["hammer", "wrench"]:  # Example tools; update with your dataset
+                        tool_boxes.append(bbox_xywh)
+
+                # Handle Unsafe Posture, Unsafe Zone, Improper Tool Use
+                for pbox in person_boxes:
+                    if is_unsafe_posture(pbox, original_frame.shape[:2]):
+                        track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "unsafe_posture"})
+                    if is_unsafe_zone(pbox, original_frame.shape[:2]):
+                        track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "unsafe_zone"})
+                    for tbox in tool_boxes:
+                        if is_improper_tool_use(pbox, tbox):
+                            track_inputs.append({"bbox": pbox, "conf": 0.9, "cls": "improper_tool_use"})
 
                 if not track_inputs:
                     continue
@@ -824,11 +712,11 @@ def process_video(video_data, temp_dir):
 
                 for obj in tracked_objects:
                     tracker_id = obj['id']
-                    label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
+                    label = obj['cls']
                     conf = obj['score']
                     bbox = obj['bbox']
                     
-                    if label is None:
+                    if label not in CONFIG["VIOLATION_LABELS"]:
                         continue
                     
                     if tracker_id not in worker_id_mapping:
@@ -837,25 +725,16 @@ def process_video(video_data, temp_dir):
                     
                     worker_id = worker_id_mapping[tracker_id]
                     
-                    # Special handling for helmet violations to ensure consistency
                     if label == "no_helmet":
-                        # Track helmet violations for this worker
                         if worker_id not in helmet_detections:
                             helmet_detections[worker_id] = []
-                        
-                        # Store this detection with frame index and confidence
                         helmet_detections[worker_id].append({
                             "frame_idx": frame_idx,
                             "confidence": conf,
                             "bbox": bbox
                         })
-                        
-                        # Only record a helmet violation if we have multiple consistent detections
                         if len(helmet_detections[worker_id]) >= CONFIG["HELMET_VALIDATION_FRAMES"]:
-                            # Calculate average confidence
                             avg_conf = sum(d["confidence"] for d in helmet_detections[worker_id]) / len(helmet_detections[worker_id])
-                            
-                            # If confidence is consistently high across multiple frames, record the violation
                             if avg_conf >= CONFIG["CONFIDENCE_THRESHOLDS"]["no_helmet"]:
                                 violation_key = (worker_id, label)
                                 if violation_key not in unique_violations:
@@ -863,7 +742,6 @@ def process_video(video_data, temp_dir):
                                     violation_frames[violation_key] = frame_idx
                                     logger.info(f"Frame {frame_idx}: Valid helmet violation for worker {worker_id} with avg conf {avg_conf:.2f}")
                     else:
-                        # Regular handling for other violations
                         violation_key = (worker_id, label)
                         if violation_key not in unique_violations:
                             unique_violations[violation_key] = current_time
@@ -900,26 +778,29 @@ def process_video(video_data, temp_dir):
                 continue
 
             frame = preprocess_frame(frame)
-            frame_tensor = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
-            frame_tensor = frame_tensor.unsqueeze(0).to(device)
+            frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            inputs = processor(images=frame_pil, return_tensors="pt").to(device)
             if device.type == "cuda":
-                frame_tensor = frame_tensor.half()
-
-            result = model(frame_tensor, device=device, conf=0.1, verbose=False)[0]
-            boxes = result.boxes
-
-            for box in boxes:
-                cls = int(box.cls)
-                conf = float(box.conf)
-                label = CONFIG["VIOLATION_LABELS"].get(cls, None)
-                if label == violation["violation"]:
+                inputs = {k: v.half() for k, v in inputs.items()}
+            with torch.no_grad():
+                outputs = model(**inputs)
+            target_sizes = torch.tensor([frame_pil.size[::-1]]).to(device)
+            result = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.1)[0]
+
+            for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
+                label_name = model.config.id2label[label.item()]
+                conf = float(score)
+                bbox = box.cpu().numpy()
+                x, y, x2, y2 = bbox
+                w, h = x2 - x, y2 - y
+                bbox_xywh = [x + w/2, y + h/2, w, h]
+                if label_name == violation["violation"]:
                     violation["confidence"] = round(conf, 2)
-                    bbox = box.xywh.cpu().numpy()[0]
                     detection = {
                         "worker_id": violation["worker_id"],
-                        "violation": label,
+                        "violation": label_name,
                         "confidence": violation["confidence"],
-                        "bounding_box": bbox,
+                        "bounding_box": bbox_xywh,
                         "timestamp": violation["timestamp"]
                     }
                     snapshot_frame = frame.copy()
@@ -933,7 +814,7 @@ def process_video(video_data, temp_dir):
                         (255, 255, 255),
                         2
                     )
-                    snapshot_filename = f"violation_{label}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
+                    snapshot_filename = f"violation_{label_name}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
                     snapshot_path = os.path.join(output_dir, snapshot_filename)
                     cv2.imwrite(
                         snapshot_path,
@@ -941,14 +822,14 @@ def process_video(video_data, temp_dir):
                         [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]]
                     )
                     snapshots.append({
-                        "violation": label,
+                        "violation": label_name,
                         "worker_id": violation["worker_id"],
                         "timestamp": violation["timestamp"],
                         "snapshot_path": snapshot_path,
                         "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
                         "confidence": violation["confidence"]
                     })
-                    logger.info(f"Captured snapshot for {label} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
+                    logger.info(f"Captured snapshot for {label_name} violation by worker {violation['worker_id']} at {violation['timestamp']:.2f}s")
                     break
 
         cap.release()
@@ -1007,7 +888,7 @@ def gradio_interface(video_file):
         if not video_file:
             return "No file uploaded.", "", "No file uploaded.", "", ""
         
-        temp_dir = tempfile.mkdtemp(prefix="Ultralytics_")
+        temp_dir = tempfile.mkdtemp(prefix="DETR_")
         logger.info(f"Created temporary directory for video processing: {temp_dir}")
 
         with open(video_file, "rb") as f:
@@ -1063,5 +944,5 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Enhanced Safety Analyzer App...")
+    logger.info("Launching Enhanced Safety Analyzer App with DETR...")
     interface.launch()