Spaces:
Sleeping
Sleeping
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'] | |
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) | |