Update app.py

app.py

@@ -1,31 +1,34 @@
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi.responses import JSONResponse, RedirectResponse
 from ultralytics import YOLO
 import cv2
 import numpy as np
 
-app = FastAPI()
+# Initialize FastAPI app
+app = FastAPI(
+    title="Door & Window Detection API",
+    description="Upload an image (.jpg or .png) to detect doors and windows.",
+    version="1.0.0"
+)
 
-# Load the custom-trained YOLOv8 model
-model = YOLO("yolov8m_custom.pt")
+# Load the YOLOv8 model once at startup
+model = YOLO("yolov8m_custom.pt")  # Replace with yolov8n_custom.pt if needed
 
-# 🔹 Add root route to prevent 404 on /
-@app.get("/", response_class=HTMLResponse)
-async def root():
-    return """
-    <h2>Door & Window Detection API</h2>
-    <p>Use <code>/predict</code> endpoint to POST a .jpg or .png image and get detections.</p>
-    """
+@app.get("/")
+async def redirect_to_docs():
+    # Redirect to Swagger UI for easy testing
+    return RedirectResponse(url="/docs")
 
-# 🔹 Change to /predict to match typical inference APIs
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
+    # Validate file extension
     if not file.filename.lower().endswith((".jpg", ".png")):
         return JSONResponse(
             content={"error": "Only JPG or PNG images are allowed"},
             status_code=400
         )
 
+    # Read file and convert to OpenCV image
     image_bytes = await file.read()
     np_array = np.frombuffer(image_bytes, np.uint8)
     image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
@@ -33,9 +36,10 @@ async def predict(file: UploadFile = File(...)):
     if image is None:
         return JSONResponse(content={"error": "Could not decode image"}, status_code=400)
 
-    # Run YOLOv8 detection
-    results = model.predict(image)
+    # Run inference with YOLOv8
+    results = model.predict(image, verbose=False)
 
+    # Process results
     detections = []
     for result in results:
         for box in result.boxes: