Update app.py

app.py

@@ -49,7 +49,9 @@ CONFIG = {
         "domain": "login"
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
-    "FRAME_SKIP": 10,  # Increased to process every 10th frame (faster)
+    "FRAME_SKIP": 15,  # Increased to process every 15th frame
+    "MAX_FRAMES": 100,  # Limit the number of frames to process
+    "FRAME_RESIZE": (640, 480),  # Downscale frames to this resolution
     "MAX_PROCESSING_TIME": 60,  # Max processing time (seconds)
     "CONFIDENCE_THRESHOLD": {  # Per-class thresholds
         "no_helmet": 0.4,
@@ -58,8 +60,7 @@ CONFIG = {
         "unsafe_zone": 0.3,
         "improper_tool_use": 0.35
     },
-    "IOU_THRESHOLD": 0.4,  # For worker tracking
-    "MIN_VIOLATION_FRAMES": 2  # Reduced to 2 frames for faster confirmation
+    "MIN_VIOLATION_FRAMES": 2  # Min frames to confirm a violation
 }
 
 # Setup logging
@@ -106,18 +107,6 @@ def draw_detections(frame, detections):
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     return frame
 
-def calculate_iou(box1, box2):
-    """Compute Intersection-over-Union for tracking."""
-    x1, y1, w1, h1 = box1
-    x2, y2, w2, h2 = box2
-    x_min = max(x1 - w1/2, x2 - w2/2)
-    y_min = max(y1 - h1/2, y2 - h2/2)
-    x_max = min(x1 + w1/2, x2 + w2/2)
-    y_max = min(y1 + h1/2, y2 + h2/2)
-    intersection = max(0, x_max - x_min) * max(0, y_max - y_min)
-    union = w1 * h1 + w2 * h2 - intersection
-    return intersection / union if union > 0 else 0
-
 # ==========================
 # Salesforce Integration
 # ==========================
@@ -235,9 +224,9 @@ def process_video(video_path):
         cap = cv2.VideoCapture(video_path)
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
         frame_count = 0
+        processed_frames = 0
         violations = []
         snapshots = []
-        workers = []
         snapshot_taken = {label: False for label in CONFIG["VIOLATION_LABELS"].values()}
         frames_to_save = []  # Store frames for snapshot saving later
         
@@ -246,10 +235,17 @@ def process_video(video_path):
             if not ret:
                 break
             
+            # Stop if max frames reached
+            if processed_frames >= CONFIG["MAX_FRAMES"]:
+                break
+            
             if frame_count % CONFIG["FRAME_SKIP"] != 0:
                 frame_count += 1
                 continue
             
+            # Downscale frame for faster inference
+            frame = cv2.resize(frame, CONFIG["FRAME_RESIZE"], interpolation=cv2.INTER_AREA)
+            
             # Run detection
             results = model(frame, device=device)
             current_time = frame_count / fps
@@ -272,10 +268,8 @@ def process_video(video_path):
                         "timestamp": current_time
                     }
                     
-                    # Simplified worker tracking
-                    worker_id = len(workers) + 1  # Assign new ID without IoU for speed
-                    workers.append({"id": worker_id, "bbox": bbox})
-                    detection["worker_id"] = worker_id
+                    # Assign a generic worker ID (skipping tracking for speed)
+                    detection["worker_id"] = f"Worker_{frame_count}"
                     violations.append(detection)
                     
                     # Store frame for snapshot if first detection of this type
@@ -288,6 +282,7 @@ def process_video(video_path):
                         snapshot_taken[label] = True
             
             frame_count += 1
+            processed_frames += 1
         
         cap.release()