kakatt / app.py
codernotme's picture
Upload 6 files
e2d9248 verified
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import io
from model_utils import predict, download_model_if_not_exists
app = FastAPI(
title="Face Shape Classification API",
description="API for classifying face shapes using the fahd9999/face_shape_classification model.",
version="1.0.0"
)
# Enable CORS for Vercel frontend
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Allow all origins for simplicity in this context, or specify Vercel domain
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.on_event("startup")
async def startup_event():
"""Ensure model is downloaded on startup."""
try:
download_model_if_not_exists()
except Exception as e:
print(f"Startup warning: Could not download model: {e}")
@app.get("/")
async def root():
return {"message": "Face Shape Classification API is running"}
@app.post("/predict")
async def predict_endpoint(file: UploadFile = File(...)):
if not file.content_type.startswith("image/"):
raise HTTPException(status_code=400, detail="File must be an image")
try:
contents = await file.read()
image_stream = io.BytesIO(contents)
result = predict(image_stream)
return JSONResponse(content=result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)