Update app.py

app.py

@@ -18,7 +18,7 @@ from multiprocessing import Pool, cpu_count
 from functools import partial
 
 # ==========================
-# Ultra-Fast Configuration
+# Optimized Configuration
 # ==========================
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
@@ -32,16 +32,16 @@ CONFIG = {
         4: "improper_tool_use"
     },
     "CLASS_COLORS": {
-        "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)
+        "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",
-        "unsafe_posture": "Unsafe Posture Violation",
+        "unsafe_posture": "Unsafe Posture",
         "unsafe_zone": "Unsafe Zone Entry",
         "improper_tool_use": "Improper Tool Use"
     },
@@ -53,21 +53,17 @@ CONFIG = {
     },
     "PUBLIC_URL_BASE": "https://huggingface.co/spaces/PrashanthB461/AI_Safety_Demo2/resolve/main/static/output/",
     "CONFIDENCE_THRESHOLDS": {
-        "no_helmet": 0.55,
-        "no_harness": 0.1,
-        "unsafe_posture": 0.1,
-        "unsafe_zone": 0.1,
-        "improper_tool_use": 0.1
+        "no_helmet": 0.6,  # Higher threshold to reduce false positives
+        "no_harness": 0.4,
+        "unsafe_posture": 0.4,
+        "unsafe_zone": 0.4,
+        "improper_tool_use": 0.4
     },
-    "IOU_THRESHOLD": 0.45,
-    "MIN_VIOLATION_FRAMES": 2,
-    "HELMET_CONFIDENCE_THRESHOLD": 0.6,
-    "WORKER_TRACKING_DURATION": 2.5,
-    "MAX_PROCESSING_TIME": 30,
-    "PARALLEL_WORKERS": 2,  # Reduced for Hugging Face Spaces
-    "CHUNK_SIZE": 20,  # Increased for faster batch processing
-    "FRAME_SAMPLE_RATE": 2,  # Process every 2nd frame
-    "MAX_FRAME_WIDTH": 640  # Resize frames to this width
+    "MIN_VIOLATION_FRAMES": 3,  # Require 3+ detections to confirm violation
+    "WORKER_TRACKING_DURATION": 3.0,  # Track workers for 3 seconds
+    "MAX_PROCESSING_TIME": 30,  # 30-second limit
+    "PARALLEL_WORKERS": max(1, cpu_count() - 1),  # Use all CPU cores
+    "CHUNK_SIZE": 8  # Frames per batch
 }
 
 # Setup logging
@@ -101,15 +97,6 @@ model = load_model()
 # ==========================
 # Optimized Helper Functions
 # ==========================
-def resize_frame(frame, max_width):
-    height, width = frame.shape[:2]
-    if width > max_width:
-        scale = max_width / width
-        new_width = int(width * scale)
-        new_height = int(height * scale)
-        frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
-    return frame
-
 def draw_detections(frame, detections):
     for det in detections:
         label = det.get("violation", "Unknown")
@@ -125,14 +112,14 @@ def draw_detections(frame, detections):
         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)
+        cv2.putText(frame, display_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     return frame
 
 def calculate_iou(box1, box2):
     x1, y1, w1, h1 = box1
     x2, y2, w2, h2 = box2
     
+    # Calculate intersection area
     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)
@@ -150,9 +137,7 @@ def calculate_iou(box1, box2):
 
 def process_frame_batch(frame_batch, frame_indices, fps):
     batch_results = []
-    start_inference = time.time()
-    results = model(frame_batch, device=device, conf=0.1, iou=CONFIG["IOU_THRESHOLD"], verbose=False)
-    logger.info(f"Inference time for batch of {len(frame_batch)} frames: {time.time() - start_inference:.2f}s")
+    results = model(frame_batch, device=device, conf=0.1, verbose=False)
     
     for idx, (result, frame_idx) in enumerate(zip(results, frame_indices)):
         current_time = frame_idx / fps
@@ -164,7 +149,7 @@ def process_frame_batch(frame_batch, frame_indices, fps):
             conf = float(box.conf)
             label = CONFIG["VIOLATION_LABELS"].get(cls, None)
             
