AutoTrain - Part 2: Recipes

from fastapi import APIRouter, File, UploadFile, BackgroundTasks
import typer
from io import BytesIO
from PIL import Image
import os
from loguru import logger

router = FastAPI()
cli = Typer()

@cli.command()
def train(config=None):
    from auto_train.classification.train import train_model
    train_model(config)

@router.post('/validate/')
async def validate(img: UploadFile = File(...)):
    from auto_train.classification.infer import infer
    image = Image.open(BytesIO(img.file.read()))
    logger.info(img.filename)
    img_path = f'test_images/from_api/{img.filename}'
    image.save(img_path)
    return infer(img_path)

if __name__ == '__main__':
    cli()

> Use `typer`'s utility to create a command line functionality
$ python main.py train --config='configs/mnist.ini'



> FastAPI will expose an endpoint localhost:8889/validate/
which will take an image and return the predictions