Update app.py

app.py

@@ -21,16 +21,9 @@ from multiprocessing import Pool, cpu_count
 from functools import partial
 import tempfile
 import shutil
+import tenacity
 
 # ========================== # Configuration and Setup # ==========================
-# Use a temporary directory for storage to avoid file system issues on Hugging Face Spaces
-TEMP_DIR = tempfile.mkdtemp(prefix="Ultralytics_")
-os.environ['YOLO_CONFIG_DIR'] = TEMP_DIR
-
-# Ensure output directory exists within temp directory
-OUTPUT_DIR = os.path.join(TEMP_DIR, "output")
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
 # Configure logging for better debugging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
@@ -193,7 +186,6 @@ class BYTETracker:
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
-    "OUTPUT_DIR": OUTPUT_DIR,
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
@@ -329,11 +321,11 @@ def calculate_safety_score(violations):
     score = max(0, 100 - total_penalty)
     return score
 
-def generate_violation_pdf(violations, score):
+def generate_violation_pdf(violations, score, output_dir):
     """Generate a PDF report for the detected violations"""
     try:
         pdf_filename = f"violations_{int(time.time())}.pdf"
-        pdf_path = os.path.join(CONFIG["OUTPUT_DIR"], pdf_filename)
+        pdf_path = os.path.join(output_dir, pdf_filename)
         pdf_file = BytesIO()
         c = canvas.Canvas(pdf_file, pagesize=letter)
         
@@ -503,23 +495,50 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
         logger.error(f"Salesforce record creation failed: {e}")
         return "N/A", "Salesforce integration failed."
 
-def process_video(video_data):
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_fixed(1),
+    retry=tenacity.retry_if_exception_type((IOError, OSError)),
+    before_sleep=lambda retry_state: logger.info(f"Retrying file access (attempt {retry_state.attempt_number}/3)...")
+)
+def verify_and_open_video(video_path):
+    """Verify file existence and readability, then open with cv2.VideoCapture"""
+    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)  # Test read access
+    
+    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
+
+def process_video(video_data, temp_dir):
     """Process video to detect safety violations"""
-    video_path = None  # Initialize video_path for cleanup
+    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  # Set YOLO config dir for this process
+    
     try:
         # Validate video data
         if not video_data:
             raise ValueError("Empty video data provided.")
         
-        # Log the size of the video data
         logger.info(f"Received video data size: {len(video_data)} bytes")
         if len(video_data) == 0:
             raise ValueError("Video data is empty.")
 
-        # Save video to a temporary file using NamedTemporaryFile
-        with tempfile.NamedTemporaryFile(suffix=".mp4", dir=TEMP_DIR, delete=False) as temp_file:
+        # Save video to a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
             temp_file.write(video_data)
-            temp_file.flush()  # Ensure all data is written to disk
+            temp_file.flush()
             video_path = temp_file.name
         logger.info(f"Video saved to temporary file: {video_path}")
 
@@ -531,18 +550,9 @@ def process_video(video_data):
             raise ValueError(f"Temporary video file is empty: {video_path}")
         logger.info(f"Temporary video file size: {file_size} bytes")
 
-        # Ensure the file is readable
-        try:
-            with open(video_path, "rb") as f:
-                f.read(1)  # Test read access
-            logger.info(f"Temporary video file is readable: {video_path}")
-        except Exception as e:
-            raise IOError(f"Cannot read temporary video file {video_path}: {str(e)}")
-
-        # Open video with OpenCV
-        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.")
+        # Open video with OpenCV (with retries)
+        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
@@ -551,7 +561,6 @@ def process_video(video_data):
         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 video is empty
         if total_frames <= 0:
             raise ValueError("Video has no frames.")
 
@@ -596,7 +605,6 @@ def process_video(video_data):
                 logger.info("No more frames to process.")
                 break
 
-            # Process batch with YOLO model
             try:
                 results = model(batch_frames, device=device, conf=0.1, verbose=False)
             except Exception as e:
@@ -659,7 +667,7 @@ def process_video(video_data):
                         detection = {
                             "worker_id": worker_id,
                             "violation": label,
-                            "confidence": round(float(conf), 2),  # Ensure confidence is a float
+                            "confidence": round(float(conf), 2),
                             "bounding_box": bbox,
                             "timestamp": current_time
                         }
@@ -678,7 +686,7 @@ def process_video(video_data):
                         )
                         
                         snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
-                        snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
+                        snapshot_path = os.path.join(output_dir, snapshot_filename)
                         
                         cv2.imwrite(
                             snapshot_path,
@@ -692,24 +700,20 @@ def process_video(video_data):
                             "timestamp": current_time,
                             "snapshot_path": snapshot_path,
                             "snapshot_url": f"{CONFIG['PUBLIC_URL_BASE']}{snapshot_filename}",
-                            "confidence": round(float(conf), 2)  # Ensure confidence is stored as float
+                            "confidence": round(float(conf), 2)
                         })
                         
                         logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
 
-        # Ensure resources are released
         cap.release()
-            
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
-        # Log the snapshots for debugging
         logger.info(f"Snapshots: {snapshots}")
 
         violations = []
         for worker_id, worker_violations in unique_violations.items():
             for label, detection_time in worker_violations.items():
-                # Find the confidence from snapshots, ensuring it's a float
                 confidence = next(
                     (float(s["confidence"]) for s in snapshots if s["worker_id"] == worker_id and s["violation"] == label),
                     0.0
@@ -722,7 +726,6 @@ def process_video(video_data):
                 }
                 violations.append(violation)
 
-        # Log the violations for debugging
         logger.info(f"Violations: {violations}")
 
         if not violations:
@@ -731,12 +734,10 @@ def process_video(video_data):
             return
 
         score = calculate_safety_score(violations)
-        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
+        pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score, output_dir)
         
-        # Push to Salesforce with fallback
         record_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        # Generate violation table with robust error handling
         violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
         violation_table += "|-----------|-----------|----------|------------|\n"
         
@@ -744,7 +745,6 @@ def process_video(video_data):
             display_name = CONFIG["DISPLAY_NAMES"].get(v.get("violation", "Unknown"), "Unknown")
             worker_id = v.get("worker_id", "Unknown")
             timestamp = v.get("timestamp", 0.0)
-            # Ensure confidence is a valid float
             try:
                 confidence = float(v.get("confidence", 0.0))
             except (ValueError, TypeError) as e:
@@ -777,44 +777,62 @@ def process_video(video_data):
         logger.error(f"Error processing video: {str(e)}", exc_info=True)
         yield f"Error processing video: {str(e)}", "", "", "", ""
     finally:
-        # Clean up the temporary video file
         if video_path and os.path.exists(video_path):
             try:
                 os.remove(video_path)
                 logger.info(f"Cleaned up temporary video file: {video_path}")
             except Exception as e:
                 logger.error(f"Failed to clean up temporary video file {video_path}: {e}")
-        # Clean up temporary directory
-        if os.path.exists(TEMP_DIR):
-            shutil.rmtree(TEMP_DIR, ignore_errors=True)
-            logger.info(f"Cleaned up temporary directory: {TEMP_DIR}")
 
 def gradio_interface(video_file):
     """Gradio interface for the video processing"""
-    if not video_file:
-        return "No file uploaded.", "", "No file uploaded.", "", ""
-        
+    temp_dir = None
+    local_video_path = None
     try:
-        # Read the video file
+        if not video_file:
+            return "No file uploaded.", "", "No file uploaded.", "", ""
+        
+        # Create a unique temporary directory for this video
+        temp_dir = tempfile.mkdtemp(prefix="Ultralytics_")
+        logger.info(f"Created temporary directory for video processing: {temp_dir}")
+
+        # Copy Gradio's video file to a local temporary file
         with open(video_file, "rb") as f:
             video_data = f.read()
-
-        # Log the size of the uploaded video file
-        logger.info(f"Uploaded video file size: {len(video_data)} bytes")
+        logger.info(f"Read Gradio video file: {video_file}, size: {len(video_data)} bytes")
+        
         if len(video_data) == 0:
             return "Uploaded video file is empty.", "", "", "", ""
 
-        # Validate FFmpeg availability
+        # Save to a local temporary file to avoid Gradio file deletion
+        with tempfile.NamedTemporaryFile(suffix=".mp4", dir=temp_dir, delete=False) as temp_file:
+            temp_file.write(video_data)
+            temp_file.flush()
+            local_video_path = temp_file.name
+        logger.info(f"Copied Gradio video to local temporary file: {local_video_path}")
+
         if not FFMPEG_AVAILABLE:
             return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", ""
 
-        # Process the video
-        for status, score, snapshots_text, record_id, details_url in process_video(video_data):
+        for status, score, snapshots_text, record_id, details_url in process_video(video_data, temp_dir):
             yield status, score, snapshots_text, record_id, details_url
             
     except Exception as e:
         logger.error(f"Error in Gradio interface: {e}", exc_info=True)
         yield f"Error: {str(e)}", "", "Error in processing.", "", ""
+    finally:
+        # Clean up the local temporary video file
+        if local_video_path and os.path.exists(local_video_path):
+            try:
+                os.remove(local_video_path)
+                logger.info(f"Cleaned up local temporary video file: {local_video_path}")
+            except Exception as e:
+                logger.error(f"Failed to clean up local temporary video file {local_video_path}: {e}")
+        
+        # Clean up the temporary directory
+        if temp_dir and os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            logger.info(f"Cleaned up temporary directory: {temp_dir}")
 
 # ========================== # Gradio Interface # ==========================
 interface = gr.Interface(