Spaces:
Sleeping
Sleeping
Fix FastAPI startup configuration
Browse files
app.py
CHANGED
|
@@ -2,7 +2,7 @@
|
|
| 2 |
Blueprint Door and Window Detection - API Server
|
| 3 |
"""
|
| 4 |
|
| 5 |
-
from fastapi import FastAPI, File, UploadFile
|
| 6 |
from fastapi.middleware.cors import CORSMiddleware
|
| 7 |
from ultralytics import YOLO
|
| 8 |
from PIL import Image
|
|
@@ -30,40 +30,36 @@ def load_model():
|
|
| 30 |
|
| 31 |
if os.path.exists(model_path):
|
| 32 |
print(f"π¦ Loading trained model from: {model_path}")
|
| 33 |
-
model = YOLO(model_path)
|
| 34 |
print("β
Model loaded successfully!")
|
| 35 |
else:
|
| 36 |
print("β οΈ Trained model not found. Using pre-trained YOLOv8n for testing.")
|
| 37 |
model = YOLO('yolov8n.pt')
|
| 38 |
print("π Using default model - train your custom model first!")
|
| 39 |
|
| 40 |
-
@app.
|
| 41 |
-
def home():
|
| 42 |
"""Health check endpoint"""
|
| 43 |
-
return
|
| 44 |
"message": "Blueprint Door and Window Detection API",
|
| 45 |
"status": "running",
|
| 46 |
"version": "1.0.0",
|
| 47 |
"model_loaded": model is not None
|
| 48 |
-
}
|
| 49 |
|
| 50 |
-
@app.
|
| 51 |
-
def detect_objects():
|
| 52 |
"""Detection endpoint that accepts image and returns bounding boxes"""
|
| 53 |
if model is None:
|
| 54 |
load_model()
|
| 55 |
|
| 56 |
-
if 'image' not in request.files:
|
| 57 |
-
return jsonify({'error': 'No image file provided'}), 400
|
| 58 |
-
|
| 59 |
-
file = request.files['image']
|
| 60 |
if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
|
| 61 |
-
|
| 62 |
|
| 63 |
try:
|
| 64 |
# Read and process image
|
| 65 |
-
|
| 66 |
-
image = Image.open(io.BytesIO(
|
| 67 |
|
| 68 |
# Run inference
|
| 69 |
results = model(image)[0]
|
|
@@ -83,15 +79,16 @@ def detect_objects():
|
|
| 83 |
round(x2-x1, 1), round(y2-y1, 1)] # [x, y, width, height]
|
| 84 |
})
|
| 85 |
|
| 86 |
-
return
|
| 87 |
'detections': detections,
|
| 88 |
'image_size': {'width': image.width, 'height': image.height},
|
| 89 |
'processing_time_ms': round(results.speed['inference'], 1)
|
| 90 |
-
}
|
| 91 |
|
| 92 |
except Exception as e:
|
| 93 |
-
|
| 94 |
|
| 95 |
-
|
|
|
|
|
|
|
| 96 |
load_model()
|
| 97 |
-
app.run(host='0.0.0.0', port=5000)
|
|
|
|
| 2 |
Blueprint Door and Window Detection - API Server
|
| 3 |
"""
|
| 4 |
|
| 5 |
+
from fastapi import FastAPI, File, UploadFile, HTTPException
|
| 6 |
from fastapi.middleware.cors import CORSMiddleware
|
| 7 |
from ultralytics import YOLO
|
| 8 |
from PIL import Image
|
|
|
|
| 30 |
|
| 31 |
if os.path.exists(model_path):
|
| 32 |
print(f"π¦ Loading trained model from: {model_path}")
|
| 33 |
+
model = YOLO(model_path)
|
| 34 |
print("β
Model loaded successfully!")
|
| 35 |
else:
|
| 36 |
print("β οΈ Trained model not found. Using pre-trained YOLOv8n for testing.")
|
| 37 |
model = YOLO('yolov8n.pt')
|
| 38 |
print("π Using default model - train your custom model first!")
|
| 39 |
|
| 40 |
+
@app.get("/")
|
| 41 |
+
async def home():
|
| 42 |
"""Health check endpoint"""
|
| 43 |
+
return {
|
| 44 |
"message": "Blueprint Door and Window Detection API",
|
| 45 |
"status": "running",
|
| 46 |
"version": "1.0.0",
|
| 47 |
"model_loaded": model is not None
|
| 48 |
+
}
|
| 49 |
|
| 50 |
+
@app.post("/detect")
|
| 51 |
+
async def detect_objects(file: UploadFile = File(...)):
|
| 52 |
"""Detection endpoint that accepts image and returns bounding boxes"""
|
| 53 |
if model is None:
|
| 54 |
load_model()
|
| 55 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
|
| 57 |
+
raise HTTPException(status_code=400, detail="Invalid file format. Only PNG/JPG allowed")
|
| 58 |
|
| 59 |
try:
|
| 60 |
# Read and process image
|
| 61 |
+
contents = await file.read()
|
| 62 |
+
image = Image.open(io.BytesIO(contents))
|
| 63 |
|
| 64 |
# Run inference
|
| 65 |
results = model(image)[0]
|
|
|
|
| 79 |
round(x2-x1, 1), round(y2-y1, 1)] # [x, y, width, height]
|
| 80 |
})
|
| 81 |
|
| 82 |
+
return {
|
| 83 |
'detections': detections,
|
| 84 |
'image_size': {'width': image.width, 'height': image.height},
|
| 85 |
'processing_time_ms': round(results.speed['inference'], 1)
|
| 86 |
+
}
|
| 87 |
|
| 88 |
except Exception as e:
|
| 89 |
+
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
|
| 90 |
|
| 91 |
+
@app.on_event("startup")
|
| 92 |
+
async def startup_event():
|
| 93 |
+
"""Load the model when the server starts"""
|
| 94 |
load_model()
|
|
|