VAE CT Slice Generator
Generate a single 2D CT slice from a segmentation mask using a VAE trained on 3D CT data. The model encodes your mask, decodes to 4 slices (3D volume), and returns slice 2 as a 2D image for the web.
How to use
- Upload a mask image — Grayscale segmentation mask (any size; it will be resized to 256×256).
- Click Run — The model encodes the mask, decodes to 4 slices, and shows slice 2.
- Download — Use the output image as needed.
The model expects a single grayscale mask. It is duplicated into 4 “slices” (matching training), then the VAE encodes → decodes to 4 slices; only slice index 2 is returned because the API cannot return 3D volumes.
API
POST /generate
Upload a mask image; get back a single slice as a data URI.
Request: multipart/form-data with a file (mask image).
Response:
{
"image": "data:image/png;base64,..."
}
POST /predict
Upload a mask image; get back the slice as base64 PNG.
Request: multipart/form-data with a file (mask image).
Response:
{
"image": "<base64 PNG string>",
"slice_index": 2
}
Model & input/output
- Model: VAE trained on 3D CT data (4-channel input/output, 256×256).
- Input (training): Label/mask as 4 slices
(1, 4, 256, 256). This Space accepts one grayscale mask, resizes to 256×256, normalizes to [0, 1], and duplicates it to 4 slices. - Output (model):
decode(embedding)has shape(1, 4, 256, 256)(4 slices). This Space returns one slice (index 2) as a 2D PNG.
Environment (optional)
| Variable | Default | Description |
|---|---|---|
MODEL_REPO |
tan200224/Synthetic-CT-Scan_VAE_Conditional |
Hugging Face repo with the model. |
MODEL_FILENAME |
mask2pic_64model_47.pt |
Filename of the checkpoint. |
Set these in your Space Settings → Variables and secrets if you use a different repo or file.
Run locally
pip install -r requirements.txt
python app.py
Open http://localhost:7860 for the Gradio UI; use /generate or /predict for the API.