[Example] Add 2D Multi-Organ Segmentation with TransUNet - #2361
[Example] Add 2D Multi-Organ Segmentation with TransUNet#2361yassienashrafwasfy wants to merge 7 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a TransUNet implementation for 2D multi-organ segmentation on the Synapse dataset, provided as both a Jupyter notebook and a Python script. The feedback focuses on improving the model's flexibility and performance: specifically, removing the hardcoded TensorFlow backend to support Keras 3's multi-backend capabilities, and updating the architecture to support non-square input images by parameterizing the encoder and feature map calculations. Additionally, it is suggested to exclude the background class from the Dice loss calculation to better prioritize small organ segmentation.
|
|
||
| import os | ||
|
|
||
| os.environ["KERAS_BACKEND"] = "tensorflow" |
There was a problem hiding this comment.
The PR description states that this example is Keras 3 compatible and backend-agnostic. However, hardcoding the backend to "tensorflow" prevents users from running this example with JAX or PyTorch. Since the model is built using keras.ops and tf.data is supported as a universal data loader in Keras 3, this line should be removed to allow multi-backend usage.
| def _build_resnet_encoder(): | ||
| """Returns a ResNet-50 feature pyramid as a plain keras.Model.""" | ||
| inputs = keras.Input(shape=(224, 224, 3)) |
There was a problem hiding this comment.
The input shape is hardcoded to (224, 224, 3) in the encoder builder. This ignores the image_size parameter passed to the TransUNet constructor. If a user initializes the model with a different image size, it will lead to a shape mismatch in the backbone.
| def _build_resnet_encoder(): | |
| """Returns a ResNet-50 feature pyramid as a plain keras.Model.""" | |
| inputs = keras.Input(shape=(224, 224, 3)) | |
| def _build_resnet_encoder(input_shape=(224, 224, 3)): | |
| """Returns a ResNet-50 feature pyramid as a plain keras.Model.""" | |
| inputs = keras.Input(shape=input_shape) |
| feature_map_size = image_size[0] // 16 | ||
| self.cnn_backbone = _build_resnet_encoder() | ||
| self.patch_embedding = PatchEmbedding( | ||
| patch_size, embedding_dim, (feature_map_size // patch_size) ** 2 | ||
| ) |
There was a problem hiding this comment.
The calculation of num_patches and the initialization of the backbone assume a square input image. To support non-square images, the height and width feature map sizes should be calculated independently.
| feature_map_size = image_size[0] // 16 | |
| self.cnn_backbone = _build_resnet_encoder() | |
| self.patch_embedding = PatchEmbedding( | |
| patch_size, embedding_dim, (feature_map_size // patch_size) ** 2 | |
| ) | |
| h_feat, w_feat = image_size[0] // 16, image_size[1] // 16 | |
| self.cnn_backbone = _build_resnet_encoder(image_size) | |
| self.patch_embedding = PatchEmbedding( | |
| patch_size, embedding_dim, (h_feat // patch_size) * (w_feat // patch_size) | |
| ) |
| fmap_size = self.image_size[0] // 16 // self.patch_size | ||
| decoded = ops.reshape(tokens, (-1, fmap_size, fmap_size, self.embedding_dim)) |
There was a problem hiding this comment.
The reshaping of tokens in the call method assumes a square feature map. This will fail if the input image is not square.
| fmap_size = self.image_size[0] // 16 // self.patch_size | |
| decoded = ops.reshape(tokens, (-1, fmap_size, fmap_size, self.embedding_dim)) | |
| h_fmap = self.image_size[0] // 16 // self.patch_size | |
| w_fmap = self.image_size[1] // 16 // self.patch_size | |
| decoded = ops.reshape(tokens, (-1, h_fmap, w_fmap, self.embedding_dim)) |
| y_pred = ops.clip(y_pred, 1e-6, 1.0) | ||
| intersection = ops.sum(y_true_oh * y_pred, axis=[1, 2]) | ||
| union = ops.sum(y_true_oh + y_pred, axis=[1, 2]) | ||
| dice_loss = 1 - ops.mean((2.0 * intersection + smooth) / (union + smooth)) |
There was a problem hiding this comment.
The combined_loss includes the background class (index 0) in the Dice loss calculation. In multi-organ segmentation, the background typically dominates the image area, which can bias the Dice score and lead the model to prioritize background accuracy over small organ segmentation. It is standard practice to exclude the background class from the Dice loss.
| dice_loss = 1 - ops.mean((2.0 * intersection + smooth) / (union + smooth)) | |
| dice_loss = 1 - ops.mean((2.0 * intersection[:, 1:] + smooth) / (union[:, 1:] + smooth)) |
2 participants
Referencing: (#2340)
This PR introduces a new Computer Vision example: 2D Multi-Organ Segmentation with TransUNet.
Overview:
This example implements TransUNet (Chen et al., 2021), a hybrid architecture that leverages both the
high-resolution spatial features of CNNs and the global context modeling of Vision Transformers (ViT). The
model is designed for medical image segmentation, specifically targeting multi-organ structures.
Key Features:
global encoding, and a cascaded upsampling decoder (CUP) with U-Net skip connections.
and .h5 file formats for training and volume-based validation.
TensorFlow, JAX, and PyTorch.
Loss, along with a custom MeanIoU wrapper for sparse labels.
predictions against ground truth overlays.