File size: 1,489 Bytes
54a9640 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
# Adapted from https://github.com/kingyiusuen/image-to-latex/blob/main/api/app.py
from http import HTTPStatus
from fastapi import FastAPI, File, UploadFile, Form
from PIL import Image
from io import BytesIO
from pix2tex.cli import LatexOCR
model = None
app = FastAPI(title='pix2tex API')
def read_imagefile(file) -> Image.Image:
image = Image.open(BytesIO(file))
return image
@app.on_event('startup')
async def load_model():
global model
if model is None:
model = LatexOCR()
@app.get('/')
def root():
'''Health check.'''
response = {
'message': HTTPStatus.OK.phrase,
'status-code': HTTPStatus.OK,
'data': {},
}
return response
@app.post('/predict/')
async def predict(file: UploadFile = File(...)) -> str:
"""Predict the Latex code from an image file.
Args:
file (UploadFile, optional): Image to predict. Defaults to File(...).
Returns:
str: Latex prediction
"""
global model
image = Image.open(file.file)
return model(image)
@app.post('/bytes/')
async def predict_from_bytes(file: bytes = File(...)) -> str: # , size: str = Form(...)
"""Predict the Latex code from a byte array
Args:
file (bytes, optional): Image as byte array. Defaults to File(...).
Returns:
str: Latex prediction
"""
global model
#size = tuple(int(a) for a in size.split(','))
image = Image.open(BytesIO(file))
return model(image, resize=False)
|