| 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 |
| |
|
|
| 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") |
