Update app.py

app.py

@@ -13,6 +13,7 @@ from io import BytesIO
 import base64
 import logging
 from retrying import retry
+import uuid
 
 # ==========================
 # Enhanced Configuration
@@ -49,9 +50,21 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 5,  # Reduced for better detection
+    "FRAME_SKIP": {
+        "no_helmet": 3,  # Lower skip for frequent violations
+        "no_harness": 2,
+        "unsafe_posture": 2,
+        "unsafe_zone": 2,
+        "improper_tool_use": 2
+    },
     "MAX_PROCESSING_TIME": 60,
-    "CONFIDENCE_THRESHOLD": 0.25,  # Lower threshold for all violations
+    "CONFIDENCE_THRESHOLDS": {
+        "no_helmet": 0.3,  # Slightly higher to reduce false positives
+        "no_harness": 0.2,
+        "unsafe_posture": 0.2,
+        "unsafe_zone": 0.2,
+        "improper_tool_use": 0.2
+    },
     "IOU_THRESHOLD": 0.4,
     "MIN_VIOLATION_FRAMES": 3  # Minimum consecutive frames to confirm violation
 }
@@ -94,7 +107,6 @@ def draw_detections(frame, detections):
         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)
@@ -113,19 +125,12 @@ def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
-    # Convert to top-left and bottom-right coordinates
     x1_min, y1_min = x1 - w1/2, y1 - h1/2
     x1_max, y1_max = x1 + w1/2, y1 + h1/2
     x2_min, y2_min = x2 - w2/2, y2 - h2/2
     x2_max, y2_max = x2 + w2/2, y2 + h2/2
     
-    # Calculate intersection
-    x_min = max(x1_min, x2_min)
-    y_min = max(y1_min, y2_min)
-    x_max = min(x1_max, x2_max)
-    y_max = min(y1_max, y2_max)
-    
-    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
+    intersection = max(0, x1_max - x1_min) * max(0, y1_max - y1_min)
     area1 = w1 * h1
     area2 = w2 * h2
     union = area1 + area2 - intersection
@@ -133,7 +138,7 @@ def calculate_iou(box1, box2):
     return intersection / union if union > 0 else 0
 
 # ==========================
-# Salesforce Integration (unchanged)
+# Salesforce Integration
 # ==========================
 @retry(stop_max_attempt_number=3, wait_fixed=2000)
 def connect_to_salesforce():
@@ -270,13 +275,7 @@ def calculate_safety_score(violations):
         "unsafe_zone": 35,
         "improper_tool_use": 25
     }
-    # Count unique violations per worker
-    unique_violations = set()
-    for v in violations:
-        key = (v["worker_id"], v["violation"])
-        unique_violations.add(key)
-    
-    total_penalty = sum(penalties.get(violation, 0) for _, violation in unique_violations)
+    total_penalty = sum(penalties.get(v["violation"], 0) for v in violations)
     score = 100 - total_penalty
     return max(score, 0)
 
@@ -302,10 +301,11 @@ def process_video(video_data):
         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()}
+        confirmed_violations = {}  # Track confirmed violations per worker
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
+        helmet_compliance = {}  # Track workers with helmets
 
         logger.info(f"Processing video with FPS: {fps}")
         logger.info(f"Looking for violations: {CONFIG['VIOLATION_LABELS']}")
@@ -315,18 +315,18 @@ def process_video(video_data):
             if not ret:
                 break
 
-            if frame_count % CONFIG["FRAME_SKIP"] != 0:
-                frame_count += 1
-                continue
-
             if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"]:
                 logger.info("Processing time limit reached")
                 break
 
             current_time = frame_count / fps
-            
+            min_frame_skip = min(CONFIG["FRAME_SKIP"].values())
+            if frame_count % min_frame_skip != 0:
+                frame_count += 1
+                continue
+
             # Run detection on this frame
-            results = model(frame, device=device)
+            results = model(frame, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], agnostic_nms=True)
             
             current_detections = []
             for result in results:
@@ -339,14 +339,14 @@ def process_video(video_data):
                     if label is None:
                         continue
                         
-                    if conf < CONFIG["CONFIDENCE_THRESHOLD"]:
+                    if conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
                         continue
 
                     bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
                     
                     current_detections.append({
                         "frame": frame_count,
-                        "violation": label,
+                        "vi臉olation": label,
                         "confidence": round(conf, 2),
                         "bounding_box": bbox,
                         "timestamp": current_time
@@ -354,7 +354,21 @@ def process_video(video_data):
 
             # Process detections and associate with workers
             for detection in current_detections:
-                # Find matching worker
+                violation_type = detection["violation"]
+                # Skip No Helmet detection if worker is compliant
+                if violation_type == "no_helmet":
+                    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 and matched_worker["id"] in helmet_compliance:
+                        continue  # Skip if worker is known to wear a helmet
+                
+                # Find or create worker
                 matched_worker = None
                 max_iou = 0
                 
@@ -365,12 +379,10 @@ def process_video(video_data):
                         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,
@@ -379,9 +391,19 @@ def process_video(video_data):
                         "last_seen": current_time
                     })
                 
-                # Add to violation history
+                # Skip if this violation type is already confirmed for this worker
+                if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
+                    continue
+                
                 detection["worker_id"] = worker_id
-                violation_history[detection["violation"]].append(detection)
+                violation_history[violation_type].append(detection)
+                
+                # Update helmet compliance (simulate by checking if No Helmet confidence is low)
+                if violation_type == "no_helmet" and conf < 0.5:
+                    helmet_compliance[worker_id] = True
+
+            # Clean up old workers
+            workers = [w for w in workers if current_time - w["last_seen"] < 5.0]
 
             frame_count += 1
 
@@ -393,30 +415,36 @@ def process_video(video_data):
             if not detections:
                 continue
                 
-            # 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
+                    # Skip if already confirmed
+                    if worker_id in confirmed_violations and violation_type in confirmed_violations[worker_id]:
+                        continue
+                        
+                    # Skip No Helmet if worker is compliant
+                    if violation_type == "no_helmet" and worker_id in helmet_compliance:
+                        continue
+                    
                     best_detection = max(worker_dets, key=lambda x: x["confidence"])
                     violations.append(best_detection)
                     
-                    # Capture snapshot if not already taken
+                    if worker_id not in confirmed_violations:
+                        confirmed_violations[worker_id] = set()
+                    confirmed_violations[worker_id].add(violation_type)
+                    
                     if not snapshot_taken[violation_type]:
-                        # 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"
@@ -430,7 +458,6 @@ def process_video(video_data):
                             })
                             snapshot_taken[violation_type] = True
 
-        # Final processing
         if not violations:
             logger.info("No persistent violations detected")
             return {