yeagerd/fastai-ch2-api

1from fastai.vision.all import *2from fastapi import FastAPI3from fastapi.middleware.cors import CORSMiddleware4from PIL import Image5import base646from io import BytesIO7from pydantic import BaseModel8import logging9 10app = FastAPI()11app.add_middleware(12    CORSMiddleware,13    allow_origins=['*'],14    allow_methods=["*"],15    allow_headers=["*"]16)17 18learn = load_learner('bear-types-model.pkl')19labels = learn.dls.vocab20 21logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')22# , filename='app.log', filemode='a'23logger = logging.getLogger(__name__)24 25@app.get("/")26def greet_json():27    logger.info('greet')28    return {"Hello": "World!"}29 30 31class PredictRequest(BaseModel):32    context: str33    data: str34 35@app.post("/predict")36def predict_json(img: PredictRequest):37    logger.info(f'predict: {img.context}')38    # https://stackoverflow.com/a/7550171339    decoded_image_data = base64.b64decode(img.data.split(",")[1])40    pil = PILImage.create(decoded_image_data)41    label,index,probs = learn.predict(pil)42    return { "data": {43        "label": label,44        "index": int(index),45        "confidences": {labels[i]: float(probs[i]) for i in range(len(labels))}46    }}47