-            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.2):
+            if label is None or conf < CONFIG["CONFIDENCE_THRESHOLDS"].get(label, 0.25):
                 continue
 
             bbox = [round(x, 2) for x in box.xywh.cpu().numpy()[0]]
@@ -182,7 +167,6 @@ def process_frame_batch(frame_batch, frame_indices, fps):
 
 def generate_violation_pdf(violations, score):
     try:
-        start_pdf = time.time()
         pdf_filename = f"violations_{int(time.time())}.pdf"
         pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
         pdf_file = BytesIO()
@@ -224,7 +208,6 @@ def generate_violation_pdf(violations, score):
         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 generation time: {time.time() - start_pdf:.2f}s")
         logger.info(f"PDF generated: {public_url}")
         return pdf_path, public_url, pdf_file
     except Exception as e:
@@ -244,11 +227,10 @@ def calculate_safety_score(violations):
     return max(score, 0)
 
 # ==========================
-# Ultra-Fast Video Processing
+# Optimized Video Processing
 # ==========================
 def process_video(video_data):
     try:
-        start_time = time.time()
         # Create temp video file
         video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
         with open(video_path, "wb") as f:
@@ -256,7 +238,6 @@ def process_video(video_data):
         logger.info(f"Video saved: {video_path}")
 
         # Open video file
-        start_read = time.time()
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise ValueError("Could not open video file")
@@ -265,163 +246,107 @@ def process_video(video_data):
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS)
         if fps <= 0:
-            fps = 30
+            fps = 30  # Default assumption if FPS not available
         duration = total_frames / fps
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
         logger.info(f"Video properties: {duration:.2f}s, {total_frames} frames, {fps:.1f} FPS, {width}x{height}")
 
-        # Check if processing will exceed time limit
-        if duration > CONFIG["MAX_PROCESSING_TIME"]:
-            logger.warning(f"Video duration {duration:.2f}s exceeds max processing time {CONFIG['MAX_PROCESSING_TIME']}s")
-            cap.release()
-            os.remove(video_path)
-            yield "Video duration too long. Please upload a shorter video.", "", "", "", ""
-            return
-
-        # Estimate processing feasibility
-        estimated_frames = total_frames // CONFIG["FRAME_SAMPLE_RATE"]
-        if estimated_frames * 0.1 > CONFIG["MAX_PROCESSING_TIME"]:  # Rough estimate: 0.1s per frame
-            logger.warning(f"Too many frames ({estimated_frames}) to process within {CONFIG['MAX_PROCESSING_TIME']}s")
-            cap.release()
-            os.remove(video_path)
-            yield "Video has too many frames to process within 30 seconds.", "", "", "", ""
-            return
-
-        # Read frames with sampling
-        frame_batches = []
-        frame_indices_batches = []
-        current_batch = []
-        current_indices = []
-        frame_count = 0
-        sampled_frame_count = 0
-
-        while True:
+        # Read all frames upfront
+        all_frames = []
+        all_indices = []
+        for frame_idx in range(total_frames):
             ret, frame = cap.read()
             if not ret:
                 break
-            if frame_count % CONFIG["FRAME_SAMPLE_RATE"] == 0:
-                frame = resize_frame(frame, CONFIG["MAX_FRAME_WIDTH"])
-                current_batch.append(frame)
-                current_indices.append(frame_count)
-                sampled_frame_count += 1
-            frame_count += 1
-
-            if len(current_batch) >= CONFIG["CHUNK_SIZE"]:
-                frame_batches.append(current_batch)
-                frame_indices_batches.append(current_indices)
-                current_batch = []
-                current_indices = []
-
-        if current_batch:
-            frame_batches.append(current_batch)
-            frame_indices_batches.append(current_indices)
-
+            all_frames.append(frame)
+            all_indices.append(frame_idx)
         cap.release()
-        logger.info(f"Frame reading time: {time.time() - start_read:.2f}s")
-        logger.info(f"Total frames: {frame_count}, Sampled frames: {sampled_frame_count}")
 
-        # Process frames in parallel
+        # Process frames in parallel batches
+        workers = []
         violations = []
         helmet_violations = {}
         snapshots = []
