app.py

"""
Blueprint Door and Window Detection - API Server
"""

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from PIL import Image
import io
import os

app = FastAPI(title="Blueprint Door Window Detector")

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize model variable
model = None

def load_model():
    """Load the trained YOLOv8 model"""
    global model
    model_path = 'models/best.pt'
    
    if os.path.exists(model_path):
        print(f"📦 Loading trained model from: {model_path}")
        model = YOLO(model_path)
        print("✅ Model loaded successfully!")
    else:
        print("⚠️  Trained model not found. Using pre-trained YOLOv8n for testing.")
        model = YOLO('yolov8n.pt')
        print("🔄 Using default model - train your custom model first!")

@app.get("/")
async def home():
    """Health check endpoint"""
    return {
        "message": "Blueprint Door and Window Detection API",
        "status": "running",
        "version": "1.0.0",
        "model_loaded": model is not None
    }

@app.post("/detect")
async def detect_objects(file: UploadFile = File(...)):
    """Detection endpoint that accepts image and returns bounding boxes"""
    if model is None:
        load_model()
    
    if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        raise HTTPException(status_code=400, detail="Invalid file format. Only PNG/JPG allowed")
    
    try:
        # Read and process image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        # Run inference
        results = model(image)[0]
        
        # Process detections
        detections = []
        for box in results.boxes:
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            confidence = float(box.conf[0])
            class_id = int(box.cls[0])
            label = 'door' if class_id == 0 else 'window'
            
            detections.append({
                'label': label,
                'confidence': round(confidence, 3),
                'bbox': [round(x1, 1), round(y1, 1), 
                        round(x2-x1, 1), round(y2-y1, 1)]  # [x, y, width, height]
            })
        
        return {
            'detections': detections,
            'image_size': {'width': image.width, 'height': image.height},
            'processing_time_ms': round(results.speed['inference'], 1)
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

@app.on_event("startup")
async def startup_event():
    """Load the model when the server starts"""
    load_model()