MAISI Mask Autoencoder

[cbe135]

During the implementation of our own Mask Autoencoder, we encountered two questions.
Specifically, our dataset masks have the characteristic of input channel of 1 and output channel of 1.
Our question relates to the different input channel and output channel setting used in pretraining the provided model weights.
The output channel count of 128 makes sense as there are 128 possible option choices for organs and disease phenomena.
In terms of the input channel count, we were wondering why the input channel was 7?
Also, in terms of fine tuning it to our dataset, we were thinking of averaging/modifying the first layer and last layer of the model into the desired input or output. Are there other recommended suggestions?
Thank you.

[Can-Zhao]

@guopengf

[guopengf]

Hi @cbe135, the input channel of the mask autoencoder is 8 (see

tutorials/generation/maisi/configs/config_maisi3d-rflow.json

Line 94 in 8b90a16

"in_channels": 8,

).
This is mainly because we use binary representation to encode the input mask, which saves memory. 8 channels can represent 2**8 (0~255) labels. Each channel represents a bit (see the following function).

tutorials/generation/maisi/scripts/utils.py

Lines 175 to 190 in 8b90a16

	def binarize_labels(x: Tensor, bits: int = 8) -> Tensor:
	"""
	Convert input tensor to binary representation.
	This function takes an input tensor and converts it to a binary representation
	using the specified number of bits.
	Args:
	x (Tensor): Input tensor with shape (B, 1, H, W, D).
	bits (int, optional): Number of bits to use for binary representation. Defaults to 8.
	Returns:
	Tensor: Binary representation of the input tensor with shape (B, bits, H, W, D).
	"""
	mask = 2 ** torch.arange(bits).to(x.device, x.dtype)
	return x.unsqueeze(-1).bitwise_and(mask).ne(0).byte().squeeze(1).permute(0, 4, 1, 2, 3)

For example, label 1 is encoded as [0, 0, 0, 0, 0, 0, 0, 1].

For your use case, is the label of your dataset covered in the pre-defined label dict?

[cbe135]

Hi @guopengf, thank you.
That makes sense.
Our mask target is hypo tumor so it is not on the list.
Would it make sense to use one of the dummies?
Does it matter which dummy is used?
Looking specifically at the label dict, is there any specific logic on how the different labels are ordered?
Thank you.

[guopengf]

The order of labels does not matter, so you can use any dummy label.

[mar-cry]

Hello @guopengf , my dataset has three categories that are not included in the 132 predefined categories in the label_dict, so I mapped these three categories to dummy6, dummy7, and dummy8. After converting them into the input format compatible with the Mask Autoencoder, encoding the input mask, and then decoding to reconstruct, I found that the reconstructed result differs greatly from the input. Is this normal? Here are illustrative examples of my input and output.

[guopengf]

Have you trained the AE with dummy labels? If not, this result should be expected. Please try to fine-tune AE with your dataset.

[mar-cry]

I don't fine-tune the mask autoencoder on our dataset. I notice that the input and output channels of mask autoencoder are different from the common image autoencoder, is there any difference between the training pipeline of two autoencoders, such as the training loss?