Flemme: A Flexible and Modular Learning Platform for Medical Images

We introduce two types of building blocks based on convolution: ConvBlock and ResConvBlock. A ConvBlock consists of a convolution, a normalization layer, and an activation function. Formally, we consider a forward pass of ConvBlock defined as:

where $\operatorname{RL}$ and $\operatorname{GN}$ denote the rectified linear unit (ReLU) and group normalization [16], which are the default choices of activation and normalization functions. Two MLPs, named Gate and Bias, are optional to introduce a gate-bias mechanism if there is a context vector embedding input $t$ [17]:

As illustrated in Fig. 1 (a), ResConvBlock contains two ConvBlocks and introduces a skip connection for residual learning:

where $\operatorname{Proj}$ is a MLP for context embedding projection, $\mathcal{F}_{C}^{1}$ and $\mathcal{F}_{C}^{2}$ are the first and second ConvBlock, respectively. Note that $\mathcal{F}_{C}^{2}$ in ResConvBlock follows a similar manner to [7] and has no activation function.

Following Dosovitskiy et al.[6], we implement ViTBlock and further introduce the shifted-window strategy [8] to construct SwinBlock for time and space complexity reduction. Specifically, we partition image patches into smaller windows, and a window-based MSA (W-MSA) is employed to capture dependencies only among patches within the same window. Relative position encoding for both 2D and 3D windows are used to further improve the modeling capability. By shifting the windows at different stages, overlaps among windows are created to increase the receptive field. Fig 1 (b) gives an intuitive overview of SwinBlock that can be formulated as the following:

where $\operatorname{LN}$ and $\operatorname{Drop}$ refers to layer normalization [18] and drop path [19], $\operatorname{W-MSA}$ represents a series of operations, including patch shift, window partition, applying window-based MSA and the reversed operations. Similar to ConvBlock, we also introduce ResSwinBlock $\mathcal{F}_{RS}$, by simply replacing $\mathcal{F}_{C}(x)$ with $\mathcal{F}_{S}(x)$ in Eqn. (3).

SSM is a linear time-invariant system to map a 1-d sequence $x(t)$ to $y(t)$ through a hidden state $h(t)$. We adopt Mamba [9, 20] as the implementation of SSM to construct MambaBlock. In addition, we traverse the 2D and 3D image patches in a cross-scan manner and jointly model the obtained sequences to enhance the spatial awareness of Mamba. As depicted in Fig. 1 (c), MambaBlock processes image patches as follows:

where $\operatorname{Proj}$ projects image patch into inner space, $\operatorname{SSM}$ indicates multiple operations including cross scanning, sequence modeling with Mamba, and scan merging. ResMambaBlock $\mathcal{F}_{RM}$ are introduced by using $\mathcal{F}_{M}(x)$ as a replacement for $\mathcal{F}_{C}(x)$ in Eqn. (3).

For all the aforementioned encoders, we also implement their U-shaped variants. The U-shaped encoder outputs feature maps from down-sampling layers, which are concatenated with the feature maps from up-sampling layers of corresponding spatial resolutions in the decoder. U-shaped networks have demonstrated superiority and become the de-facto standard in various imaging tasks [3, 14, 21, 22, 23]. Fig. 4 gives an illustration of a U-shaped structure and horizontal feature integration.

Because the building blocks are capable of handling context embeddings, the encoder and decoder can also incorporate an additional context embedding as input. To summarize, the basic encoding and decoding processes can be formulated as follows:

where $\mathcal{E}$ and $\mathcal{D}$ represent the encoder and decoder, $\mathcal{P}(\cdot)$ and $\mathcal{P}^{-1}(\cdot)$ refer to the patching encoding and decoding operations, $\operatorname{Down}$ and $\operatorname{Up}$ refer to the down-sampling and up-sampling layers, and $\operatorname{Middle}$ refers to the middle layers, respectively.

Our base architecture follows an encoder-decoder style with optional components for encoding context as mentioned in Sec. II-B. A forward pass of based architecture is defined as:

where $\mathcal{C}_{e}$ and $\mathcal{C}_{d}$ are functions of encoding context $c$ for encoder $\mathcal{E}$ and decoder $\mathcal{D}$, $\mathcal{T}(\cdot)$ computes the sinusoidal positional embedding, $z$ is the latent embedding computed by the encoder, $c$ and $t$ are input contexts. Usually, $c$ is the input condition such as a prompt image or class label and $t$ is a time step index. Base architecture is not directly trainable due to the lack of loss functions and serves as a foundational structure for other architectures as illustrated in Fig. 3, including the Segmentation Model (SeM), the Auto-Encoder (AE) and the De-noising Diffusion Probabilistic Model (DDPM), for medical image segmentation, reconstruction, and generation, respectively. We explain these architectures in the remaining part of this section. To simplify the notation, we omit the optional input contexts in the following formulas to simplify the notation.

