Update app.py

app.py

@@ -12,8 +12,10 @@ from ultralytics import YOLO
 from tracker import BYTETracker
 from utils import (
     preprocess_frame, draw_detections, calculate_safety_score,
-    generate_violation_pdf, push_report_to_salesforce, verify_and_open_video,
-    blur_faces
+    generateಸ
+
+System: generate_violation_pdf, push_report_to_salesforce, verify_and_open_video,
+    blur_faces, clean_output_directory
 )
 from config import CONFIG, check_ffmpeg
 
@@ -29,13 +31,39 @@ FFMPEG_AVAILABLE = check_ffmpeg()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
-# Load YOLO model
-def load_model():
+def setup_static_folder():
+    """Ensure static folder and model weights are available."""
+    static_dir = "static"
+    output_dir = os.path.join(static_dir, "output")
+    
+    os.makedirs(static_dir, exist_ok=True)
+    os.makedirs(output_dir, exist_ok=True)
+    logger.info(f"Static directory ensured: {static_dir}")
+    logger.info(f"Output directory ensured: {output_dir}")
+
+    model_path = CONFIG["MODEL_PATH"]
+    fallback_model = CONFIG["FALLBACK_MODEL"]
+    
+    if not os.path.isfile(model_path):
+        logger.warning(f"Custom model {model_path} not found. Falling back to {fallback_model}.")
+        if not os.path.isfile(fallback_model):
+            logger.info(f"Downloading fallback model: {fallback_model}")
+            try:
+                torch.hub.download_url_to_file(
+                    'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt',
+                    fallback_model
+                )
+                logger.info(f"Downloaded {fallback_model}")
+            except Exception as e:
+                logger.error(f"Failed to download {fallback_model}: {e}")
+                raise
+    else:
+        logger.info(f"Using custom model: {model_path}")
+
+    return model_path if os.path.isfile(model_path) else fallback_model
+
+def load_model(model_path):
     try:
-        model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
-        if not os.path.isfile(model_path):
-            logger.info(f"Downloading fallback model: {model_path}")
-            torch.hub.download_url_to_file('https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt', model_path)
         model = YOLO(model_path).to(device)
         if device.type == "cuda":
             model.model.half()
@@ -45,12 +73,13 @@ def load_model():
         logger.error(f"Failed to load model: {e}")
         raise
 
-model = load_model()
+model_path = setup_static_folder()
+model = load_model(model_path)
 
 async def process_video(video_data, temp_dir, progress=gr.Progress()):
+    clean_output_directory()
     output_dir = os.path.join(temp_dir, "output")
     os.makedirs(output_dir, exist_ok=True)
-    os.environ['YOLO_CONFIG_DIR'] = temp_dir
     video_path = None
 
     try:
@@ -127,7 +156,8 @@ async def process_video(video_data, temp_dir, progress=gr.Progress()):
                     tracked_objects = tracker.update(
                         np.array([t["bbox"] for t in track_inputs]),
                         np.array([t["conf"] for t in track_inputs]),
-                        np.array([t["cls"] for t in track_inputs])
+                        np.array([t["cls"] for t in track_inputs]),
+                        current_time
                     )
                     logger.info(f"Frame {frame_idx}: Detected {len(tracked_objects)} workers")
 
@@ -149,12 +179,12 @@ async def process_video(video_data, temp_dir, progress=gr.Progress()):
                                 "violation": label,
                                 "timestamp": current_time,
                                 "confidence": round(obj['score'], 2),
-                                "frame_idx": frame_idx
+                                "frame_idx": frame_idx,
+                                "bounding_box": obj['bbox']
                             })
 
         cap.release()
 
-        # Capture snapshots with face blurring
         cap = cv2.VideoCapture(video_path)
         for violation in violations:
             frame_idx = violation["frame_idx"]
@@ -163,16 +193,16 @@ async def process_video(video_data, temp_dir, progress=gr.Progress()):
             if not ret:
                 continue
             frame = preprocess_frame(frame)
-            frame = blur_faces(frame)  # Add face blurring for privacy
+            frame = blur_faces(frame)
             snapshot_frame = draw_detections(frame, [{
                 "worker_id": violation["worker_id"],
                 "violation": violation["violation"],
                 "confidence": violation["confidence"],
-                "bounding_box": violation.get("bounding_box", [0, 0, 0, 0]),
+                "bounding_box": violation["bounding_box"],
                 "timestamp": violation["timestamp"]
             }])
             snapshot_filename = f"violation_{violation['violation']}_worker{violation['worker_id']}_{int(violation['timestamp']*100)}.jpg"
-            snapshot_path = os.path.join(output_dir, snapshot_filename)
+            snapshot_path = os.path.join("static/output", snapshot_filename)
             cv2.imwrite(snapshot_path, snapshot_frame, [cv2.IMWRITE_JPEG_QUALITY, CONFIG["SNAPSHOT_QUALITY"]])
             snapshots.append({
                 "violation": violation["violation"],
@@ -186,7 +216,7 @@ async def process_video(video_data, temp_dir, progress=gr.Progress()):
         cap.release()
 
         score = calculate_safety_score(violations)
-        pdf_path, pdf_url, pdf_file = await generate_violation_pdf(violations, score, output_dir)
+        pdf_path, pdf_url, pdf_file = await generate_violation_pdf(violations, score, "static/output")
         record_id, final_pdf_url = await push_report_to_salesforce(violations, score, pdf_path, pdf_file)
 
         violation_table = "| Violation | Worker ID | Time (s) | Confidence |\n|-----------|-----------|----------|------------|\n"
@@ -221,7 +251,7 @@ async def gradio_interface(video_file=None, stream_url=None):
         if video_file:
             with open(video_file, "rb") as f:
                 video_data = f.read()
-            for result in await process_video(video_data, temp_dir):
+            async for result in process_video(video_data, temp_dir):
                 yield result
         elif stream_url:
             cap = cv2.VideoCapture(stream_url)
@@ -242,12 +272,11 @@ async def gradio_interface(video_file=None, stream_url=None):
             writer.release()
             with open(temp_file, "rb") as f:
                 video_data = f.read()
-            for result in await process_video(video_data, temp_dir):
+            async for result in process_video(video_data, temp_dir):
                 yield result
     finally:
         shutil.rmtree(temp_dir, ignore_errors=True)
 
-# Gradio Interface
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=[