Spaces:
Sleeping
Sleeping
| """Inference API for local use and Docker Spaces. No training imports.""" | |
| import io | |
| import os | |
| from contextlib import asynccontextmanager | |
| from fastapi import FastAPI, File, HTTPException, UploadFile | |
| from PIL import Image, UnidentifiedImageError | |
| MAX_UPLOAD = 8 * 1024 * 1024 | |
| def create_app(recognizer=None): | |
| async def lifespan(app): | |
| if recognizer is None: | |
| threads = int(os.getenv("MATH2LATEX_THREADS", "2")) | |
| if os.getenv("MATH2LATEX_BACKEND", "formulanet") == "formulanet": | |
| from image_to_latex.formulanet import FormulaNetRecognizer | |
| app.state.recognizer = FormulaNetRecognizer( | |
| os.getenv("MATH2LATEX_MODEL_DIR", "artifacts/formulanet"), threads | |
| ) | |
| else: | |
| from image_to_latex.inference import Recognizer | |
| app.state.recognizer = Recognizer( | |
| os.getenv("MATH2LATEX_MODEL_DIR", "artifacts/latex-ocr"), | |
| threads, | |
| max_tokens=int(os.getenv("MATH2LATEX_MAX_TOKENS", "192")), | |
| quantized=os.getenv("MATH2LATEX_INT8", "0") == "1", | |
| ) | |
| else: | |
| app.state.recognizer = recognizer | |
| yield | |
| app = FastAPI( | |
| title="Math2LaTeX", | |
| description="Convert a cropped formula into LaTeX.", | |
| lifespan=lifespan, | |
| ) | |
| def health(): | |
| return {"status": "ready", "backend": app.state.recognizer.backend} | |
| def predict(file: UploadFile = File(...)): | |
| payload = file.file.read(MAX_UPLOAD + 1) | |
| if len(payload) > MAX_UPLOAD: | |
| raise HTTPException(413, "Maximum upload size is 8 MiB") | |
| try: | |
| with Image.open(io.BytesIO(payload)) as image: | |
| result = app.state.recognizer.predict(image) | |
| return {"data": {"pred": result["latex"], **result}} | |
| except ( | |
| UnidentifiedImageError, | |
| ValueError, | |
| OSError, | |
| Image.DecompressionBombError, | |
| ) as exc: | |
| raise HTTPException(400, str(exc)) from exc | |
| return app | |
| app = create_app() | |