tan200224/Synthetic-CT-Scan_VAE_Conditional

1"""2VAE Hugging Face Space app.3Upload a mask image -> encode -> decode -> return one slice (slice 2 of 4).4"""5import base646import io7import logging8from fastapi import FastAPI, HTTPException, UploadFile9 10logging.basicConfig(level=logging.INFO)11logger = logging.getLogger(__name__)12from fastapi.middleware.cors import CORSMiddleware13from fastapi.responses import JSONResponse14import gradio as gr15from PIL import Image16import uvicorn17 18from inference import inference_to_png, OUTPUT_SLICE_INDEX19 20 21# --- Gradio UI ---22def run_inference(mask: Image.Image) -> Image.Image:23    if mask is None:24        raise gr.Error("Please upload a mask image.")25    png_bytes = inference_to_png(mask, slice_index=OUTPUT_SLICE_INDEX)26    return Image.open(io.BytesIO(png_bytes)).convert("L")27 28 29demo = gr.Interface(30    fn=run_inference,31    inputs=gr.Image(label="Mask (grayscale)", type="pil"),32    outputs=gr.Image(label=f"Output slice {OUTPUT_SLICE_INDEX} of 4"),33    title="VAE CT Slice Generator",34    description=(35        "Upload a **mask** image (grayscale). The model encodes it, decodes to 4 slices (3D CT), "36        f"and returns **slice {OUTPUT_SLICE_INDEX}** as a 2D image for the web."37    ),38)39 40 41# --- FastAPI app (Gradio mounted at /) ---42app = FastAPI(title="VAE CT Slice API")43 44# CORS: allow your website (and others) to call /predict and /generate from the browser45app.add_middleware(46    CORSMiddleware,47    allow_origins=["*"],48    allow_credentials=True,49    allow_methods=["GET", "POST", "OPTIONS"],50    allow_headers=["*"],51)52 53 54@app.post("/generate")55async def generate(file: UploadFile):56    """Upload mask -> return single slice as data URI (same shape as diffusion /generate)."""57    try:58        raw = await file.read()59        logger.info("[request /generate] INPUT: filename=%s content_type=%s raw_bytes=%s",60                    file.filename, file.content_type, len(raw))61        mask = Image.open(io.BytesIO(raw)).convert("L")62        logger.info("[request /generate] image opened: size=%s mode=%s", mask.size, mask.mode)63        png_bytes = inference_to_png(mask, slice_index=OUTPUT_SLICE_INDEX)64        b64 = base64.b64encode(png_bytes).decode("ascii")65        return JSONResponse(content={"image": f"data:image/png;base64,{b64}"})66    except Exception as e:67        raise HTTPException(status_code=400, detail=str(e))68 69 70@app.post("/predict")71async def predict(file: UploadFile):72    """Upload mask -> return single slice as base64 PNG in JSON."""73    try:74        raw = await file.read()75        logger.info("[request /predict] INPUT: filename=%s content_type=%s raw_bytes=%s",76                    file.filename, file.content_type, len(raw))77        mask = Image.open(io.BytesIO(raw)).convert("L")78        logger.info("[request /predict] image opened: size=%s mode=%s", mask.size, mask.mode)79        png_bytes = inference_to_png(mask, slice_index=OUTPUT_SLICE_INDEX)80        b64 = base64.b64encode(png_bytes).decode("ascii")81        return JSONResponse(content={"image": b64, "slice_index": OUTPUT_SLICE_INDEX})82    except Exception as e:83        raise HTTPException(status_code=400, detail=str(e))84 85 86@app.get("/")87def root():88    return {"status": "ok", "message": "VAE CT slice API. Use /generate or /predict with a mask image."}89 90 91app = gr.mount_gradio_app(app, demo, path="/")92 93 94if __name__ == "__main__":95    uvicorn.run(app, host="0.0.0.0", port=7860)96