Mask VAE Autoencoder inference weird result

[cbe135]

Mask VAE inference seems to like to add a box to inference results. For SlidingWindow inference, the result will be many small boxes in the background. Specifically, the issue seems to be at the encoder end, seeing that with a diffusion-model generated latent, the decoding side generates a normal mask.

Following are in the order of original masks and autoencoder outputs. The autoencoder applied uses default weights mask_generation_autoencoder.pt.

[Can-Zhao]

Regarding to Mask VAE, @guopengf may answer the question. My intuition is that the issue comes from Gaussian smoothing during stitching.

[cbe135]

Thank you for the suggestion. I've attempted an inference with overlap of 0. The generated result is attached below.

[guopengf]

Hi @cbe135， would you like to share the settings of SlidingWindow in your code? In addition, would you like to elaborate on your use case? Are you trying to use the mask autoencoder standalone?

[cbe135]

Hi yes, my SlidingWindow uses the following setting
roi_size=[48,48,48],
sw_batch_size=1,
progress=True,
mode="gaussian",
overlap=0.6666,
device=device,
sw_device=device,

I'm trying to generate masks with mask diffusion model. We are adding about 10+ organs. Thus, we would need to finetune both mask diffusion (need mask latents) and mask autoencoder. Also, since we are adding more organs than the available dummy variables, we would need to change the output_channel of the autoencoder. This error was discovered in trying to generate the mask latents.

[Can-Zhao]

@guopengf