Update app.py

app.py

@@ -2,21 +2,6 @@ import os
 import shutil
 import sys
 import warnings
-
-# Suppress all warnings
-warnings.filterwarnings("ignore")
-
-# Clean up and set YOLO_CONFIG_DIR before any Ultralytics imports
-yolo_config_dir = "/tmp/Ultralytics"
-if os.path.exists(yolo_config_dir):
-    shutil.rmtree(yolo_config_dir)
-os.makedirs(yolo_config_dir, exist_ok=True)
-os.environ["YOLO_CONFIG_DIR"] = yolo_config_dir
-
-# Redirect stdout temporarily to suppress Ultralytics initialization messages
-original_stdout = sys.stdout
-sys.stdout = open(os.devnull, 'w')
-
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from contextlib import asynccontextmanager
 import torch
@@ -31,79 +16,61 @@ from datetime import datetime
 import logging
 import pandas as pd
 
-# Restore stdout
-sys.stdout = original_stdout
+# --- Initial Configuration ---
+# Suppress all warnings
+warnings.filterwarnings("ignore")
+
+# Clean up and set YOLO_CONFIG_DIR
+yolo_config_dir = "/tmp/Ultralytics"
+if os.path.exists(yolo_config_dir):
+    shutil.rmtree(yolo_config_dir)
+os.makedirs(yolo_config_dir, exist_ok=True)
+os.environ["YOLO_CONFIG_DIR"] = yolo_config_dir
 
-# --- Configuration ---
+# --- Logging Configuration ---
+# Configure logging before any other operations
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler()
-    ]
+    handlers=[logging.StreamHandler()]
 )
+
+# Get logger instance
 logger = logging.getLogger(__name__)
 
-# Suppress Ultralytics and other unnecessary logging
+# Suppress third-party logging
 logging.getLogger("ultralytics").setLevel(logging.WARNING)
 logging.getLogger("PIL").setLevel(logging.WARNING)
 logging.getLogger("matplotlib").setLevel(logging.WARNING)
+logging.getLogger("urllib3").setLevel(logging.WARNING)
 
-# Log environment variable for debugging
+# Log environment variable once
 logger.info(f"YOLO_CONFIG_DIR set to: {os.getenv('YOLO_CONFIG_DIR')}")
 
+# --- Environment Variables ---
 RTSP_URL = os.getenv("RTSP_URL", "rtsp://localhost:8554/stream")
 SALESFORCE_URL = os.getenv("SALESFORCE_URL", "https://your_salesforce_instance_url")
 SALESFORCE_TOKEN = os.getenv("SALESFORCE_TOKEN", "your_salesforce_access_token")
 HUGGINGFACE_API_URL = os.getenv("HUGGINGFACE_API_URL", "https://api-inference.huggingface.co/models/PrashanthB461/SafetyViolationAI1")
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN", "your_huggingface_api_token")
 
-# Global model instance
+# --- Global Model Instance ---
 yolo_model = None
 
-# --- Lifespan Handler ---
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Startup logic
-    global yolo_model
-    logger.info("FastAPI application starting up")
-    try:
-        # Suppress Ultralytics initialization messages
-        with open(os.devnull, 'w') as f:
-            original_stdout = sys.stdout
-            sys.stdout = f
-            yolo_model = YOLOv8Model()
-            sys.stdout = original_stdout
-        logger.info("YOLOv8 model initialized successfully")
-    except Exception as e:
-        logger.error(f"Startup error: {e}")
-        raise
-    yield
-    # Shutdown logic
-    logger.info("FastAPI application shutting down")
-    yolo_model = None
-
-# --- Initialize FastAPI app ---
-app = FastAPI(
-    lifespan=lifespan,
-    title="Safety Violation Detection API",
-    description="API for detecting safety violations using YOLOv8",
-    version="1.0.0"
-)
-
-# --- YOLOv8 Model ---
+# --- Model Class ---
 class YOLOv8Model:
     def __init__(self, model_path='yolov8n.pt'):
         try:
             logger.info("Initializing YOLOv8 model")
-            # Ensure the config directory exists
-            os.makedirs(os.environ["YOLO_CONFIG_DIR"], exist_ok=True)
             
