Spaces:
Sleeping
Sleeping
File size: 2,845 Bytes
14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 56bf267 14969c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
"""
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()
|