Update app.py

app.py

@@ -19,14 +19,33 @@ from retrying import retry
 import uuid
 from multiprocessing import Pool, cpu_count
 from functools import partial
+import tempfile
+import shutil
 
 # ========================== # Configuration and Setup # ==========================
-os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
-os.makedirs('/tmp/Ultralytics', exist_ok=True)
+# Use a temporary directory for storage
+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)
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
+# Check for FFmpeg availability
+def check_ffmpeg():
+    try:
+        subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
+        logger.info("FFmpeg is available.")
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        logger.error("FFmpeg is not installed or not found in PATH. Video processing may fail.")
+        return False
+
+FFMPEG_AVAILABLE = check_ffmpeg()
+
 # ========================== # ByteTrack Implementation # ==========================
 class BYTETracker:
     def __init__(self, track_thresh=0.3, track_buffer=30, match_thresh=0.7, frame_rate=30):
@@ -89,7 +108,6 @@ class BYTETracker:
                 })
             else:
                 # Create new track
-                # Check if this detection might be the same worker from a different angle
                 same_worker = False
                 for worker_id, last_pos in self.last_positions.items():
                     if self._is_same_worker([x, y], last_pos):
@@ -174,7 +192,7 @@ class BYTETracker:
 CONFIG = {
     "MODEL_PATH": "yolov8_safety.pt",
     "FALLBACK_MODEL": "yolov8n.pt",
-    "OUTPUT_DIR": "static/output",
+    "OUTPUT_DIR": OUTPUT_DIR,
     "VIOLATION_LABELS": {
         0: "no_helmet",
         1: "no_harness",
@@ -211,17 +229,17 @@ CONFIG = {
         "improper_tool_use": 0.3
     },
     "MIN_VIOLATION_FRAMES": 1,
-    "VIOLATION_COOLDOWN": 30.0,  # Increased cooldown period
+    "VIOLATION_COOLDOWN": 30.0,
     "WORKER_TRACKING_DURATION": 5.0,
     "MAX_PROCESSING_TIME": 60,
-    "FRAME_SKIP": 2,  # Skip more frames for faster processing
-    "BATCH_SIZE": 16,
+    "FRAME_SKIP": 2,
+    "BATCH_SIZE": 8,  # Reduced batch size to lower memory usage
     "PARALLEL_WORKERS": max(1, cpu_count() - 1),
     "TRACK_BUFFER": 30,
     "TRACK_THRESH": 0.3,
     "MATCH_THRESH": 0.7,
-    "SNAPSHOT_QUALITY": 95,  # Higher quality for better visibility
-    "MAX_WORKER_DISTANCE": 100  # Maximum pixel distance to consider same worker
+    "SNAPSHOT_QUALITY": 95,
+    "MAX_WORKER_DISTANCE": 100
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -255,7 +273,7 @@ def preprocess_frame(frame):
     return frame
 
 def draw_detections(frame, detections):
-    """Draw bounding boxes and labels on detection frame with improved visibility"""
+    """Draw bounding boxes and labels on detection frame"""
     result_frame = frame.copy()
     
     for det in detections:
@@ -271,16 +289,13 @@ def draw_detections(frame, detections):
         
         color = CONFIG["CLASS_COLORS"].get(label, (0, 0, 255))
         
-        # Draw thicker rectangle with border
         cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 3)
         
-        # Add black background behind text
         display_text = f"{CONFIG['DISPLAY_NAMES'].get(label, label)} (Worker {worker_id})"
         text_size = cv2.getTextSize(display_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
         cv2.rectangle(result_frame, (x1, y1-text_size[1]-10), (x1+text_size[0]+10, y1), (0, 0, 0), -1)
         cv2.putText(result_frame, display_text, (x1+5, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
         
-        # Add confidence score
         conf_text = f"Conf: {confidence:.2f}"
         cv2.putText(result_frame, conf_text, (x1+5, y2+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
         
@@ -296,7 +311,6 @@ def calculate_safety_score(violations):
         "improper_tool_use": 25
     }
     
-    # Count unique violation types per worker
     worker_violations = {}
     for v in violations:
         worker_id = v.get("worker_id", "Unknown")
@@ -306,7 +320,6 @@ def calculate_safety_score(violations):
             worker_violations[worker_id] = set()
         worker_violations[worker_id].add(violation_type)
     
-    # Calculate total penalty
     total_penalty = 0
     for worker_violations_set in worker_violations.values():
         worker_penalty = sum(penalties.get(v, 0) for v in worker_violations_set)
@@ -323,26 +336,21 @@ def generate_violation_pdf(violations, score):
         pdf_file = BytesIO()
         c = canvas.Canvas(pdf_file, pagesize=letter)
         
-        # Title
         c.setFont("Helvetica-Bold", 16)
         c.drawString(1 * inch, 10 * inch, "Worksite Safety Violation Report")
         
-        # Basic Information
         c.setFont("Helvetica", 12)
         c.drawString(1 * inch, 9.5 * inch, f"Date: {time.strftime('%Y-%m-%d')}")
         c.drawString(1 * inch, 9.2 * inch, f"Time: {time.strftime('%H:%M:%S')}")
         
-        # Safety Score
         c.setFont("Helvetica-Bold", 14)
         c.drawString(1 * inch, 8.7 * inch, f"Safety Compliance Score: {score}%")
 
-        # Violation Summary
         y_position = 8.2 * inch
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Summary:")
         y_position -= 0.3 * inch
         
-        # Group violations by worker
         worker_violations = {}
         for v in violations:
             worker_id = v.get("worker_id", "Unknown")
@@ -361,7 +369,6 @@ def generate_violation_pdf(violations, score):
             c.drawString(1 * inch, y_position, f"{key}: {value}")
             y_position -= 0.25 * inch
 
-        # Detailed Violations by Worker
         y_position -= 0.5 * inch
         c.setFont("Helvetica-Bold", 12)
         c.drawString(1 * inch, y_position, "Violations by Worker:")
@@ -389,7 +396,6 @@ def generate_violation_pdf(violations, score):
         c.save()
         pdf_file.seek(0)
 
-        # Save PDF file
         with open(pdf_path, "wb") as f:
             f.write(pdf_file.getvalue())
             
@@ -445,7 +451,6 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
     try:
         sf = connect_to_salesforce()
         
-        # Format violations for Salesforce
         violations_text = ""
         for v in violations:
             display_name = CONFIG['DISPLAY_NAMES'].get(v.get('violation', 'Unknown'), 'Unknown')
@@ -494,24 +499,26 @@ def push_report_to_salesforce(violations, score, pdf_path, pdf_file):
 
         return record_id, pdf_url
     except Exception as e:
-        logger.error(f"Salesforce record creation failed: {e}", exc_info=True)
-        return None, ""
+        logger.error(f"Salesforce record creation failed: {e}")
+        return "N/A", "Salesforce integration failed."
 
 def process_video(video_data):
     """Process video to detect safety violations"""
     try:
-        os.makedirs(CONFIG["OUTPUT_DIR"], exist_ok=True)
-        logger.info(f"Output directory ensured: {CONFIG['OUTPUT_DIR']}")
+        # Validate video data
+        if not video_data:
+            raise ValueError("Empty video data provided.")
 
-        video_path = os.path.join(CONFIG["OUTPUT_DIR"], f"temp_{int(time.time())}.mp4")
-        with open(video_path, "wb") as f:
+        # Save video to a temporary file
+        video_fd, video_path = tempfile.mkstemp(suffix=".mp4", dir=TEMP_DIR)
+        with os.fdopen(video_fd, "wb") as f:
             f.write(video_data)
         logger.info(f"Video saved: {video_path}")
 
+        # Open video with OpenCV
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
-            os.remove(video_path)
-            raise ValueError("Could not open video file")
+            raise ValueError("Could not open video file. Ensure the video format is supported (e.g., MP4) and FFmpeg is installed.")
 
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         fps = cap.get(cv2.CAP_PROP_FPS) or 30
@@ -520,6 +527,10 @@ 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.")
+
         tracker = BYTETracker(
             track_thresh=CONFIG["TRACK_THRESH"],
             track_buffer=CONFIG["TRACK_BUFFER"],
@@ -527,8 +538,7 @@ def process_video(video_data):
             frame_rate=fps
         )
 
-        # Track unique violations by worker ID
-        unique_violations = {}  # {worker_id: {violation_type: first_detection_time}}
+        unique_violations = {}
         snapshots = []
         start_time = time.time()
         frame_skip = CONFIG["FRAME_SKIP"]
@@ -545,11 +555,11 @@ def process_video(video_data):
                 
                 ret, frame = cap.read()
                 if not ret:
+                    logger.warning(f"Failed to read frame {frame_idx}. Skipping.")
                     break
                 
                 frame = preprocess_frame(frame)
                 
-                # Skip frames if needed
                 for _ in range(frame_skip - 1):
                     if not cap.grab():
                         break
@@ -559,15 +569,19 @@ def process_video(video_data):
                 processed_frames += 1
 
             if not batch_frames:
+                logger.info("No more frames to process.")
                 break
 
             # Process batch with YOLO model
-            results = model(batch_frames, device=device, conf=0.1, verbose=False)
-            
+            try:
+                results = model(batch_frames, device=device, conf=0.1, verbose=False)
+            except Exception as e:
+                logger.error(f"Model inference failed: {e}")
+                raise ValueError(f"Failed to process video frames with YOLO model: {str(e)}")
+
             for i, (result, frame_idx) in enumerate(zip(results, batch_indices)):
                 current_time = frame_idx / fps
                 
-                # Update progress every second
                 if time.time() - start_time > 1.0:
                     progress = (processed_frames / total_frames) * 100
                     yield f"Processing video... {progress:.1f}% complete (Frame {processed_frames}/{total_frames})", "", "", "", ""
@@ -603,7 +617,6 @@ def process_video(video_data):
                     np.array([t["cls"] for t in track_inputs])
                 )
 
-                # Process tracked objects for violations
                 for obj in tracked_objects:
                     worker_id = obj['id']
                     label = CONFIG["VIOLATION_LABELS"].get(int(obj['cls']), None)
@@ -613,16 +626,12 @@ def process_video(video_data):
                     if label is None:
                         continue
                         
-                    # Initialize worker if not seen before
                     if worker_id not in unique_violations:
                         unique_violations[worker_id] = {}
                     
-                    # Check if this violation type has been recorded for this worker
                     if label not in unique_violations[worker_id]:
-                        # This is a new violation type for this worker
                         unique_violations[worker_id][label] = current_time
                         
-                        # Create detection object
                         detection = {
                             "worker_id": worker_id,
                             "violation": label,
@@ -631,11 +640,9 @@ def process_video(video_data):
                             "timestamp": current_time
                         }
                         
-                        # Take snapshot for the new violation
                         snapshot_frame = batch_frames[i].copy()
                         snapshot_frame = draw_detections(snapshot_frame, [detection])
                         
-                        # Add timestamp to snapshot
                         cv2.putText(
                             snapshot_frame,
                             f"Time: {current_time:.2f}s",
@@ -646,7 +653,6 @@ def process_video(video_data):
                             2
                         )
                         
-                        # Save snapshot with high quality
                         snapshot_filename = f"violation_{label}_worker{worker_id}_{int(current_time*100)}.jpg"
                         snapshot_path = os.path.join(CONFIG["OUTPUT_DIR"], snapshot_filename)
                         
@@ -666,6 +672,7 @@ def process_video(video_data):
                         
                         logger.info(f"Captured snapshot for {label} violation by worker {worker_id} at {current_time:.2f}s")
 
+        # Ensure resources are released
         cap.release()
         if os.path.exists(video_path):
             os.remove(video_path)
@@ -673,14 +680,14 @@ def process_video(video_data):
         processing_time = time.time() - start_time
         logger.info(f"Processing complete in {processing_time:.2f}s")
 
-        # Convert tracked violations to final violation list
         violations = []
         for worker_id, worker_violations in unique_violations.items():
             for label, detection_time in worker_violations.items():
                 violation = {
                     "worker_id": worker_id,
                     "violation": label,
-                    "timestamp": detection_time
+                    "timestamp": detection_time,
+                    "confidence": next((s["confidence"] for s in snapshots if s["worker_id"] == worker_id and s["violation"] == label), 0.0)
                 }
                 violations.append(violation)
 
@@ -689,16 +696,12 @@ def process_video(video_data):
             yield "No violations detected in the video.", "Safety Score: 100%", "No snapshots captured.", "N/A", "N/A"
             return
 
-        # Calculate safety score
         score = calculate_safety_score(violations)
-        
-        # Generate PDF report
         pdf_path, pdf_url, pdf_file = generate_violation_pdf(violations, score)
         
-        # Push report to Salesforce
-        report_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
+        # Push to Salesforce with fallback
+        record_id, final_pdf_url = push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
-        # Format violations table for display
         violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n"
         violation_table += "|-----------|-----------|----------|------------|\n"
         
@@ -710,7 +713,6 @@ def process_video(video_data):
             
             violation_table += f"| {display_name} | {worker_id} | {timestamp:.2f} | {confidence:.2f} |\n"
 
-        # Format snapshots for display
         snapshots_text = ""
         for s in snapshots:
             display_name = CONFIG["DISPLAY_NAMES"].get(s["violation"], "Unknown")
@@ -727,15 +729,19 @@ def process_video(video_data):
             violation_table,
             f"Safety Score: {score}%",
             snapshots_text,
-            f"Salesforce Record ID: {report_id or 'N/A'}",
-            final_pdf_url or "N/A"
+            f"Salesforce Record ID: {record_id}",
+            final_pdf_url
         )
 
     except Exception as e:
-        logger.error(f"Error processing video: {e}", exc_info=True)
+        logger.error(f"Error processing video: {str(e)}", exc_info=True)
         if 'video_path' in locals() and os.path.exists(video_path):
             os.remove(video_path)
-        yield f"Error processing video: {e}", "", "", "", ""
+        yield f"Error processing video: {str(e)}", "", "", "", ""
+    finally:
+        # Clean up temporary directory
+        if os.path.exists(TEMP_DIR):
+            shutil.rmtree(TEMP_DIR, ignore_errors=True)
 
 def gradio_interface(video_file):
     """Gradio interface for the video processing"""
@@ -746,6 +752,10 @@ def gradio_interface(video_file):
         with open(video_file, "rb") as f:
             video_data = f.read()
 
+        # Validate FFmpeg availability
+        if not FFMPEG_AVAILABLE:
+            return "FFmpeg is not available in the environment. Please install FFmpeg to process videos.", "", "", "", ""
+
         for status, score, snapshots_text, record_id, details_url in process_video(video_data):
             yield status, score, snapshots_text, record_id, details_url