File size: 1,349 Bytes
3c4cb64 0635f64 3c4cb64 0635f64 3c4cb64 0635f64 c5fba9a 3c4cb64 c5fba9a 3c4cb64 c5fba9a 3c4cb64 0635f64 | 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 | import io
from PIL import Image
from fastapi.responses import RedirectResponse
import torch
from pydantic import BaseModel, Field
from contextlib import asynccontextmanager
from fastapi import FastAPI, Response
from src import utils
class GenerateRequest(BaseModel):
number: int = Field(ge=0, le=9, description="El número a generar (0-9)")
@asynccontextmanager
async def lifespan(app: FastAPI):
print("\n\n\n====================\n\n. INICIO API REST \n")
app.state.model = utils.load_best_model()
yield # <- aquí la API queda disponible
# Empieza el teardown
print("\n\n\n====================\n\n. FIN API REST \n")
app = FastAPI(
title="Generador de Imágenes MNIST",
lifespan=lifespan,
swagger_ui_parameters={"defaultModelsExpandDepth": -1}
)
@app.post("/generateNumber", tags=["Generación"])
def generateNumber(body: GenerateRequest):
generated_tensor = app.state.model.generate_number(body.number)
img_array = (generated_tensor * 255).astype("uint8")
img = Image.fromarray(img_array).resize((280, 280))
buf = io.BytesIO()
img.save(buf, format="PNG")
byte_im = buf.getvalue()
return Response(content=byte_im, media_type="image/png")
@app.get("/", include_in_schema=False)
async def root():
return RedirectResponse(url="/docs") |