HalimElsa/BackEnd-VAE

1# api_server.py2from fastapi import FastAPI, HTTPException3from fastapi.responses import StreamingResponse, JSONResponse4from fastapi.middleware.cors import CORSMiddleware5import numpy as np6from io import BytesIO7from generator_api import generate_handwriting8from PIL import Image9import uvicorn10 11app = FastAPI()12 13# Enable CORS for frontend connection14app.add_middleware(15    CORSMiddleware,16    allow_origins=["*"],17    allow_credentials=True,18    allow_methods=["*"],19    allow_headers=["*"],20)21 22 23@app.get("/generate")24def generate_image(latent_x: float, latent_y: float, digit_label: int):25    """Generate a handwriting sample from latent values and digit label."""26    print(f"[request] latent_x={latent_x}, latent_y={latent_y}, digit_label={digit_label}")27 28    if digit_label < 0 or digit_label > 9:29        raise HTTPException(status_code=400, detail="digit_label must be between 0 and 9")30 31    # 1. Generate image 28x2832    try:33        img = generate_handwriting(latent_x, latent_y, digit_label)34    except Exception as e:35        print(f"[error] generate_handwriting failed: {e}")36        raise HTTPException(status_code=500, detail="Model failed to generate image")37 38    if not isinstance(img, np.ndarray):39        raise HTTPException(status_code=500, detail="Model did not return an array")40 41    print(f"[generate] image shape: {img.shape}")42 43    # 2. Convert float -> uint8 (0-255)44    img_uint8 = (img * 255).astype(np.uint8)45    print(f"[convert] uint8 min={img_uint8.min()}, max={img_uint8.max()}")46 47    # 3. Convert to PNG in-memory48    pil_img = Image.fromarray(img_uint8)49    buffer = BytesIO()50    pil_img.save(buffer, format="PNG")51    buffer.seek(0)52    print(f"[response] PNG size={buffer.getbuffer().nbytes} bytes")53 54    # 4. Send back as PNG stream55    return StreamingResponse(buffer, media_type="image/png")56 57 58@app.get("/health")59def health_check():60    """Simple health endpoint for readiness checks."""61    return {"status": "ok"}62 63 64if __name__ == "__main__":65    uvicorn.run(app, host="0.0.0.0", port=7860)66