shekzee's picture
Update app.py
6e77308 verified
raw
history blame
820 Bytes
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np
import io
app = FastAPI()
# Load model and class names
model = load_model("hf_keras_model.keras")
class_names = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']
@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
contents = await file.read()
image = Image.open(io.BytesIO(contents)).convert("RGB")
image = image.resize((150, 150))
img_array = np.array(image) / 255.0
img_array = np.expand_dims(img_array, axis=0)
preds = model.predict(img_array)[0]
confidences = {class_names[i]: float(preds[i]) for i in range(len(class_names))}
return JSONResponse(content=confidences)