from fastai.vision.all import *
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from PIL import Image
import io
import base64

# Load the model
learn = load_learner('model.pkl')

app = FastAPI()

class ImageData(BaseModel):
    image: str

def predict_image(img):
    img = img.convert("L")
    img = img.resize((28, 28))
    img = np.array(img)
    return f"{learn.predict(img)[0][0]:.2f}"

@app.post("/predict")
async def predict(data: ImageData):
    try:
        image_data = base64.b64decode(data.image)
        img = Image.open(io.BytesIO(image_data))
        probability = predict_image(img)
        return {"probability": probability}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))