As shown in Fig. 3 (a), SeM is a simple extension of base architecture by introducing segmentation loss. It takes raw images as inputs and gives the predictions through a forward pass defined in Eqn. (7). Loss is computed over prediction and ground truth:

where $y^{\prime}=\mathcal{M}(x)$ is the prediction, $y$ specifies the ground-truth. By default, segmentation loss $\mathcal{L}_{\text{seg}}$ is a combination of cross-entropy and Dice loss, denoted by $\mathcal{L}_{\text{CE}}$ and $\mathcal{L}_{\text{Dice}}$, respectively.

We implement auto-encoder (AE) for medical image restoration, which builds reconstruction loss over predictions and inputs, indicating an unsupervised representation learning:

where $\mathcal{L}_{\text{MSE}}$ is the mean squared error, serving as the default reconstruction loss. As shown in Fig. 3 (b), AE can be regularized by distribution loss to learn a more continuous latent representation and have a certain capability of generation [13]. The learning process of AE can also be supervised by a ground-truth target.

Fig. 3 (c) illustrates the pipeline of DDPM [14], in which base architecture is used to construct a noise predictor, denoted as $\epsilon$-model. In the diffusion process, we add random Gaussian noise to the input image to get $t$-th step noisy images, which can be formulated as the following:

where $\alpha_{t}$ and $\beta_{t}$ are hyper-parameters related to noise schedules. The $\epsilon$-model takes noisy images as inputs and computes loss over the noise and prediction:

where $\epsilon^{\prime}=\mathcal{M}_{\epsilon}(x_{t},t)$ is the predicted noise. In the reversed diffusion process, we gradually remove the predicted noise from the current noisy image to recover the image of the last time step $t-1$. Refer to [14] for more Details about the de-noising process.

Inspired by Ronneberger et al.[3], who propose to connect the feature maps from encoder and decoder for horizontal feature fusion, we introduce a generic hierarchical architecture to refine and fuse features vertically. In specific, we assign a building block and a feature decoding block for each stage in the decoder to generate multi-resolution predictions, denoted as $x_{1}^{\prime},x_{2}^{\prime},\cdots,x_{n}^{\prime}$. We introduce a combination layer, where all predictions are weighted with a learnable position encoding and integrated to produce the final output:

In the above, $\mathop{\mathchoice{\leavevmode\hbox to12.22pt{\vbox to10pt{\pgfpicture\makeatletter\hbox{\hskip 6.1111pt\lower-5.00006pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } {{}{}}{{}}{}{{{}}{}{}{}{}{}{}{}{}}{}\pgfsys@moveto{-1.29166pt}{0.0pt}\pgfsys@moveto{1.55379pt}{0.0pt}\pgfsys@curveto{1.55379pt}{1.5715pt}{0.27985pt}{2.84544pt}{-1.29166pt}{2.84544pt}\pgfsys@curveto{-2.86316pt}{2.84544pt}{-4.1371pt}{1.5715pt}{-4.1371pt}{0.0pt}\pgfsys@curveto{-4.1371pt}{-1.5715pt}{-2.86316pt}{-2.84544pt}{-1.29166pt}{-2.84544pt}\pgfsys@curveto{0.27985pt}{-2.84544pt}{1.55379pt}{-1.5715pt}{1.55379pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{-1.29166pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}} {{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-6.1111pt}{5.00006pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{$\displaystyle\sum$}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{ {}{}{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}{\leavevmode\hbox to12.22pt{\vbox to10pt{\pgfpicture\makeatletter\hbox{\hskip 6.1111pt\lower-5.00006pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } {{}{}}{{}}{}{{{}}{}{}{}{}{}{}{}{}}{}\pgfsys@moveto{-1.29166pt}{0.0pt}\pgfsys@moveto{0.41544pt}{0.0pt}\pgfsys@curveto{0.41544pt}{0.94281pt}{-0.34885pt}{1.70709pt}{-1.29166pt}{1.70709pt}\pgfsys@curveto{-2.23447pt}{1.70709pt}{-2.99875pt}{0.94281pt}{-2.99875pt}{0.0pt}\pgfsys@curveto{-2.99875pt}{-0.94281pt}{-2.23447pt}{-1.70709pt}{-1.29166pt}{-1.70709pt}\pgfsys@curveto{-0.34885pt}{-1.70709pt}{0.41544pt}{-0.94281pt}{0.41544pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{-1.29166pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}} {{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-6.1111pt}{5.00006pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{$\textstyle\sum$}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{ {}{}{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}{\leavevmode\hbox to9.06pt{\vbox to7pt{\pgfpicture\makeatletter\hbox{\hskip 4.52776pt\lower-3.50005pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } {{}{}}{{}}{}{{{}}{}{}{}{}{}{}{}{}}{}\pgfsys@moveto{-0.86111pt}{0.0pt}\pgfsys@moveto{0.2768pt}{0.0pt}\pgfsys@curveto{0.2768pt}{0.62845pt}{-0.23267pt}{1.13791pt}{-0.86111pt}{1.13791pt}\pgfsys@curveto{-1.48956pt}{1.13791pt}{-1.99902pt}{0.62845pt}{-1.99902pt}{0.0pt}\pgfsys@curveto{-1.99902pt}{-0.62845pt}{-1.48956pt}{-1.13791pt}{-0.86111pt}{-1.13791pt}\pgfsys@curveto{-0.23267pt}{-1.13791pt}{0.2768pt}{-0.62845pt}{0.2768pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{-0.86111pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}} {{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-4.52776pt}{3.50005pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{$\scriptstyle\sum$}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{ {}{}{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}{\leavevmode\hbox to6.94pt{\vbox to5pt{\pgfpicture\makeatletter\hbox{\hskip 3.4722pt\lower-2.50003pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } {{}{}}{{}}{}{{{}}{}{}{}{}{}{}{}{}}{}\pgfsys@moveto{-0.64583pt}{0.0pt}\pgfsys@moveto{0.20772pt}{0.0pt}\pgfsys@curveto{0.20772pt}{0.4714pt}{-0.17442pt}{0.85355pt}{-0.64583pt}{0.85355pt}\pgfsys@curveto{-1.11723pt}{0.85355pt}{-1.49937pt}{0.4714pt}{-1.49937pt}{0.0pt}\pgfsys@curveto{-1.49937pt}{-0.4714pt}{-1.11723pt}{-0.85355pt}{-0.64583pt}{-0.85355pt}\pgfsys@curveto{-0.17442pt}{-0.85355pt}{0.20772pt}{-0.4714pt}{0.20772pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{-0.64583pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}} {{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.4722pt}{2.50003pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{$\scriptscriptstyle\sum$}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{ {}{}{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}}\displaylimits$ denotes concatenation, $\mathcal{S}(x,y)$ scales $x$ to $y$’s size. Meanwhile, the target $x$ is scaled to the corresponding shapes of predictions to construct a pyramid loss for feature refinements:

where $\mathcal{L}_{\text{model}}$ is the original non-hierarchical loss. The overall loss is computed through the following equation:

The hierarchical architecture is an improvement of our base architecture aiming at discovering and combining local fine details and global structures, both of which are important for high-quality segmentation and reconstruction. Therefore, we derive the hierarchical versions of SeM and AE, denoted as H-SeM and H-AE, respectively. A vertical feature fusion for DDPM is not recommended, because we predict noise instead of reconstructing the image in the reverse diffusion process. Noise usually doesn’t contain clear global structures, and scaling the noise map may cause severe loss of details. Fig. 4 gives an example of a segmentation model constructed with H-SeM and a U-shaped encoder using ConvBlock.

We construct segmentation models with SeM architecture and different encoders including ResNet [7], U-Net [3], CAtten-U [14], Swin-U [21] and Mamba-U [23], which are constructed with ResConvBlock, ConvBlock, ConvBlock plus MSA, SwinBlock and MambaBlock, respectively. In the above, U-Net, CAtten-U, Swin-U and Mamba-U are U-shaped encoders. The corresponding hierarchical models are constructed using H-SeM. Note that the huge GPU memory requirements prevent us from training models with CAtten-U on 3D image patches. Models trained on 2D and 3D image datasets contain 2 and 3 down-sampling layers, respectively. The number of middle layers is set to 2. We employ a hybrid loss function that combines Dice and cross-entropy loss for all segmentation models, which are trained using Adam [35] optimizer for 500 epochs. The learning rate starts from $3\times 10^{-4}$ and follows a linear decaying schedule.

We construct reconstruction models with AE architecture and different encoders including ResNet, U-Net, Atten-U, Swin-U, and Mamba-U. H-AE architecture is used to construct their hierarchical versions. We train all reconstruction models for 50 epochs with L1 loss. Other settings remain the same as segmentation models.

We construct generation models with DDPM architecture and use different encoders for noise prediction, which are ResNet, U-Net, Swin-U, and Mamba-U. We use noisy images as input conditions for MRI generation. All models are trained for 200 epochs with L2 loss. The sampling process is accelerated through [15]. For a given condition, we integrate 5 or 10 generated samples as the final result, denoted as DDPM (5) and DDPM (10). Other settings remain the same as reconstruction models.

All experiments about 2D segmentation and reconstruction are conducted on a server with 8 NVIDIA 3090 GPUs. Models for generation and 3D segmentation are trained on 4 NVIDIA A800 GPUs.