-        last_progress_time = start_time
+        start_time = time.time()
+
+        # Split frames into batches
+        frame_batches = [all_frames[i:i + CONFIG["CHUNK_SIZE"]] for i in range(0, len(all_frames), CONFIG["CHUNK_SIZE"])]
+        frame_indices_batches = [all_indices[i:i + CONFIG["CHUNK_SIZE"]] for i in range(0, len(all_indices), CONFIG["CHUNK_SIZE"])]
 
-        start_parallel = time.time()
+        # Process batches in parallel
         with Pool(processes=CONFIG["PARALLEL_WORKERS"]) as pool:
             process_func = partial(process_frame_batch, fps=fps)
             results = pool.starmap(process_func, zip(frame_batches, frame_indices_batches))
-            
-            # Flatten and sort results
-            all_detections = []
-            for batch_result in results:
-                all_detections.extend(batch_result)
-            all_detections.sort(key=lambda x: x[0])
-
-        logger.info(f"Parallel processing time: {time.time() - start_parallel:.2f}s")
-
-        # Worker tracking
-        start_tracking = time.time()
-        workers = []
-        for frame_idx, detections in all_detections:
-            current_time = frame_idx / fps
-
-            # Update progress every second
-            if time.time() - last_progress_time > 1.0:
-                progress = (frame_idx / frame_count) * 100
-                yield f"Processing video... {progress:.1f}% complete (Frame {frame_idx}/{frame_count})", "", "", "", ""
-                last_progress_time = time.time()
-
-            # Early termination if time limit approached
-            if time.time() - start_time > CONFIG["MAX_PROCESSING_TIME"] - 2:
-                logger.warning("Approaching max processing time, terminating early")
-                break
-
-            for detection in detections:
-                worker_id = None
-                max_iou = 0
-                for idx, worker in enumerate(workers):
-                    iou = calculate_iou(detection["bounding_box"], worker["bbox"])
-                    if iou > max_iou and iou > CONFIG["IOU_THRESHOLD"]:
-                        max_iou = iou
-                        worker_id = worker["id"]
-                        workers[idx]["bbox"] = detection["bounding_box"]
-                        workers[idx]["last_seen"] = current_time
-
-                if worker_id is None:
-                    worker_id = len(workers) + 1
-                    workers.append({
-                        "id": worker_id,
-                        "bbox": detection["bounding_box"],
-                        "first_seen": current_time,
-                        "last_seen": current_time
-                    })
 
-                detection["worker_id"] = worker_id
+        # Flatten results and track workers
+        for batch_result in results:
+            for frame_idx, detections in batch_result:
+                current_time = frame_idx / fps
+
+                for detection in detections:
+                    # Worker tracking
+                    worker_id = None
+                    max_iou = 0
+                    for idx, worker in enumerate(workers):
+                        iou = calculate_iou(detection["bounding_box"], worker["bbox"])
+                        if iou > max_iou and iou > 0.4:  # IOU threshold
+                            max_iou = iou
+                            worker_id = worker["id"]
+                            workers[idx]["bbox"] = detection["bounding_box"]
+                            workers[idx]["last_seen"] = current_time
+
+                    if worker_id is None:
+                        worker_id = len(workers) + 1
+                        workers.append({
+                            "id": worker_id,
+                            "bbox": detection["bounding_box"],
+                            "first_seen": current_time,
+                            "last_seen": current_time
+                        })
 
-                if detection["violation"] == "no_helmet":
-                    if worker_id not in helmet_violations:
-                        helmet_violations[worker_id] = []
-                    helmet_violations[worker_id].append(detection)
-                else:
-                    violations.append(detection)
+                    detection["worker_id"] = worker_id
 
-            workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
+                    # Track helmet violations separately
+                    if detection["violation"] == "no_helmet":
+                        if worker_id not in helmet_violations:
+                            helmet_violations[worker_id] = []
+                        helmet_violations[worker_id].append(detection)
+                    else:
+                        violations.append(detection)
 
-        logger.info(f"Worker tracking time: {time.time() - start_tracking:.2f}s")
+                # Remove inactive workers
+                workers = [w for w in workers if current_time - w["last_seen"] < CONFIG["WORKER_TRACKING_DURATION"]]
 
-        # Process helmet violations
-        start_snapshot = time.time()
+        # Confirm helmet violations (require multiple detections)
         for worker_id, detections in helmet_violations.items():
             if len(detections) >= CONFIG["MIN_VIOLATION_FRAMES"]:
                 best_detection = max(detections, key=lambda x: x["confidence"])
-                if best_detection["confidence"] >= CONFIG["HELMET_CONFIDENCE_THRESHOLD"]:
-                    violations.append(best_detection)
-                    
-                    # Capture snapshot
-                    cap = cv2.VideoCapture(video_path)
-                    cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
-                    ret, snapshot_frame = cap.read()
-                    if ret:
-                        snapshot_frame = draw_detections(snapshot_frame, [best_detection])
-                        snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
-                        cv2.imwrite(snapshot_path, snapshot_frame)
-                        snapshots.append({
-                            "violation": "no_helmet",
-                            "frame": best_detection["frame"],
-                            "snapshot_path": snapshot_path,
-                            "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
-                        })
-                    cap.release()
-
-        logger.info(f"Snapshot generation time: {time.time() - start_snapshot:.2f}s")
+                violations.append(best_detection)
+
+                # Capture snapshot
+                cap = cv2.VideoCapture(video_path)
+                cap.set(cv2.CAP_PROP_POS_FRAMES, best_detection["frame"])
+                ret, snapshot_frame = cap.read()
+                if ret:
+                    snapshot_frame = draw_detections(snapshot_frame, [best_detection])
+                    snapshot_filename = f"no_helmet_{best_detection['frame']}.jpg"
+                    snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                    cv2.imwrite(snapshot_path, snapshot_frame)
+                    snapshots.append({
+                        "violation": "no_helmet",
+                        "frame": best_detection["frame"],
+                        "snapshot_path": snapshot_path,
+                        "snapshot_base64": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}"
+                    })
+                cap.release()
 
         os.remove(video_path)
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s. {len(violations)} violations found.")
 
+        # Generate results
         if not violations:
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
@@ -429,6 +354,7 @@ def process_video(video_data):
         score = calculate_safety_score(violations)
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
 
+        # Generate violation table
         violation_table = "| Violation              | Timestamp (s) | Confidence | Worker ID |\n"
         violation_table += "|------------------------|---------------|------------|-----------|\n"
         for v in sorted(violations, key=lambda x: x["timestamp"]):
@@ -436,12 +362,13 @@ def process_video(video_data):
             row = f"| {display_name:<22} | {v.get('timestamp', 0.0):.2f}        | {v.get('confidence', 0.0):.2f}      | {v.get('worker_id', 'N/A')} |\n"
             violation_table += row
 
+        # Generate snapshots text
         snapshots_text = "\n".join(
             f"- Snapshot for {CONFIG['DISPLAY_NAMES'].get(s['violation'], 'Unknown')} at frame {s['frame']}: ![]({s['snapshot_base64']})"
             for s in snapshots
         ) if snapshots else "No snapshots captured."
 
-        start_salesforce = time.time()
+        # Push to Salesforce
         try:
             sf = connect_to_salesforce()
             record_data = {
@@ -458,8 +385,6 @@ def process_video(video_data):
         except Exception as e:
             logger.error(f"Salesforce integration failed: {e}")
             record_id = "N/A (Salesforce error)"
-        
-        logger.info(f"Salesforce integration time: {time.time() - start_salesforce:.2f}s")
 
         yield (
             violation_table,
@@ -533,11 +458,11 @@ interface = gr.Interface(
         gr.Textbox(label="Salesforce Record ID"),
         gr.Textbox(label="Violation Details URL")
     ],
-    title="Ultra-Fast Safety Violation Analyzer",
-    description="Upload site videos to detect safety violations in under 30 seconds (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use).",
+    title="Worksite Safety Violation Analyzer",
+    description="Upload site videos to detect safety violations (No Helmet, No Harness, Unsafe Posture, Unsafe Zone, Improper Tool Use).",
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    logger.info("Launching Ultra-Fast Safety Analyzer App...")
+    logger.info("Launching Safety Analyzer App...")
     interface.launch()