Update app.py

app.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
+from contextlib import asynccontextmanager
 import torch
 from ultralytics import YOLO
 import cv2
@@ -21,8 +22,23 @@ 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")
 
+# --- Lifespan Handler ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup logic
+    logger.info("FastAPI application starting up")
+    try:
+        model = YOLOv8Model()
+        logger.info("Startup: YOLOv8 model initialized")
+    except Exception as e:
+        logger.error(f"Startup error: {e}")
+        raise
+    yield
+    # Shutdown logic (if needed)
+    logger.info("FastAPI application shutting down")
+
 # --- Initialize FastAPI app ---
-app = FastAPI()
+app = FastAPI(lifespan=lifespan)
 
 # --- YOLOv8 Model ---
 class YOLOv8Model:
@@ -170,15 +186,4 @@ async def upload_image(file: UploadFile = File(...)):
 
 @app.get("/health_check/")
 async def health_check():
-    return {"status": "Running smoothly"}
-
-@app.on_event("startup")
-async def startup_event():
-    logger.info("FastAPI application starting up")
-    # Initialize any resources (e.g., check RTSP connection, model load)
-    try:
-        model = YOLOv8Model()
-        logger.info("Startup: YOLOv8 model initialized")
-    except Exception as e:
-        logger.error(f"Startup error: {e}")
-        raise
+    return {"status": "Running smoothly"}