-            # Suppress model loading messages
-            with open(os.devnull, 'w') as f:
-                original_stdout = sys.stdout
-                sys.stdout = f
-                self.model = YOLO(model_path)  # Load YOLOv8 model
+            # Temporarily suppress stdout for model loading
+            original_stdout = sys.stdout
+            sys.stdout = open(os.devnull, 'w')
+            
+            try:
+                self.model = YOLO(model_path)
+            finally:
+                sys.stdout.close()
                 sys.stdout = original_stdout
             
             logger.info("YOLOv8 model loaded successfully")
@@ -113,14 +80,13 @@ class YOLOv8Model:
 
     def predict(self, image):
         try:
-            # Run prediction
-            results = self.model(image)  # Inference
-            return results.pandas().xyxy[0]  # Bounding boxes, class names, confidence score
+            results = self.model(image)
+            return results.pandas().xyxy[0]
         except Exception as e:
             logger.error(f"Prediction error: {e}")
             raise
 
-# --- Preprocessing RTSP Frame ---
+# --- Frame Processing Functions ---
 def preprocess_frame(frame):
     try:
         img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -130,13 +96,13 @@ def preprocess_frame(frame):
         logger.error(f"Frame preprocessing error: {e}")
         raise
 
-# --- RTSP Stream Handler ---
 def capture_rtsp_frames(rtsp_url: str):
     try:
         cap = cv2.VideoCapture(rtsp_url)
         if not cap.isOpened():
             logger.error(f"Failed to open RTSP stream: {rtsp_url}")
             raise ValueError("RTSP stream not accessible")
+        
         while cap.isOpened():
             ret, frame = cap.read()
             if ret:
@@ -150,14 +116,14 @@ def capture_rtsp_frames(rtsp_url: str):
         logger.error(f"RTSP capture error: {e}")
         raise
 
-# --- Save Violations and Snapshots ---
+# --- Violation Handling Functions ---
 def save_snapshot(frame):
     try:
         filename = f"snapshot_{int(time.time())}.jpg"
         snapshot_path = Path("/snapshots") / filename
         os.makedirs("/snapshots", exist_ok=True)
         cv2.imwrite(str(snapshot_path), frame)
-        return f"/snapshots/{filename}"  # Relative path for containerized env
+        return f"/snapshots/{filename}"
     except Exception as e:
         logger.error(f"Snapshot saving error: {e}")
         raise
@@ -176,12 +142,9 @@ def log_violation(violation_data):
         logger.error(f"Violation logging error: {e}")
         raise
 
-# --- Notification System ---
 def send_alert(violation):
     logger.info(f"Alert! {violation['violation_type']} detected. Severity: {violation['severity']}")
-    # Placeholder for email/SMS notification logic
 
-# --- Salesforce Integration ---
 def create_salesforce_violation_record(violation_data):
     try:
         salesforce_url = f"{SALESFORCE_URL}/services/data/v60.0/sobjects/Safety_Violation_Log__c/"
@@ -206,13 +169,36 @@ def create_salesforce_violation_record(violation_data):
         logger.error(f"Salesforce integration error: {e}")
         raise
 
-# --- API Routes ---
+# --- FastAPI Application ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global yolo_model
+    logger.info("FastAPI application starting up")
+    try:
+        yolo_model = YOLOv8Model()
+        logger.info("YOLOv8 model initialized successfully")
+    except Exception as e:
+        logger.error(f"Startup error: {e}")
+        raise
+    yield
+    logger.info("FastAPI application shutting down")
+    yolo_model = None
+
+app = FastAPI(
+    lifespan=lifespan,
+    title="Safety Violation Detection API",
+    description="API for detecting safety violations using YOLOv8",
+    version="1.0.0"
+)
+
+# --- API Endpoints ---
 @app.post("/detect_violation/")
 async def detect_violation():
     try:
         global yolo_model
         if yolo_model is None:
             raise HTTPException(status_code=500, detail="YOLO model not initialized")
+        
         for frame, timestamp in capture_rtsp_frames(RTSP_URL):
             frame_processed = preprocess_frame(frame)
             results = yolo_model.predict(frame_processed)
@@ -242,6 +228,7 @@ async def upload_image(file: UploadFile = File(...)):
         global yolo_model
         if yolo_model is None:
             raise HTTPException(status_code=500, detail="YOLO model not initialized")
+        
         image_data = await file.read()
         image = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
         results = yolo_model.predict(image)