Update app.py

app.py

@@ -28,6 +28,13 @@ CONFIG = {
         3: "unsafe_zone",
         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) # Yellow
+    },
     "DISPLAY_NAMES": {
         "no_helmet": "No Helmet Violation",
         "no_harness": "No Harness Violation",
@@ -42,17 +49,11 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 10,  # Reduced for better detection
-    "MAX_PROCESSING_TIME": 45,
-    "CONFIDENCE_THRESHOLD": {
-        "no_helmet": 0.4,
-        "no_harness": 0.35,
-        "unsafe_posture": 0.3,
-        "unsafe_zone": 0.3,
-        "improper_tool_use": 0.35
-    },
+    "FRAME_SKIP": 5,  # Reduced for better detection
+    "MAX_PROCESSING_TIME": 60,
+    "CONFIDENCE_THRESHOLD": 0.25,  # Lower threshold for all violations
     "IOU_THRESHOLD": 0.4,
-    "MIN_VIOLATION_DURATION": 2  # seconds
+    "MIN_VIOLATION_FRAMES": 3  # Minimum consecutive frames to confirm violation
 }
 
 # Setup logging
@@ -86,6 +87,27 @@ model = load_model()
 # ==========================
 # Enhanced Helper Functions
 # ==========================
+def draw_detections(frame, detections):
+    """Draw bounding boxes and labels on frame"""
+    for det in detections:
+        label = det["violation"]
+        confidence = det["confidence"]
+        x, y, w, h = det["bounding_box"]
+        
+        # Convert from center coordinates to corner coordinates
+        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))
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        
+        display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)}: {confidence:.2f}"
+        cv2.putText(frame, display_text, (x1, y1-10), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    return frame
+
 def calculate_iou(box1, box2):
     """Calculate Intersection over Union (IoU) for two bounding boxes."""
     x1, y1, w1, h1 = box1
@@ -110,15 +132,6 @@ def calculate_iou(box1, box2):
     
     return intersection / union if union > 0 else 0
 
-def is_violation_persistent(violation_type, worker_id, violations, fps):
-    """Check if a violation persists for the required duration."""
-    violation_times = [v['timestamp'] for v in violations 
-                      if v['violation'] == violation_type and v['worker_id'] == worker_id]
-    if len(violation_times) < 2:
-        return False
-    duration = max(violation_times) - min(violation_times)
-    return duration >= CONFIG["MIN_VIOLATION_DURATION"]
-
 # ==========================
 # Salesforce Integration (unchanged)
 # ==========================
@@ -281,16 +294,18 @@ def process_video(video_data):
         if not video.isOpened():
             raise ValueError("Could not open video file")
 
-        violations, snapshots = [], []
+        violations = []
+        snapshots = []
         frame_count = 0
         start_time = time.time()
         fps = video.get(cv2.CAP_PROP_FPS)
         if fps <= 0:
             fps = 30  # Default assumption if FPS cannot be determined
 
+        # Structure to track workers and their violations
+        workers = []
+        violation_history = {label: [] for label in CONFIG["VIOLATION_LABELS"].values()}
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
-        workers = []  # List to track workers
-        violation_history = []  # Track all potential violations before filtering
 
         logger.info(f"Processing video with FPS: {fps}")
         logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
@@ -308,31 +323,28 @@ def process_video(video_data):
                 logger.info("Processing time limit reached")
                 break
 
+            current_time = frame_count / fps
+            
             # Run detection on this frame
             results = model(frame, device=device)
-            current_time = frame_count / fps
             
+            current_detections = []
             for result in results:
                 boxes = result.boxes
-                logger.debug(f"Frame {frame_count}: Found {len(boxes)} potential detections")
-                
                 for box in boxes:
                     cls = int(box.cls)
                     conf = float(box.conf)
                     label = CONFIG["VIOLATION_LABELS"].get(cls, None)
                     
                     if label is None:
-                        continue  # Skip unknown classes
+                        continue
                         
-                    conf_threshold = CONFIG["CONFIDENCE_THRESHOLD"].get(label, 0.3)
-                    if conf < conf_threshold:
-                        logger.debug(f"Skipping {label} with low confidence: {conf:.2f} < {conf_threshold}")
+                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
                         continue
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     
-                    # Store potential violation (will filter later)
-                    violation_history.append({
+                    current_detections.append({
                         "frame": frame_count,
                         "violation": label,
                         "confidence": round(conf, 2),
@@ -340,72 +352,73 @@ def process_video(video_data):
                         "timestamp": current_time
                     })
 
+            # Process detections and associate with workers
+            for detection in current_detections:
+                # Find matching worker
+                matched_worker = None
+                max_iou = 0
+                
+                for worker in workers:
+                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
+                        max_iou = iou
+                        matched_worker = worker
+
+                if matched_worker:
+                    # Update worker's position
+                    matched_worker["bbox"] = detection["bounding_box"]
+                    matched_worker["last_seen"] = current_time
+                    worker_id = matched_worker["id"]
+                else:
+                    # New worker
+                    worker_id = len(workers) + 1
+                    workers.append({
+                        "id": worker_id,
+                        "bbox": detection["bounding_box"],
+                        "first_seen": current_time,
+                        "last_seen": current_time
+                    })
+                
+                # Add to violation history
+                detection["worker_id"] = worker_id
+                violation_history[detection["violation"]].append(detection)
+
             frame_count += 1
 
         video.release()
         os.remove(video_path)
         
-        # Process violation history to track workers and persistent violations
-        workers = []
-        for v in violation_history:
-            # Find matching worker
-            matched_worker = None
-            max_iou = 0
-            
-            for worker in workers:
-                iou = calculate_iou(v["bounding_box"], worker["bbox"])
-                if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
-                    max_iou = iou
-                    matched_worker = worker
-
-            if matched_worker:
-                # Update worker's violation history
-                matched_worker["violations"].append(v)
-                matched_worker["bbox"] = v["bounding_box"]
-                matched_worker["last_seen"] = v["timestamp"]
-                v["worker_id"] = matched_worker["id"]
-            else:
-                # New worker
-                worker_id = len(workers) + 1
-                workers.append({
-                    "id": worker_id,
-                    "bbox": v["bounding_box"],
-                    "violations": [v],
-                    "first_seen": v["timestamp"],
-                    "last_seen": v["timestamp"]
-                })
-                v["worker_id"] = worker_id
-
-        # Filter violations to only include those that persist for minimum duration
-        final_violations = []
-        for worker in workers:
-            # Group violations by type
-            violations_by_type = {}
-            for v in worker["violations"]:
-                if v["violation"] not in violations_by_type:
-                    violations_by_type[v["violation"]] = []
-                violations_by_type[v["violation"]].append(v)
-            
-            # Check each violation type for persistence
-            for violation_type, v_list in violations_by_type.items():
-                if len(v_list) < 2:
-                    continue  # Need multiple detections to check duration
+        # Process violation history to confirm persistent violations
+        for violation_type, detections in violation_history.items():
+            if not detections:
+                continue
                 
-                duration = max(v["timestamp"] for v in v_list) - min(v["timestamp"] for v in v_list)
-                if duration >= CONFIG["MIN_VIOLATION_DURATION"]:
+            # Group by worker
+            worker_violations = {}
+            for det in detections:
+                if det["worker_id"] not in worker_violations:
+                    worker_violations[det["worker_id"]] = []
+                worker_violations[det["worker_id"]].append(det)
+            
+            # Check each worker's violations for persistence
+            for worker_id, worker_dets in worker_violations.items():
+                if len(worker_dets) >= CONFIG["MIN_VIOLATION_FRAMES"]:
                     # Take the highest confidence detection
-                    best_detection = max(v_list, key=lambda x: x["confidence"])
-                    final_violations.append(best_detection)
+                    best_detection = max(worker_dets, key=lambda x: x["confidence"])
+                    violations.append(best_detection)
                     
                     # Capture snapshot if not already taken
                     if not snapshot_taken[violation_type]:
-                        # We need to get the frame for this violation
+                        # Get the frame for this violation
                         cap = cv2.VideoCapture(video_path)
                         cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
                         ret, snapshot_frame = cap.read()
                         cap.release()
                         
                         if ret:
+                            # Draw detections on snapshot
+                            snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                            
                             snapshot_filename = f"{violation_type}_{best_detection['frame']}.jpg"
                             snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
                             cv2.imwrite(snapshot_path, snapshot_frame)
@@ -418,7 +431,7 @@ def process_video(video_data):
                             snapshot_taken[violation_type] = True
 
         # Final processing
-        if not final_violations:
+        if not violations:
             logger.info("No persistent violations detected")
             return {
                 "violations": [],
@@ -429,12 +442,12 @@ def process_video(video_data):
                 "message": "No violations detected in the video."
             }
 
-        score = calculate_safety_score(final_violations)
-        pdf_path, pdf_url, pdf_file = generate_violation_pdf(final_violations, score)
-        report_id, final_pdf_url = push_report_to_salesforce(final_violations, score, pdf_path, pdf_file)
+        score = calculate_safety_score(violations)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
         return {
-            "violations": final_violations,
+            "violations": violations,
             "snapshots": snapshots,
             "score": score,
             "salesforce_record_id": report_id,
@@ -453,7 +466,7 @@ def process_video(video_data):
         }
 
 # ==========================
-# Gradio Interface (unchanged)
+# Gradio Interface
 # ==========================
 def gradio_interface(video_file):
     if not video_file: