Paper Title: Explainable and Controllable Motion Curve Guided Cardiac Ultrasound Video Generation.







Cardiac Ultrasound Video Generation




 


1National-Regional Key Technology Engineering Laboratory for Medical Ultrasound, School of Biomedical Engineering, Medical School,
Shenzhen University, China

xinyang@szu.edu.cn; nidong@szu.edu.cn 

2Medical Ultrasound Image Computing (MUSIC) Lab, Shenzhen University, China

3Marshall Laboratory of Biomedical Engineering, Shenzhen University, China

4Faculty of Applied Sciences Macao Polytechnic University, Macao, China 

5Shenzhen RayShape Medical Technology Co., Ltd, China




J.Yu and R.Chen et al.


Explainable and Controllable Motion Curve Guided Cardiac Ultrasound Video Generation
    Junxuan Yu1,2,3Junxuan Yu and Rusi Chen contribute equally to this work. Rusi Chen1,2,3⋆ Yongsong Zhou1,2,3 Yanlin Chen1,2,3 Yaofei Duan4 Yuhao Huang1,2,3 Han Zhou1,2,3,5 Tan Tao4 Xin Yang1,2,3() Dong Ni1,2,3()
    August 5, 2024
======================================================================================================================================================================================================================



§ ABSTRACT 
Echocardiography video is a primary modality for diagnosing heart diseases, but the limited data poses challenges for both clinical teaching and machine learning training.
Recently, video generative models have emerged as a promising strategy to alleviate this issue.
However, previous methods often relied on holistic conditions during generation, hindering the flexible movement control over specific cardiac structures.
In this context, we propose an explainable and controllable method for echocardiography video generation, taking an initial frame and a motion curve as guidance. 
Our contributions are three-fold.
First, we extract motion information from each heart substructure to construct motion curves, enabling the diffusion model to synthesize customized echocardiography videos by modifying these curves.
Second, we propose the structure-to-motion alignment module, which can map semantic features onto motion curves across cardiac structures.
Third, The position-aware attention mechanism is designed to enhance video consistency utilizing Gaussian masks with structural position information.
Extensive experiments on three echocardiography datasets show that our method outperforms others regarding fidelity and consistency. The full code will be released at https://github.com/mlmi-2024-72/ECMhttps://github.com/mlmi-2024-72/ECM.







§ INTRODUCTION

Echocardiography is a primary method that relies on dynamic video to obtain structural information for clinical diagnoses <cit.>. 
However, training radiologists with diagnostic skills and establishing machine learning models both suffer from limitations on video resources.
Recently, video generation models have demonstrated a promising ability to solve this problem, owing to their powerful capability in modeling data distribution <cit.>.
Several studies about video generation based on specific conditions have been investigated.
The typical ones relied on canny edges or depth maps <cit.> extracted from additional videos as conditions. Nevertheless, these conditions are non-editable and lack motion-driving information.
In contrast, Shi et al. <cit.> proposed the Motion-I2V framework to predict dense optical flow and guide video generation, maintaining both spatial and motion consistency.
Wang et al. <cit.> then employed a self-tracking training method.
Specifically, they specified the box positions of the first and last frames along with motion paths, to control the movement of objects.
However, these paths are basic and coarse, inadequate for ultrasound video synthesis that demands an accurate representation of intricate structure motions.





In the field of echocardiography video generation, 
Zhou et al. <cit.> proposed the OnUVS framework to synthesize ultrasound videos by animating source images and leveraging motion information from driving video. 
Nevertheless, OnUVS faced challenges due to the anatomical structure gap between the source image and the driving video, making accurate motion control difficult.
In addition, some studies have mined the structural information of the heart to guide the movement of videos.
For instance, Reynaud et al. <cit.> developed a Generative Adversarial Network (GAN) capable of generating echocardiography videos corresponding to the left ventricular ejection fractions (LVEFs).
They further employed a cascade video diffusion model conditioned on randomly sampled frames, enhancing the synthesis quality of echocardiography video <cit.>.
Van et al. <cit.> utilized segmentation masks of end-diastolic (ED) frames as a condition to generate four-chamber heart videos.
However, both LVEFs and ED masks are relatively sparse conditions, leading to an imbalance of information between intricate motions and limited conditions. 
This sparsity poses challenges for effectively controlling fine-grained cardiac structures, thus limiting the ability to capture the full complexity of heart movements and dynamics.

To address the above issues, we propose an explainable and controllable motion curve guided video diffusion model (ECM) that can synthesize video guided by the initial frame and motion curves (Fig. <ref>). 
Our contributions are threefold: 
(1) We innovatively mine the motion of the echocardiography video to obtain motion curves, which fully reflect the movement of each cardiac structure. This easily controlled approach enables the customization of videos through the modification (scaling and replacing) of the initial motion curve.
(2) As the curve lacks category information of the structure, we propose a Structure-to-Motion alignment mechanism. This mechanism extracts the semantic features of each cardiac structure and maps them with motion curve features, aiming to align visual and motion information effectively.
(3) We design position-aware attention masks based on the position of the cardiac structure movement, effectively enhancing the motion consistency of each structure.
To the best of our knowledge, ECM is the first study to apply cardiac motion curve guidance in echocardiography video generation. Extensive experimental results show that the proposed ECM is a flexible, controllable, and reliable method.




§ METHODOLOGY




Fig. <ref> shows the framework of ECM. During the training process, ECM takes an original video as input. 
A pretrained variational auto-encoder (VAE) from Stable Diffusion (SD) <cit.> is then utilized to downsample the video into latent features. 
By gradually adding noise to the latent and then learning to denoise it, the model can obtain the generated video with a pretrained VAE decoder. 
Remarkably, we develop the structure-to-motion alignment module to match cardiac structures with motion, yielding aligned features that condition the ECM model. 
Additionally, we employ Gaussian masks for each structure using a Position-aware Attention Mechanism, incorporating these into the spatial layers of the diffusion model. 
During the testing process, ECM generates echocardiography videos using an initial frame and motion curves as input. Users can customize the videos by replacing the initial frame or adjusting (scaling or replacing) the motion curves, highlighting the controllability and interpretability of our method.




 §.§ Extraction of Motion Curves

Previous studies <cit.> faced challenges in controlling the motion of specific cardiac substructures since they relied solely on the single sparse condition (LVEFs), while our approach aims to provide fine-grained control over the motion of each cardiac substructure.
Therefore, we extract motion curves for the key structures (e.g., Left Ventricle, Left Atrium, Mitral Valve, etc.) as conditions.

As illustrated in Fig. <ref>, the process of extracting motion curves is as follows: 
(a) We employ a well-trained anatomy detector (average accuracy=85%) to identify each substructure in each frame of the echocardiogram video. 
(b) Subsequently, we utilize the pixel coordinates of each substructure's bounding box (bbox) as the basis for encoding the motion curves, which form the basis for encoding the motion curves, represented as 
f_c^m∈ℝ^B × N × C × ( 4 × 2),
where N represents the number of frames and C represents the categories of the substructure. Notably, any missed detected structures would be treated as learnable parameters, initialized by the network.
(c) Since cardiac motion is periodic, we employ Fourier transformation (FT) to transform the pixel coordinates into a high-dimensional feature representation, denoted as 
f_c^m∈ℝ^B × N × C × E, 
where E represents the dimensionality of the features obtained from FT. 
(d) Finally, the motion embedding is passed through several multi-layer perceptron (MLP) layers. Overall, these motion curve features can be formulated as f_c^m∈ℝ^B × N × C × 1024.



 


 §.§ Structure-to-Motion Alignment Module


Although the motion curves of the echocardiography videos are captured, it is challenging for the model to distinguish the relationship between the motion curves and the semantic information of each substructure.
Recently, GLIGEN <cit.> has demonstrated its effectiveness in combining caption and bbox information, enhancing visual-language understanding to enable fine-grained control over specific objects in natural images.
However, our preliminary experiments suggested that GLIGEN struggles to effectively represent and interpret texts related to cardiac structures in echocardiography.



Innovatively, we replace texts with the cardiac structure features, denoted as f_i^s. 
As shown in Fig. <ref>, to obtain f_i^s, 
the regions of interest (ROI) that correspond to the cardiac substructures in the initial frame are cropped by the well-pretrained detector as mentioned above. 
Then, a pretrained CLIP image encoder <cit.> is utilized to transform the ROIs into structural embedding features.
Then, these features pass through several MLP layers,
resulting in an output denoted as f_c^s∈ℝ^B × N × C × 1024, which has the same shape with the extracted motion curve features f_i^m mentioned in Sec. <ref>.
Consequently, structural features and their corresponding motion curve features from the same category are concatenated to create aligned features. The aligned features are formulated as:

    F_c^a = Concat(f_c^s, f_c^m).


Note that any undetected structures are replaced with general features from the dataset.
Next, the concatenated features pass through an additional MLP layer for further integration.
Finally, to introduce the aligned motion curve features to guide the generation of echocardiography video, we then mapped 
F_c^a to the intermediate spatial layers of the UNet<cit.> via a cross-attention layer implementing cross attention, formulated as:

    CrossAtten(Q, K, V) = softmax( QK^T/√(d)) · V,

where Q, K and V represent the query, key and value respectively in the attention mechanism. Here, we regard noise latent feature F_l^m∈ℝ^B × 64 × 64 as query and the structure-to-motion embedding features as key and value. 






 §.§ Position-aware Attention Mechanism

Enhancing the consistency of cardiac motion is crucial in the task of cardiac video generation. 
To address this, we aim to inject the positional information of the cardiac structure into the cross-attention mechanism. Specifically, we design Gaussian masks that are generated based on the position of the cardiac structures. 
The Gaussian masks are defined as:

    M_g(x, y) = 1/2πσ^2exp(-(x - μ_x)^2 + (y - μ_y)^2/2σ^2),

where x and y represent the spatial positions of the mask.
μ_x and μ_y denote the center positions of the Gaussian distribution, which are the coordinates of the four corner points of the detected bbox. 
σ is the standard deviation to control the distribution width. 
Here, we set σ to 10 pixels as default.
Subsequently, the Gaussian masks are resized to match the dimensions of the latent feature map. 
The Gaussian-weighted cross-attention mechanism is as follows:

    CrossAttenMask(Q, K, V, M_g) = softmax( QK^T/√(d_k)⊙M_g ) V,

where ⊙ is the element-wise multiplication. Overall, this position-aware attention mechanism effectively integrates positional information of cardiac structures, enhancing the consistency and realism of generated cardiac motion.





§ EXPERIMENTS


Datasets and Implementations. 
To assess the performance of ECM, we gathered data from three sources: two private dataset from multiple hospitals and the publicly available dataset named EchoNet-Dynamic <cit.>.
The in-house dataset includes 144 apical four-chamber (A4C) and 100 apical two-chamber (A2C) heart videos, and the public one comprises 10,030 labeled A4C echocardiography videos.
For both datasets, the videos were randomly split into training (90%) and testing (10%) sets. During training, 12-frame clips were randomly sampled from each video, with a sampling interval ranging from 1 to 4 frames. 
For testing, videos were truncated according to different sampling intervals. 
The input videos from the private datasets were resized to 256x256 pixels, while those from the EchoNet-Dynamic dataset were kept at their original resolution of 112x112 pixels.
All methods were implemented in PyTorch using an NVIDIA RTX 4090 GPU under same settings. The Adam optimizer was used with a learning rate of 5e-3 and 60,000 training steps.

Evaluation Metrics.
Our evaluation metrics cover both image-level and video-level evaluation.
The image-level assessment includes Structural Similarity Index (SSIM) <cit.>, Mean Absolute Error (MAE) <cit.>, Peak Signal-to-Noise Ratio (PSNR) <cit.>, Fréchet Inception Distance (FID) <cit.> and Learned Perceptual Image Patch Similarity (LPIPS) <cit.>.
For video-level assessment, we only considered the commonly-used Fréchet Video Distance (FVD) <cit.>.
Notably, to assess the consistency of cardiac structures between the generated and target videos, we introduce a new indicator calculated by the Intersection over Union (IoU) between the bboxes of the original and synthesized videos.




Method Comparison.
The quantitative comparison of ECM and other methods is reported in Table <ref>.
The ECM model achieves the best performance across all metrics, indicating superior image quality and high fidelity compared to other methods.
It can be seen that the SEG Diffusion method, generating echocardiography videos without any control conditions, achieves poor quality.
Notably, ECM model markedly outperforms the SD method under bbox/text conditions, achieving 67.8%↓ in FVD. 
This demonstrates that synthesizing 2D images with rich conditions and stitching them into videos does not yield satisfactory outcomes. 
The integration of a position-aware attention mechanism significantly contributes to our superior performance and enhanced video consistency.
Furthermore, the ECM model shows better performance than 3D ControlNet. 
This disparity arises from the fact that natural images typically have clearly defined control conditions, whereas the motion in echocardiography videos is inherently more complex.
Similarly, for the EchoNet-Dynamic dataset, guided by motion curves, our proposed ECM outperforms the strong competitor <cit.> across most metrics, indicating its superior generation capacity.










Ablation Study.
We conducted the ablation study to test the contribution of each component in Table <ref>. 
The ECM model consistently outperforms others across most metrics for both the A4C and A2C datasets, achieving outstanding FVD scores of 189.27 and 208.12, and impressive IoU scores of 0.766 and 0.828, respectively.
This highlights that ECM effectively enhances video generation consistency while maintaining strong image quality.
Additionally, incorporating motion curves instead of traditional textual control leads to notable improvements. 
Specifically, using motion curves, the FVD decreased from 307.50 to 259.77, and the IoU increased from 0.691 to 0.748.
This indicates that motion curves provide precise control over the motion in echocardiography videos.
It can also be observed that the Structure-to-Motion alignment module (`+S2M') and the position-aware attention mechanism (`+Att') enhance image quality and video motion consistency, respectively. Consequently, the final ECM model demonstrates excellent performance at both the image and video levels.



Qualitative Results. 
Fig. <ref> demonstrates that ECM generates videos closely matching the target in both the A4C and EchoNet-Dynamic datasets by inputting the initial frame and motion curves. It effectively mimics heart structure motions such as chamber dilation, contraction, and diastolic opening and systolic closing
Furthermore, Fig. <ref> (a) illustrates ECM's controllability. Starting from the same initial frame, the second-row video uses the original motion curves from target video A, while the third-row video uses replaced curves from target video B. The mitral valve's systole and diastole differ after replacing the motion curves (as shown by yellow arrows), indicating ECM's effective control over motion curves.
More visualization results, including generated videos with replaced and scaled motion curves, can be found at http://106.53.160.189:8000/anonymous link.




§ CONCLUSION

In this study, We have presented ECM for generating echocardiography videos guided by motion curves, which reflect the movement of each cardiac structure. ECM enables to customize the generated videos by adjusting (scaling and replacing) the initial motion curve. Besides, to link the structure features with corresponding movement information, we propose the structure-to-motion alignment mechanism. Moreover, attention masks based on the position of the anatomical structures are introduced to enhance the motion consistency of each structure. Overall, our proposed ECM achieves state-of-the-art performance for generating echocardiography videos in terms of fidelity and consistency.




  §.§.§ 

This work was supported by the grant from National Natural Science Foundation of China (12326619, 62101343, 62171290), Science and Technology Planning Project of Guangdong Province (2023A0505020002), Science and Technology Development Fund of Macao (0021/2022/AGJ), and Shenzhen-Hong Kong Joint Research Program (SGDX20201103095613036).



  §.§.§ 

The authors have no competing interests to declare that are relevant to the content of this article.


splncs04




 Paper Title: Goal-Oriented Semantic Communication for Wireless Image Transmission via Stable Diffusion.





Goal-Oriented Semantic Communication for Wireless Image Transmission via Stable Diffusion 
 
    Nan Li, Yansha Deng 
Department of Engineering, King's College London.

Email: {nan.3.li, yansha.deng}@kcl.ac.uk
    August, 2024
====================================================================================================================



§ ABSTRACT 
Efficient image transmission is essential for seamless communication and collaboration within the visually-driven digital landscape. To achieve low latency and high-quality image reconstruction over a bandwidth-constrained noisy wireless channel, we propose a stable diffusion (SD)-based goal-oriented semantic communication (GSC) framework. In this framework, we design a semantic autoencoder that effectively extracts semantic information from images to reduce the transmission data size while ensuring high-quality reconstruction. Recognizing the impact of wireless channel noise on semantic information transmission, we propose an SD-based denoiser for GSC (SD-GSC) conditional on instantaneous channel gain to remove the channel noise from the received noisy semantic information under known channel. For scenarios with unknown channel, we further propose a parallel SD denoiser for GSC (PSD-GSC)  to jointly learn the distribution of channel gains and denoise the received semantic information. Experimental results show that SD-GSC outperforms state-of-the-art ADJSCC and Latent-Diff DNSC, with the Peak Signal-to-Noise Ratio (PSNR) improvement by 7 dB and 5 dB, and the Fréchet Inception Distance (FID) reduction by 16 and 20, respectively. Additionally, PSD-GSC archives PSNR improvement of 2 dB and FID reduction of 6 compared to MMSE equalizer-enhanced SD-GSC.



Goal-oriented semantic communication, Stable diffusion model, and Image transmission.






§ INTRODUCTION

The growing reliance on visual applications in domains such as edge computing and vehicular networks highlights the importance of seamless, reliable, and fast sharing of visual information <cit.>. To achieve efficient image transmission over bandwidth-limited wireless networks, it is necessary to reduce the transmitted data size while ensuring high-quality image reconstruction <cit.>. Goal-oriented semantic communication (GSC) has emerged as a promising scheme that focuses on extracting and transmitting the important semantic information relevant to the goal rather than full bits <cit.>. By selectively transmitting the semantic information, GSC allows for a certain degree of pixel-level error tolerance and evaluates the perception of the reconstructed image, potentially leading to more efficient image transmission <cit.>. 

To enhance the efficiency of GSC system, deep learning-based joint source-channel coding (JSCC) schemes have attracted increasing attention. A JSCC-based semantic communication architecture has been developed in <cit.> for high-resolution image transmission over additive white Gaussian noise (AWGN) and fading channels. Subsequently, a multiple-description JSCC scheme was designed in <cit.> for bandwidth-agile image transmission. However, these schemes could not understand the semantic content of the transmitted images. To solve this, the work <cit.> designed a deep reinforcement learning-based semantic encoder and introduced a Generative Adversarial Nets-based semantic decoder with attention mechanism to achieve high-quality image reconstruction. 
Nonetheless, the aforementioned works assumed that the channel remains consistent during training and inference, which can cause severe performance degradation under practical channel conditions with dynamism. To address this issue, an attention mechanism-based deep JSCC (ADJSCC) scheme was proposed in <cit.>, which can dynamically adjust the SNR during training to adapt to changing channel conditions and remove channel noises. However, adaptive denoising in practical channel conditions with dynamism remains a major challenge. 

Generative models, e.g., Denoising Diffusion Probabilistic Models (DDPM) <cit.>, have demonstrated their capability in denoising noisy or imperfect input data across various generation tasks, such as image <cit.> and video generation <cit.>. Inspired by this, Latent-Diff DNSC was developed in <cit.> to characterize noisy channels for efficient image transmission, which demonstrates the effectiveness of DDPM in removing noise from received noisy semantic information. However, this scheme assumes perfect channel estimation for AWGN channels, which poses challenges for accurate denoising in practical channel conditions, especially in the presence of imperfect channel estimation.


To effectively combat the impact of wireless channel noise on semantic information transmission, we propose a stable diffusion (SD)-based GSC framework. The main contributions are summarized in the following:

    
  * We propose an SD-based GSC framework with a semantic encoder at the transmitter, and a semantic denoiser and a semantic decoder at the receiver. The semantic encoder extracts semantic information from the image; the semantic denoiser denoises the received noisy semantic information; the semantic decoder reconstructs the image from the denoised semantic information.
    
  * We propose an SD-based semantic denoiser for GSC (SD-GSC) that treats instantaneous channel gain as an input condition of SD to identify and remove noise under known channel. For unknown channel, we propose a parallel SD-based semantic denoiser for GSC (PSD-GSC) that jointly estimates the instantaneous channel gain and denoises the received noisy semantic information.
    
  * We conduct extensive experiments to examine the Peak SNR (PSNR) and Fréchet Inception Distance (FID). Compared to state-of-the-art ADJSCC and Latent-Diff DNSC, SD-GSC improves the PSNR by 7 dB and 5 dB and reduces the FID by 16 and 20, respectively. Additionally, PSD-GSC improves the PSNR by 2 dB and reduces the FID by 6 compared to the MMSE equalizer-enhanced SD-GSC.







§ SYSTEM MODEL AND PROBLEM FORMULATION

We consider an image transmission task over Rayleigh fading channels. To achieve efficient image transmission, we introduce a stable diffusion (SD)-based goal-oriented semantic communication (GSC) framework, that consists of a semantic encoder at the transmitter and a semantic denoiser and a semantic decoder at the receiver, as shown in Fig. <ref>. The semantic encoder extracts the semantic information from the input image, and the semantic decoder then reconstructs the image from the received semantic information. Importantly, the semantic denoiser aims to remove the noise from the received noisy semantic information to enhance the reconstruction capacity of the semantic autoencoder. Specifically, we propose an SD denoiser under known channel, and we further introduce two alternative denoisers for scenarios with unknown channel:

    
  * MMSE equalizer-enhanced SD (MSD) denoiser: MMSE equalizer is used to estimate the instantaneous channel gain, which is then fed into the SD denoiser,
    
  * Parallel SD (PSD) denoiser: Two parallel SD modules are used to jointly learn the instantaneous channel gain and denoise the channel noise.

The details of each module are presented one by one below.

At the transmitter, the input data x∈ℝ^H × W × C is fed into the semantic encoder to extract the semantic information, thereby yielding a latent semantic vector z∈ℝ^d, 

    z = ℰ (x),

where ℰ(·) is the operation of semantic encoder. 

Subsequently, the semantic information z is transmitted through a wireless channel, experiencing transmission impairments such as distortion and noise. The received noisy semantic information is mathematically represented as

    z' = h * z + n,

where h denotes the channel gain between the transmitter and receiver; and n∼𝒞𝒩(0,σ^2) represents the additive white Gaussian noise (AWGN) with the noise power of σ^2. 

At the receiver, the noisy semantic information z' is progressively denoised using the semantic denoiser to obtain a denoised semantic vector as

    z̃ = ℱ_ϵ_θ(z'),

where ℱ_ϵ_θ(·) is the denosing operation with parameter ϵ_θ. 

The denoised semantic vector z̃ is then fed into the semantic decoder to generate the recovered image data x̃∈ℝ^H × W × C as 

    x̃ =  𝒟(z̃),

where 𝒟(·) is the operation of the semantic encoder.

With the aim to ensure the fidelity of image reconstruction, we introduce the mean squared error (MSE) to measure the difference between the transmitted image at the transmitter and the recovered image at the receiver, and formulate the optimization problem as 

    min1/N∑_i=1^Nx_i - x̃_i^2,

where N is the size of the image dataset 𝒳.




§ ARCHITECTURE OF SD-BASED GSC

In this section, we present our design of semantic autoencoder and semantic denoiser under perfect and imperfect channel estimation, respectively.


 §.§ Semantic Autoencoder

We design a deep neural network-based semantic autoencoder that can encode input image into a low-dimensional latent semantic space to extract important semantic information. The detailed neural network structure of our designed semantic autoencoder is present in Table <ref>. 

To better reconstruct the image, we incorporate MSE and Kullback-Leibler (KL) divergence into the training process of the semantic autoencoder. The MSE facilitates accurate image reconstruction, while the KL divergence serves as a regularization term, guiding the data distribution in the latent space towards approximating a unit Gaussian distribution. The loss function of the semantic autoencoder is defined as the sum of MSE and KL divergence via

    ℒ = 1/N∑_i=1^Nx_i - x̃_i^2 + λℒ_KL,

where λ is to control the importance of KL divergence with respect to the MSE, and the KL divergence is denoted as 

    ℒ_KL = -1/2∑_i=1^N (1 + log(σ_i^2) - μ_i^2 - σ_i^2),

where μ_i and σ_i represent the mean and standard deviation of the latent space, respectively.
0.05in


 §.§ SD Denoiser under Known Channel

To handle the wireless channel noise during semantic information transmission, Latent-Diff DNSC used DDPM for image generation. However, the uncontrolled generation process of DDPM may lead to a decrease in image fidelity during reconstruction <cit.>. Furthermore, the normalization of the received semantic information in Latent-Diff DNSC cannot always guarantee that the normalized data adhere to the standard Gaussian distribution, potentially degrading the reconstruction quality. To mitigate these challenges, we propose an SD Denoiser that uses the received noisy semantic information and instantaneous channel gain as control conditions to effectively remove channel noise, as shown in Fig. <ref>. Note that SD and DDPM share the same forward process.


  §.§.§ Forward Process

The forward process (i.e., training process) is executed by iteratively adding Gaussian noise to the initial distribution of semantic information z_0 ∼ p(z) over T time steps to gradually approach an isotropic Gaussian distribution z_T ∼𝒩(0, 𝐈). In DDPM <cit.>, at time step t ∈ [0, T], the forward process of z_t is expressed as the following noisy process:

    z_t=√(α̅_t)z_0+ √((1-α̅_t))ϵ,

where z_0 ∼ p(z) is the  initial distribution
of semantic information, α_t=1-β_t and α̅_t=∏_i=1^t(1-α_i), and β_t is the noise scheduling function and is typically modeled as a monotonically increasing linear function of time step t, and ϵ∼𝒩(0, 𝐈).  

From the score-based perspective (i.e., the gradient of the log probability density with respect to the data z_t at each noise scale β_t), the forward process can be expressed as the following linear stochastic differentiable equation (SDE) <cit.>:

    d^(f)z_t = - β_t/2z_t dt + √(β_t)dw,

where w is the d-dimensional Wiener process. 


  §.§.§ Reverse Process

The reverse process (i.e., inference process) is to recover the semantic information z_0 starting from the sample z_T. 

DDPM: According to (<ref>), the reverse process of DDPM can be achieved by the corresponding reverse SDE as <cit.>

    d^(r)z_t =  [ - β_t/2z_t - β_t ∇_z_tlog p (z_t)  ] d t + √(β_t)dw,

where d t is the reverse time step; dw is the standard Wiener process in reverse; and log p_t (z_t) is the score function.

Since ∇_z_tlog p (z_t) = ∇_z_tlog p(z_t |z_0), we can approximate ∇_z_tlog p (z_t) ≃s_θ(z_t, t) for the reverse process in (<ref>) by solving the following minimization problem during the training in the forward process <cit.>: 

    θ^* =θargmin𝔼_z_t, z_0[s_θ(z_t, t)-∇_z_tlog p(z_t | z_0)_2^2],

where the trained score network s_θ(z_t, t) can be denoted by using Tweedie’s identity <cit.> as

    s_θ(z_t, t) = ∇_z_tlog p(z_t)  = - 1/√(1-α̅_t)ϵ_θ(z_t, t),

where ϵ_θ(z_t, t) is the learned noise estimator of semantic information with parameter θ at time step t.  



SD: Obviously, the above reverse process of DDPM from a random sample z_0 ∼ p(z) cannot ensure the fidelity of image reconstruction. To allow for more controlled and guided image generation, we designed an SD Denoiser conditional on the received noisy semantic information z' and instantaneous channel gain h. Therefore, it is straightforward to modify (<ref>) to derive the reverse process of SD from the posterior distribution



    d^(r)z_t     = [ - β_t/2z_t - β_t∇_z_tlog p(z' |z_t) ] d t + √(β_t)dw
        = [ - β_t/2z_t - β_t (∇_z_tlog p (z_t)  
       + ∇_z_tlog p(z'|z_t ,h) )] d t +√(β_t)dw.


By discretizing the reverse process in (<ref>), we have

    z_t-1 =    1/√(α_t)(z_t + β_t [s_θ(z_t, t) + ∇_z_tlog p(z'|z_t,h) ] ) 
        + √(β_t)𝒩(0, 𝐈),

where s_θ(z_t, t) is the trained score network in (<ref>).

To solve the reverse process in (<ref>), the main challenge lies in the posterior distribution p(z'|z_t,h). However, the relationship between z' and z_0 is known but the relationship between z_t and z' is unknown. To tackle this issue, we express p(z' |z_t) as 

    p(z' |z_t)=∫ p(z' |z_0 ) p(z_0 |z_t) dz_0.

where the mean of p(z_0|z_t) can be approximated by a delta function as

    p(z_0 |z_t) ≃δ_E[z_0 |z_t](z_0),


To estimate the E[z_0|z_t], we can use the well-trained noise estimator ϵ_θ(z_t, t) in the forward process (<ref>) to obtain the estimation E[z_0|z_t] = ẑ_0(t) as

    ẑ_0(t) =1/√(α_t)(z_t- √(1-α̅_t-1)ϵ_θ(z_t,t))
        .


Using (<ref>), the approximation p(z_0|z_t) leads to the following formula for the gradient of the log-likelihood:

    ∇_z_tlog p(z'|z_t,h) = - 1-α̅_t/(1 -α_t)(1-α̅_t-1)∇_z_t ||z'-h * z_t||^2.



To fasten the convergence of our proposed SD denoiser for GSC (SD-GSC), we use a two-stage training process. We first solely train the semantic autoencoder, then we freeze its parameters and feed the latent semantic vector z into the semantic denoiser for training. The training and inference process is detailed in Algorithm 1.




 §.§ PSD Denoiser under Unknown Channel

Notably, Algorithm 1 is only applicable when the instantaneous channel gain h is known, and hence cannot be directly used for the scenarios (e.g., Massive MIMO communications systems) with imperfect estimation of h. To solve this issue, we propose a parallel SD-GSC (PSD-GSC) to jointly estimate the channel gain and remove the noise. 


  §.§.§ Forward Process

Since z and h are independent, the posterior probability is given by 

    p(z,h|z') = p(z'|z,h) p(z) p(h).

Therefore, we can train two separate forward processes for z and h, respectively. Similar to (<ref>) and (<ref>), the score network of h and  E[h_0|h_t]= ĥ_0(t) can be estimated using

    s_ϑ(h_t, t) = ∇_h_tlog p(h_t)  = - 1/√(1-α̅_t)ϵ_ϑ(h_t,t),

and

    ĥ_0(t) =1/√(α_t)(h_t- √(1-α̅_t-1)ϵ_ϑ(z_t,t))
        ,

where ϵ_ϑ(h_t, t) is the learned noise estimator of channel gain with parameter θ at time step t.


  §.§.§ Reverse Process

Similar to the reverse process in (<ref>), the two reverse processes of PSD-GSC can be expressed as



    d^(r)z_t =    [ - β_t/2z_t - β_t (∇_z_tlog p (z_t) 
        + ∇_z_tlog p(z'|z_t, h_t) )] d t + √(β_t)dw,

and



    d^(r)h_t =    [ - β_t/2h_t - β_t (∇_h_tlog p (h_t) 
        + ∇_h_tlog p(z'|z_t, h_t) ) ] d t + √(β_t)dw.

Correspondingly, the above two reverse processes can be discretized in a similar way as (<ref>) and expressed as 



    z_t-1 =  1/√(α_t)(z_t +    β_t [∇_z_tlog p(z'|z_t,h_t)  
         +  s_θ(z_t, t)] ) 
         + √(β_t)𝒩(0, 𝐈),

and 



    h_t-1 =  1/√(α_t)(h_t +    β_t [∇_h_tlog p(z'|z_t,h_t)  
         +  s_ϑ(h_t, t)] ) 
         + √(β_t)𝒩(0, 𝐈).



Due to the sparse structure of wireless channel, we use ℓ_1 regularization to sparse the channel gain by augmenting the diffusion prior thereby 
better stabilize the reconstruction. The minimization strategy for the channel gain then becomes

    h'_t-1 = h_t-1 - α (||z'-h_t-1 * z_t-1||_2 + ϕ ||h_t-1||),

where ϕ is the regularization strength. 

The detailed process of our proposed PSD denoiser for GSC (PSD-GSC) under unknown channel is presented in Algorithm 2. 







§ PERFORMANCE EVALUATION




We implement the image transmission task under a Rayleigh fading channel, using the FFHQ dataset <cit.>, which consists of 70,000 high-quality human face images at a resolution of 256×256 pixels, that is created as a benchmark for generative tasks. In the implementation of the semantic autoencoder, the encoder downsamples the feature dimensions by half, starting from the 2nd block, with channel dimensions specified as [64, 128, 256, 512], and then the features are embedded into a latent representation of dimension d=4096. Conversely, the decoder upsamples the feature dimensions by double for the first 3 blocks, with channel dimensions set to be [512, 256, 128, 64]. For the semantic denoiser, we set the time step T=1000 and use a linear variance scheme to determine the hyperparameters of {β_1, β_2, ⋯, β_T}. 

Baselines: To demonstrate the effectiveness of our proposed frame, we perform a comparative analysis against the state-of-the-art baselines: ADJSCC <cit.> and Latent-Diff DNSC <cit.>. We also compare with MSD denoiser-based GSC (MSD-GSC) as a benchmark for channel gain estimation to highlight the effectiveness of our PSD-GSC in learning channel gains. 

Metrics: We evaluate two performances: peak signal-to-noise ratio (PSNR) to evaluate image reconstruction quality, with higher values indicating better reconstruction quality; and Frechet inception distance (FID) to assess the similarity between real and reconstructed image distributions, with lower values indicating better reconstruction quality. 




 §.§ Known Channel

Fig. <ref> presents the PSNR and the FID performance across various SNRs under known channel.
As expected, both PSNR and FID performances improve as the SNR increases, owing to the fact that better wireless channel conditions facilitate higher-quality image reconstruction. 
Interestingly, the Latent-Diff DNSC scheme outperforms the ADJSCC scheme in both PSNR and FID at all SNRs, due to diffusion model can efficiently capture the wireless channel characteristics and remove the noise from the received noisy semantic information, resulting in a more accurate reconstruction.
Importantly, our proposed SD-GSC outperforms two baseline schemes in both PSNR and FID metrics. Specifically, compared to Latent-Diff DNSC and ADJSCC, the PSNR of our proposed SD-GSC achieves approximately 5 dB and 7dB improvement, and the FID of our proposed SD-GSC achieves around 16 and 30 reduction, respectively. This is because integrating the instantaneous channel gain will enable better characterization of channel gain and mitigation of distortion due to channel noise, resulting in a more controlled and guided image
generation than the uncontrolled generation process of DDPM used in the Latent-Diff DNSC scheme. 


 §.§ Unkown Channel

To validate the effectiveness of our proposed PSD-GSC for channel estimation as well as efficient image transmission, we also compare our proposed PSD-GSC with MSD-GSC and SD-GSC.  
Fig. <ref> plots the PSNR and the FID for various SNRs under unknown channel. As expected, our proposed SD-GSC with known h outperforms PSD-GSC and MSD-GSC with unknown h in terms of both PSNR and FID, due to the errors in unknown channel estimation. Specifically, our proposed PSD-GSC and MSD-GSC experience degradation in PSNR by approximately 1 dB and 3 dB, and degradation in FID by 2 and 8, respectively. This is because channel estimation errors result in inaccurate channel gain inputs of the diffusion model, which decreases its denoising and high-quality image reconstruction capabilities. Importantly, PSD-GSC outperforms MSD-GSC  both in PSNR and FID. This is due to the probabilistic modeling-based diffusion model used in PSD-GSC can well estimate nonlinear channels and capture the complex properties of wireless channels, for better image reconstruction.






§ CONCLUSION

In this paper, we developed a stable diffusion (SD)-based goal-oriented semantic communication framework (SD-GSC) for efficient image transmission over Rayleigh fading channels. The main goal is to reduce the bandwidth requirement while ensuring high-quality image reconstruction at the receiver. To this end, we first designed a semantic autoencoder to extract semantic information to reduce the transmitted data size while accurately reconstructing the image from the received semantic information. To combat the negative impact of wireless channel noise on the transmitted semantic information, we designed an SD-based semantic denoiser to better characterize and mitigate the channel noise, leading to significant improvements in both PSNR and FID metrics compared to the state-of-the-art baselines under known channel. 
For unknown channel scenarios, we further designed a parallel SD-GSC (PSD-GSC) to jointly learn the channel gain and denoise the received semantic information, showcasing the capability of our proposed PSD-GSC in channel estimation and efficient image transmission.



IEEEtran




 Paper Title: Empowering Snapshot Compressive Imaging: Spatial-Spectral State Space Model with Across-Scanning and Local Enhancement.


Enhancing Ethereum Fraud Detection via Generative and Contrastive Self-supervision
    Chenxiang Jin, Jiajun Zhou^∗, Chenxuan Xie, Shanqing Yu, Qi Xuan, Senior Member, IEEE, Xiaoniu Yang


This work was supported in part by the Key R&D Program of Zhejiang under Grants 2022C01018 and 2024C01025, by the National Natural Science Foundation of China under Grants 62103374 and U21B2001. (Corresponding author: Jiajun Zhou.)
C. Jin, C. Xie, S. Yu and Q. Xuan are with the Institute of Cyberspace Security, College of Information Engineering, Zhejiang University of Technology, Hangzhou 310023, China, with the Binjiang Institute of Artificial Intelligence, ZJUT, Hangzhou 310056, China. E-mail:{jincxiang, 221122030330, yushanqing, xuanqi}@zjut.edu.cn.
J. Zhou are with the Institute of Cyberspace Security, College of Computer Science and Technology, Zhejiang University of Technology, Hangzhou 310023, China, with the Binjiang Institute of Artificial Intelligence, ZJUT, Hangzhou 310056, China. E-mail: jjzhou@zjut.edu.cn.


X. Yang is with the Institute of Cyberspace Security, Zhejiang University of Technology, Hangzhou 310023, China, and also with Science and Technology on Communication Information Security Control Laboratory, the 36th Research Institute of China Electronics Technology Group Corporation, Jiaxing 314033, China. Email: yxn2117@126.com.










    August 5, 2024
===================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================




§ ABSTRACT 
Snapshot Compressive Imaging (SCI) relies on decoding algorithms such as CNN or Transformer to reconstruct the hyperspectral image (HSI) from its compressed measurement. Although existing CNN and Transformer-based methods have proven effective, CNNs are limited by their inadequate modeling of long-range dependencies, while Transformer ones face high computational costs due to quadratic complexity.
Recent Mamba models have demonstrated superior performance over CNN and Transformer-based architectures in some visual tasks, but these models have not fully utilized the local similarities in both spatial and spectral dimensions.
Moreover, the long-sequence modeling capability of SSM may offer an advantage in processing the numerous spectral bands for HSI reconstruction, which has not yet been explored.
In this paper, we introduce a State Space Model with Across-Scanning and Local Enhancement, named ASLE-SSM, that employs a Spatial-Spectral SSM for global-local balanced context encoding and cross-channel interaction promoting. Specifically, we introduce local scanning in the spatial dimension to balance the global and local receptive fields, and then propose our across-scanning method based on spatial-spectral local cubes to leverage local similarities between adjacent spectral bands and pixels to guide the reconstruction process. These two scanning mechanisms extract the HSI's local features while balancing the global perspective without any additional costs. 
Experimental results illustrate ASLE-SSM's superiority over existing state-of-the-art methods, with an inference speed 2.4× faster than Transformer-based MST and saving 0.12 (M) of parameters, achieving the lowest computational cost and parameter count.







§ INTRODUCTION






r0.37
	

	
 
		
    < g r a p h i c s >

	

	

	PSNR-FLOPs comparisons of ASLE-SSM and previous End-to-End Transformer-based SOTA methods.
	

	
 


Hyperspectral images (HSIs) contain a large number of spectral bands to store rich detail information and thus are widely used in various fields such as object tracking <cit.>, remote sensing <cit.>,  super-resolution <cit.>, and medical image processing <cit.>.
Traditional imaging systems used for collecting HSIs involve scanning along the spatial or spectral dimensions but this imaging process is time-consuming and can not effectively capture dynamic objects. 
Recently, snapshot compressive imaging (SCI) systems have been developed <cit.> as an alternative method for obtaining HSIs at video rate. 
SCI systems capture multiple image bands in a single snapshot, compressing the information of HSIs into a 2D measurement. <cit.>. 
Among these systems, Coded Aperture Snapshot Spectral Imaging (CASSI) has emerged and become mainstream <cit.>. 
CASSI utilizes a physical mask to modulate signals from HSIs, shifting these modulated images to different spatial locations on the detector plane via a disperser <cit.>. 
However, the performance of the CASSI system heavily depends on the efficiency of the reconstruction algorithm.


In recent years, numerous reconstruction algorithms have been proposed, including model-based methods and deep learning-based methods.
Traditional model-based methods typically attempt to model the inverse process of image reconstruction, using manually designed priors to derive the solution to an optimization problem.
Although these methods are highly interpretable, they often require manual parameter tuning and tend to be slow. 
With the success of deep learning methods in various visual tasks, CNN and Transformer-based methods have been developed and applied to HSI reconstruction, achieving higher image reconstruction quality and speed than traditional algorithms. 
The CNN-based methods <cit.> have an advantage in extracting local features from images, but they lack the ability to model long-range dependencies effectively. 
Transformer-based methods <cit.> realize a global receptive field, resulting in superior performance compared to CNNs but higher computational costs. 
Although some mechanisms can mitigate this issue, they do not fundamentally change the quadratic complexity problem caused by the attention mechanism in the Transformer architecture. 



The limitation of Transformers in spectral SCI remains unresolved. However, in the realm of large language model and RGB image vision, an advanced solution for this problem exists, known as Mamba <cit.>.
Through data-dependent selective SSMs with parallel scanning and hardware-wise strategies, Mamba establishes a global receptive field and achieves near-linear complexity.
Mamba-based methods have been shown to outperform previous methods in various tasks due to their excellent long sequence modeling capabilities and lower computational costs <cit.>. 

Subsequently, many works have begun to explore the application and performance of the Mamba model in various visual tasks. Representative works, such as the visual state space model <cit.> and vision Mamba <cit.>, have demonstrated Mamba's superiority for image processing, particularly in obtaining fewer parameters and faster speed. However, most Mamba-based models only conduct the scanning process along the spatial dimension, lacking cross-channel interaction, despite the importance of spectral information for HSI reconstruction. 
Recently, some Mamba-based methods have made progress in developing an SSM mechanism along the channel (spectral) dimension for certain visual tasks <cit.>, achieving better results compared to scanning only along the spatial dimension. However, this scanning method does not consider the inherently local features of HSIs in both spatial and spectral dimensions.


In this paper, we introduce ASLE-SSM, a spatial-spectral SSM with across-scanning and local enhancement, to improve HSI reconstruction in SCI. 
Our ASLE-SSM aims to strike a balance between global and local information utilization in HSIs. 
ASLE-SSM comprises two main components: spatial SSM and spectral-spatial cross SSM. 
The spatial SSM employs a combination of global and local scanning methods, enhancing row-by-row scanning with both global and local modeling. 
Meanwhile, the spectral-spatial cross SSM incorporates an innovative across-scanning mechanism, leveraging spatial-spectral local cubes to enhance the module's ability to capture local details.
Overall, our contributions are summarized as follows:


    
  * We present a novel spectral-spatial state space model for HSI reconstruction, which can simultaneously integrate information from both spectral and spatial dimensions to adequately model spatial context and spectral features. 
   
    
  * We propose an across-scanning method that integrates local features from both spatial and spectral dimensions within a single scan. Additionally, we introduce a local scanning approach within the spatial SSM to further enhance local feature extraction.
   
    
  * Our network outperforms state-of-the-art (SOTA) methods on simulation and real datasets, delivering visually superior results. Importantly, it boasts the minimum parameter count, lowest computational cost, and fastest inference speed among current SOTA methods.
    








 
 


§ PRELIMINARIES AND RELATED WORKS

 
 


 §.§ CASSI System

 
 
CASSI consists of a hardware encoder and a software decoder. The encoder compresses a 3D data cube into a 2D snapshot measurement and the decoder reconstructs HSIs from the measurement. 
An imaging process CASSI system is illustrated in Fig. <ref>. The 3D HSI data cube is denoted as 𝐅∈ℝ^H × W × N_λ where H, W, and N_λ represent the height, width, and number of wavelengths of the 3D HSI, respectively. 
A physical mask 𝐌∈ℝ^H × W is used to modulate the 3D HSI 𝐅 as:
    
    𝐅^'(:,:,n_λ) = 𝐅^'(:,:,n_λ) ⊙𝐌,

where ⊙ represents the element-wise product and 𝐅^'∈ℝ^H × W × N_λ is the modulated 3D HSIs. 
After the modulated cube passes the disperser, 𝐅^' is tilted and considered to be sheared along the y-axis.
We use 𝐅^”∈ℝ^H × (W+d(N_λ-1)) × N_λ to denote the tilted cube and assume n_λ to be the reference wavelength. That is, image 𝐅^'(:,:,n_λ_c) is not sheared along the y-axis. Then we have: 
    
    𝐅^”(u,v,n_λ) = 𝐅^' (x,y + d(λ_n - λ_c),n_λ),

where (u, v) indicates the coordinate system on the detector plane, and λ_n is the wavelength of channel n_λ. Here,
d(λ_n - λ_c) signifies the spatial shifting for channel n_λ. Finally, the 2D compressed measurement 𝐘∈ℝ^H × [W+d(N_λ-1)] can be obtained by
    
    𝐘 = ∑_n_λ=1^N_λ𝐅^”(:,:,n_λ) + 𝐍,

where 𝐍∈ℝ^H × [W + d(N_λ-1)] represents the noise of CASSI caused by the sensing detector.

 
 


 §.§ State Space Models

 
 
    Recently, researchers have made significant progress in state space models (SSMs), leading to the emergence of various methods and architectures based on SSMs.
    Among them, the structured state space model (S4) <cit.> becomes a pioneer work for its advantage in modeling long-range dependencies. 
    S4 maps a 1D sequence x(t) ∈ℝ→ y(t) ∈ℝ by an implicit latent state h(t)∈ℝ^N and decodes it into an output sequence, which can be formulated as:
        
    h^'(t)=𝐀 h(t)+𝐁 x(t), 
        y(t)=𝐂 h(t)+𝐃 x(t),

    where N is the state size, A∈ℝ^N× N, B∈ℝ^N × 1, C∈ℝ^1× N, and D∈ℝ.
    After that, a discretization process is adopted to integrate Eq. (<ref>) into the deep learning algorithm, using the  zero-order hold (ZOH) rule, which is defined as follows:
        
    A   =  exp(ΔA),
    B   =(ΔA)^-1( exp(A)-I)·ΔB,

    where A and 𝐁 represent the discretized forms of the parameters 𝐀 and 𝐁, respectively. Δ is denoted as the timescale parameter to transform the continuous parameters 𝐀 and 𝐁 into discrete ones. 
    
     However, the constant dynamics in S4 prevent it from correctly selecting the context information. By introducing a selective mechanism, selective structured state space models (S6), also called Mamba <cit.> have been proposed to realize the data-dependent dynamics, thereby effectively filtering the context information and achieving a better performance in language and visual tasks.

 
 


 §.§ Mamba in Vision

 
 
    Motivated by Mamba's excellent performance in language tasks, recent research has begun exploring its application in vision tasks and has achieved promising results. 
    To address the unidirectional modeling and lack of positional awareness in the original Mamba model, 
    Vision Mamba (Vim) <cit.> incorporates the bidirectional SSMs for data-dependent global visual context modeling and position embeddings for location-aware visual recognition. 
    VMamba <cit.> introduces the Cross-Scan Module (CSM) to bridge the gap between 1-D array scanning and 2-D plain traversing. 
    However, Vim and VMamba have not explored the role of SSM in the channel dimension.
    
    More recently, some works have begun to apply the SSM mechanism to the channel dimension.
    CU-Mamba <cit.> employs a Channel SSM component to preserve channel correlation features. 
    S^2Mamba <cit.> designs an extra scanning module on the spectral dimension for retrieving semantic properties in HSI data by a bi-directional interaction between each band.
    However, these methods directly conduct the scanning process along the spectral dimension and do not utilize the local similarity of HSIs in both spatial and spectral dimensions.

 
 


§ THE PROPOSED ASLE-SSM METHOD

 
 


 §.§ Architecture of ASLE-SSM

 
 
The overall architecture of our ASLE-SSM is depicted in Fig. <ref>. 
ASLE-SSM adopts a U-shaped network with the spatial-spectral SSM block as the basic unit. 
The structure consists of the encoder, bottleneck, and decoder. Given the vectorized measurement 𝐲∈ℝ^H(W + d(N_λ - 1)), we reverse the dispersion and shift back 𝐘 to initialize the input signal 𝐗∈ℝ^H × W × N_λ:
    
    𝐗(x,y,n_λ) = 𝐘 (x,y - d(λ_n - λ_c)).

We first concatenate the input 𝐗_0 with the physical mask 𝐘∈ℝ^H × W × N_λ and then employ a 1×1 convolution layer to integrate information from the original signal and the physical mask. Moreover, the integrated input 𝐗_0 undergoes a 3×3 convolution embedding into a feature 𝐗_0^'∈ℝ^H × W × C. Subsequently, 𝐗_0^' passes through the downsample and upsample process. After passing through the decoder, we employ a 3×3 convolution mapping to obtain the residual information. By computing the sum of the integrated input and residual output, we finally obtain the reconstructed HSIs. In our implementation, we change the number of spatial-spectral SSM blocks in the encoder, bottleneck, and decoder to realize the small, middle, and large sizes of ASLE-SSM models. 



 
 


 §.§ Local Enhancement SSM (LE-SSM) Mechanism

 
 
    Although Mamba is effective in modeling long-range dependencies, directly using the SSM mechanism often leads to less effective in extracting local features. 
    Recent models tend to incorporate CNNs into the network structure, utilizing their local feature extraction capabilities to improve the performance of original SSMs. 
    However, such methods would increase the parameter count, diminishing the lightweight advantage of the original SSM mechanism. 
    While directly adding convolutional layers can enhance the model's performance, it does not fundamentally solve the problem.

    Inspired by LocalMamba <cit.>, we introduce a simplified local mechanism in the spectral dimension. 
    We use two scanning directions as the global branch, scanning the image row by row in both forward and backward directions while using the other two scanning directions as local branches to unfold and scan local patches. 
    The specific details of the scanning process are illustrated in Fig. <ref>(d). 
    Given the input 𝐗∈ℝ^H × W × C, we first generate four copies of input data with different arrangements.
    For the global branch, we directly flatten the feature map into 𝐗_1∈ℝ^L × C, where L =  H × W, and then reverse it to generate 𝐗_2. For the local branch, we first set the size of our patch flattening each patch into a 1D sequence, and then linking each patch together to generate 𝐗_3. 
    Similarly, we reverse 𝐗_3 to obtain 𝐗_4. 
    After the generation of input data in four directions, we conduct the process as follows:
        
    𝐗^' = Concat(𝐗_1,𝐗_2,𝐗_3,𝐗_4),
       𝐗̂ = LE-Spatial-SSM(𝐗^').

    The structure of LE-Spatial-SSM is shown in Fig. <ref>(b). 
    It is worth noting that our method does not employ any CNNs or spatial and channel attention mechanisms based on MLPs, ensuring that our method does not add any extra parameters or time consumption compared to the global cross-scan mechanism <cit.>.
    However, it also means we need to carefully choose the size of local patches, or else the reconstruction result will even become worse. 
    In our implementation, we use patches of size 4×4 to achieve a relatively better performance.

 
   


 §.§ Across-Scanning SSM (AS-SSM) Mechanism

 
 
Our LE-SSM mechanism establishes a local-enhanced global receptive field but does not fully exploit the rich spectral information inherent in HSIs, leading to suboptimal reconstruction.
Recently, some methods <cit.> have utilized and integrated channel information by directly scanning along the channel dimension (spectral dimension), but this scanning method does not consider the local features in both channel and spatial dimensions. 
In other words, neighboring pixels and bands often have certain similarities, but scanning the spectral dimension directly may lead to a distance increase between data points that are highly correlated. 

To address the above issues, we propose a local scanning method that integrates local features from both spatial and spectral dimensions. This mechanism attempts to exploit the local properties of HSIs in both spatial and spectral dimensions while leveraging the rich spectral information to guide the reconstruction process. 
First, we rearrange the input 𝐗∈ℝ^H × W × C into patches along the spatial context by dividing the input HSIs into a series of smaller data cubes 𝐒∈ℝ^P × P × C, which will be flattened into sequences for scanning. 
After that, we organize the scanning order within each sequence based on our spatial-spectral local cubes 𝐂∈ℝ^h × w × c.
As shown in Fig. <ref>(e), different colors indicate each spatial-spectral local cube. The scanning order of each data point is also depicted. The process can be formulated as:
    
    𝐂 = Divide(𝐗),
       𝐂^' = AS-Spectral-Spatial-SSM(𝐂),

where 𝐂^' denotes the output patch of AS-SSM. After the SSM mechanism, we reorganize the patches back into the same order as the input 𝐗. 
By scanning the data points within each data cube, we can leverage the local similarity between adjacent spectral bands and pixels, thus achieving better results than previous methods that directly scan along the spectral dimension.



 
 


§ EXPERIMENT





 §.§ Experimental Settings



The effectiveness of our ASLE-SSM network is verified on both simulation and real datasets, by using 28 wavelengths ranging from 450nm to 650nm, acquired via spectral interpolation manipulation.

Simulation HSI Data.
Two simulated HSI datasets CAVE <cit.> and KAIST <cit.> are utilized to conduct the simulation experiment. 
The CAVE dataset contains 32 HSIs with a spatial size of 512×512, while the KAIST dataset includes 30 HSIs with a spatial size of 2704× 3376. 
Following the experimental settings of TSA-Net <cit.>, we designate CAVE as the training set and select 10 scenes from the KAIST as the testing set.

Real HSI Data. 
We adopt the real HSIs obtained by the CASSI system, as developed in TSA-Net <cit.> for testing.

Evaluation Metrics. 
We select peak signal-to-noise ratio (PSNR) <cit.> and structural similarity (SSIM) <cit.> as the evaluation metrics to assess the quality of HSI reconstruction.

Implementation Details. 
Our ASLE-SSM is implemented using PyTorch and we conduct model training with Adam optimizer (β_1 = 0.9 and β_2 = 0.999) for 300 epochs.
The learning rate is set to 4×10^-4 at the beginning and is halved every 50 epochs during the training procedure. 
Image patches with size 256×256 are randomly cropped from the 3D HSI cubes with 28 channels for training in the simulation and real experiments. 
The shifting step d in the dispersion process is set to 2 and the batch size is 2. 
We apply the data augmentation techniques including random flipping and rotation. The models are trained on an RTX 4090 GPU.






 §.§ Simulation Results




r0.58


Comparisons of different frameworks on PSNR, SSIM, FLOPs, Params, and Inference Time on the test dataset.



 1.1

	0.55!
	
	
				color3Framework     MST-S <cit.>      CST-S <cit.>      ASLE-SSM-S    

				PSNR    34.14      34.53      34.74  

				SSIM     0.938       0.945      0.946  

				Params (M)     0.93       1.20     0.81   

				FLOPS (G)    12.96       11.67    7.01    

                    Inference Time(s)    0.022       0.041    0.009    

				
	
 



We conduct a comprehensive comparative analysis of our ASLE-SSM method and SOTA HSI reconstruction algorithms in Params, FLOPs, PSNR, and SSIM. 
We compared the results of ASLE-SSM with 12 SOTA methods, including two model-based methods (GAP-TV <cit.>, and DeSCI <cit.>), two CNN-based methods (λ-net <cit.>, TSA-Net <cit.>), three Transformer-based methods (HDNet <cit.>, MST-S, MST-M, MST-L <cit.>, and CST-S, CST-M, CST-L, CST-L+ <cit.>), two Plug-and-Play (PnP) methods (DIP-HSI <cit.>, DiffSCI <cit.>) and three deep unfolding methods (ADMM-Net <cit.>, GAP-Net <cit.>, DGSMP <cit.>) across 10 simulation scenes. The quantitative results of different models are listed in Table <ref>.

As shown in Table <ref>, our ASLE-SSM-L obtains 35.74dB of PSNR and 0.958 of SSIM, achieving the best reconstruction performance among SOTAs.
When compared to CNN-based and Transformer-based methods such as TSA-net <cit.>, HDNet <cit.>, MST-L <cit.>, and CST-L+ <cit.>, our ASLE-SSM-L demonstrates notable improvements over these methods, which are 4.39dB, 1.08dB, 0.67dB, and 0.03dB on average, respectively. The spectral density curves in Fig. <ref> also demonstrate our superiority.

We also conduct comparisons of our ASLE-SSM-S with previous Transformer-based models MST <cit.> and CST <cit.> in the small size.
As shown in Table <ref>, we not only achieve better performance in PSNR and SSIM compared to MST and CST but also obtain a more lightweight model with faster speed.
Specifically, the ASLE-SSM-S is 2.4 times faster than MST-S and 4.5 times faster than CST-S, while having 0.12 (M) fewer parameters than the MST-S, which means the smallest parameter count. From Fig. <ref>, we can also observe that our ASLE-SSM not only achieves better results, but the value of FLOPs is also significantly lower than that of the corresponding MST and CST models.


Except for quantitative comparisons, we also visualize the results of simulation experiments. 
We choose one of 28 channels from Scene 1 and Scene 2, against the simulation results obtained from twelve SOTA approaches. 
Visually, as illustrated in Fig. <ref>, our method yields smoother and cleaner details. 
The experimental results indicate the capability of our method to generate high-quality HSIs. 
Specifically, in Scene 1, our reconstruction results better generate the curvature of the chin compared to other models. 
In Scene 2, our model clearly reconstructs the shapes of the blocks and ropes. 
The heart-like pattern on the blocks is also well-preserved.





 §.§ Real Data Results



r0.6
    

    

        
    < g r a p h i c s >

         
        
    < g r a p h i c s >

    

    

    Spectral Density Curves.
    

    




We further apply our method to reconstruct real HSIs. We retrain our model on the CAVE <cit.> and KAIST <cit.> datasets, conducting tests on real scenes for experiments. Following the configurations established in MST <cit.>, We introduce 11-bit shot noise during training to simulate real measurement conditions. As shown in Fig. <ref>, we compare the reconstructed images from two real scenes (Scene 3 and Scene 4) using our ASLE-SSM and nine SOTA methods. Our model achieved superior results when compared to other models. Specifically, In Scene 3, our image preserves the most information compared to the results of other models. A noticeable protrusion consistent with the original image can be observed in the center of our result and other methods fail to reproduce. 

Additionally, our images do not exhibit significant artifacts or ghosting. In Scene 4, We can observe that the face shape in the original image is relatively rounded. However, the reconstruction results of previous methods generate an almond-shaped face similar to the compressed measurements, indicating that these methods did not learn this information. In contrast, our results clearly show a fuller face on both sides. Additionally, previous reconstruction methods often only reconstruct one side of the mouth, while our image noticeably reconstructs the whole mouth, though there is a slight misalignment between the two sides. Furthermore, our images do not exhibit nonexistent features that appear in previous methods' results, such as noses, wrinkles, or gaps that are not present in the original RGB image. These results demonstrate that our ASLE-SSM surpasses other methods in reconstructing image details and suppressing noise in real scenes.













 §.§ Ablation Study



Break-down Ablation. We conduct an ablation study by removing the LE-Spatial Mechanism and AS-Spectral-Spatial Mechanism from the ASLE-SSM network. The experiment results are shown in Table <ref>. Our goal is to explore the contributions of each mechanism to the results. The baseline barely adopts the U-shaped framework without any SSM modules. 
r0.45

Ablation study of ASLE-SSM's components



 1.1

	0.43!
	
	 
                    color3 Architecture     Params (M)     PSNR     SSIM  

                    Baseline w/o SSM 
                        0.18             27.41            0.678  

                    with LE-SSM                                     0.62     34.48     0.942   

                with AS-SSM                                        0.37            31.16          0.863    

                    with LE + AS SSM                            0.81             34.74         0.946    

                    
 


When we respectively employ LE-SSM and AS-SSM blocks, our model achieves 7.07 dB and 3.75 dB improvements in PSNR over the baseline, respectively. Employing both mechanisms simultaneously can further enhance the model's performance. Here we need to point out that, in consideration of maintaining lightweight, we only utilized a single-direction scan in AS-SSM. In fact, we could also deploy a four-direction scanning similar to LE-SSM. However, each scan direction implies additional computational costs so we decided not to adopt this way.

r0.45


Comparisons of AS mechanism and common scanning.



 1.1

	0.42!
	
	 
                    color3 Architecture                     PSNR     SSIM  

                    w/o Spectral SSM 
                     34.48            0.942  

                    with AS Mechanism                                         34.74         0.946    

                with Common Scanning                                 34.55         0.943    

                    
 


Across-Scanning Mechanism. To demonstrate the advantages of our across-scanning mechanism compared to common spectral scanning, we conduct the experiment on model performance under both methods. As illustrated in Table <ref>, when we adopt common scanning, it yields a 0.07 dB improvement. In contrast, when we apply the proposed across-scanning mechanism, our method achieves a significant increase of 0.26 dB compared to not using spectral SSM. Both methods were experimented with under the condition of a single scan. The superiority of our method mainly comes from paying more attention to adjacent bands and pixels, thus better extracting the local features of HSIs along both spatial and spectral dimensions. The results verify the role of our AS mechanism.


r0.45


Comparisons of LE mechanism and global scanning.



 1.1

	0.38!
	
	 
                    color3 Architecture                     PSNR     SSIM  

                    w/o Spatial SSM 
                     31.16           0.863  

                    with LE Machenism                                         34.48         0.942    

                with Global Scanning                                 34.40         0.940    

                    
 



Local-Enhancement Mechanism. Additionally, we conduct an ablation study of our LE mechanism on the CAVE and KAIST datasets to assess its contributions. As shown in Table <ref>, The evaluations begin with the model without SSM in the spatial dimension, which served as our initial baseline with a PSNR of 31.16 dB. Specifically, the addition of the global four-direction scanning along the spatial context yields a significant increase of 3.32 dB in PSNR and 0.077 in SSIM, demonstrating its effectiveness in modeling the global spatial context. 

We further apply our LE Machenism along the spatial dimension and achieve a 0.08 dB increase compared to the global scanning. our method only changes the scanning order and does not introduce any additional parameter consumption or computational cost to achieve a relatively better performance. In other words, we do not use CNN layers or spatial attention mechanisms as in other methods <cit.>. In fact, Without increasing any additional cost and with the reconstructed results already achieving a relatively good level, 0.08 dB in PSNR and 0.002 in SSIM are significant improvements. The ablation experiment demonstrates the importance of local features in HSIs and verifies the effectiveness of our LE Mechanism.






§ CONCLUSION



In this paper, we address two limitations of previous SSM-based models by our proposed mechanisms: insufficient utilization of channel (spectral) information, and limited local detail extraction capability. By designing a novel across-scanning mechanism, we fully utilize the rich spectral information in HSIs to guide the image reconstruction process. At the same time, our scanning method based on patched data cubes pays more attention to the local similarity of HSIs between adjacent bands and pixels, achieving better results compared to simply scanning along the channel dimension. By introducing local enhancement in the Spatial SSM, we integrate spatial context from both global and local perspectives, achieving better results than using a global SSM alone.

Despite the significant progress of our method, we also acknowledge that as the first framework to adopt the SSM mechanism in HSI reconstruction, our network still has room for further improvement. For example, we could design a lightweight and effective way to fuse local and global branch information in the spatial dimension, rather than simply adding them together. However, the key challenge lies not in improving the effect itself, but in controlling computational costs. We will explore this aspect in the future with the aim of designing a more efficient network architecture.



10

ot_1
M. H. Kim, T. A. Harvey, D. S. Kittle, H. Rushmeier, R. O. Prum J. Dorsey, and D. J. Brady.
3d imaging spectroscopy for measuring hyperspectral patterns on solid objects.
ACM Transactions on on Graphics, 2012.

ot_2
Zhihong Pan, Glenn Healey, Manish Prasad, and Bruce Tromberg.
Face recognition in hyperspectral images.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 25(12):1552–1560, 2003.

ot_3
Hien Van Nguyen, Amit Banerjee, and Rama Chellappa.
Tracking via object reflectance using a hyperspectral video camera.
In IEEE Computer Society Conference on Computer Vision and Pattern Recognition-Workshops, pages 44–51. IEEE, 2010.

chen2017denoising
Yongyong Chen, Yanwen Guo, Yongli Wang, Dong Wang, Chong Peng, and Guoping He.
Denoising of hyperspectral images using nonconvex low rank matrix approximation.
IEEE Transactions on Geoscience and Remote Sensing, 55(9):5366–5380, 2017.

he2020non
Wei He, Quanming Yao, Chao Li, Naoto Yokoya, Qibin Zhao, Hongyan Zhang, and Liangpei Zhang.
Non-local meets global: An iterative paradigm for hyperspectral image restoration.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(4):2089–2107, 2020.

yuan2018hyperspectral
Qiangqiang Yuan, Qiang Zhang, Jie Li, Huanfeng Shen, and Liangpei Zhang.
Hyperspectral image denoising employing a spatial–spectral deep residual convolutional neural network.
IEEE Transactions on Geoscience and Remote Sensing, 57(2):1205–1218, 2018.

dong2021model
Weisheng Dong, Chen Zhou, Fangfang Wu, Jinjian Wu, Guangming Shi, and Xin Li.
Model-guided deep hyperspectral image super-resolution.
IEEE Transactions on Image Processing, 30:5754–5768, 2021.

sr_2
Hongyuan Wang, Lizhi Wang, Chang Chen, Xue Hu, Fenglong Song, and Hua Huang.
Learning spectral-wise correlation for spectral super-resolution: Where similarity meets particularity.
In Proceedings of the 31st ACM International Conference on Multimedia, pages 7676–7685, 2023.

sr_3
Xian-Hua Han and Yen-Wei Chen.
Deep residual network of spectral and spatial fusion for hyperspectral image super-resolution.
In 2019 IEEE Fifth International Conference on Multimedia Big Data (BigMM), pages 266–270. IEEE, 2019.

mi_1
V Backman, Michael B Wallace, LT Perelman, JT Arendt, R Gurjar, MG Müller, Q Zhang, G Zonios, E Kline, T McGillican, et al.
Detection of preinvasive cancer cells.
Nature, 406(6791):35–36, 2000.

mi_2
Guolan Lu and Baowei Fei.
Medical hyperspectral imaging: a review.
Journal of Biomedical Optics, 2014.

mi_3
Ziyi Meng, Mu Qiao, Jiawei Ma, Zhenming Yu, Kun Xu, and Xin Yuan.
Snapshot multispectral endomicroscopy.
Optics Letters, 2020.

sci_1
Patrick Llull, Xuejun Liao, Xin Yuan, Jianbo Yang, David Kittle, Lawrence Carin, Guillermo Sapiro, and David J Brady.
Coded aperture compressive temporal imaging.
Optics Express, 2013.

sci_2
Ashwin Wagadarikar, Renu John, Rebecca Willett, and David Brady.
Single disperser design for coded aperture snapshot spectral imaging.
Applied Optics, 2008.

sci_3
Ashwin A Wagadarikar, Nikos P Pitsianis, Xiaobai Sun, and David J Brady.
Video rate spectral imaging using a coded aperture snapshot spectral imager.
Optics Express, 2009.

sci_4
Xin Yuan, Tsung-Han Tsai, Ruoyu Zhu, Patrick Llull, David Brady, and Lawrence Carin.
Compressive hyperspectral imaging with side information.
IEEE Journal of selected topics in Signal Processing, 9(6):964–976, 2015.

sci_5
Xun Cao, Tao Yue, Xing Lin, Stephen Lin, Xin Yuan, Qionghai Dai, Lawrence Carin, and David J Brady.
Computational snapshot multispectral cameras: Toward dynamic capture of the spectral world.
IEEE Signal Processing Magazine, 33(5):95–108, 2016.

sci_6
Hao Du, Xin Tong, Xun Cao, and Stephen Lin.
A prism-based system for multispectral video acquisition.
In 2009 IEEE 12th International Conference on Computer Vision, pages 175–182. IEEE, 2009.

gap_tv
Xin Yuan.
Generalized alternating projection based total variation minimization for compressive sensing.
In 2016 IEEE International conference on image processing (ICIP), pages 2539–2543. IEEE, 2016.

icme_sci_1
Yangke Ying, Jin Wang, Yunhui Shi, and Baocai Yin.
Dual-domain feature learning and memory-enhanced unfolding network for spectral compressive imaging.
In 2023 IEEE International Conference on Multimedia and Expo (ICME), pages 1589–1594. IEEE, 2023.

mst
Yuanhao Cai, Jing Lin, Xiaowan Hu, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, and Luc Van Gool.
Mask-guided spectral-wise transformer for efficient hyperspectral image reconstruction.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 17502–17511, 2022.

icme_sci_2
Shumian Yang, Xinxin Xiang, Fenghua Tong, Dawei Zhao, and Xin Li.
Image compressed sensing using multi-scale characteristic residual learning.
In 2023 IEEE International Conference on Multimedia and Expo (ICME), pages 1595–1600. IEEE, 2023.

lambda
Xin Miao, Xin Yuan, Yunchen Pu, and Vassilis Athitsos.
λ-net: Reconstruct hyperspectral images from a snapshot measurement.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4059–4069, 2019.

tsa_net
Ziyi Meng, Jiawei Ma, and Xin Yuan.
End-to-end low cost compressive spectral imaging with spatial-spectral self-attention.
In Proceedings of the European Conference on Computer Vision, pages 187–204. Springer, 2020.

mst_pp
Yuanhao Cai, Jing Lin, Zudi Lin, Haoqian Wang, Yulun Zhang, Hanspeter Pfister, Radu Timofte, and Luc Van Gool.
Mst++: Multi-stage spectral-wise transformer for efficient spectral reconstruction.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 745–755, 2022.

dauhst
Yuanhao Cai, Jing Lin, Haoqian Wang, Xin Yuan, Henghui Ding, Yulun Zhang, Radu Timofte, and Luc V Gool.
Degradation-aware unfolding half-shuffle transformer for spectral compressive imaging.
In Proceedings of the Advances in Neural Information Processing Systems, volume 35, pages 37749–37761, 2022.

cst
Yuanhao Cai, Jing Lin, Xiaowan Hu, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, and Luc Van Gool.
Coarse-to-fine sparse transformer for hyperspectral image reconstruction.
In Proceedings of the European Conference on Computer Vision, pages 686–704. Springer, 2022.

gu2023mamba
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
arXiv preprint arXiv:2312.00752, 2023.

mamba_survey
Xiao Wang, Shiao Wang, Yuhe Ding, Yuehang Li, Wentao Wu, Yao Rong, Weizhe Kong, Ju Huang, Shihao Li, Haoxiang Yang, et al.
State space model for new-generation network alternative to transformers: A survey.
arXiv preprint arXiv:2404.09516, 2024.

vmamba2
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang.
Vision mamba: Efficient visual representation learning with bidirectional state space model.
arXiv preprint arXiv:2401.09417, 2024.

mambair
Hang Guo, Jinmin Li, Tao Dai, Zhihao Ouyang, Xudong Ren, and Shu-Tao Xia.
Mambair: A simple baseline for image restoration with state-space model.
arXiv preprint arXiv:2402.15648, 2024.

lightm-unet
Weibin Liao, Yinghao Zhu, Xinyuan Wang, Cehngwei Pan, Yasha Wang, and Liantao Ma.
Lightm-unet: Mamba assists in lightweight unet for medical image segmentation.
arXiv preprint arXiv:2403.05246, 2024.

shi2024vmambair
Yuan Shi, Bin Xia, Xiaoyu Jin, Xing Wang, Tianyu Zhao, Xin Xia, Xuefeng Xiao, and Wenming Yang.
Vmambair: Visual state space model for image restoration.
arXiv preprint arXiv:2403.11423, 2024.

vision_mamba
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang.
Vision mamba: Efficient visual representation learning with bidirectional state space model.
arXiv preprint arXiv:2401.09417, 2024.

s2mamba
Guanchun Wang, Xiangrong Zhang, Zelin Peng, Tianyang Zhang, Xiuping Jia, and Licheng Jiao.
S2mamba: A spatial-spectral state space model for hyperspectral image classification.
arXiv preprint arXiv:2404.18213, 2024.

huang2024spectral
Lingbo Huang, Yushi Chen, and Xin He.
Spectral-spatial mamba for hyperspectral image classification.
arXiv preprint arXiv:2404.18401, 2024.

s4gu2021
Albert Gu, Karan Goel, and Christopher Re.
Efficiently modeling long sequences with structured state spaces.
In ICLR, 2021.

mambagu2023mamba
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
arXiv preprint arXiv:2312.00752, 2023.

liu2024vmamba
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu.
Vmamba: Visual state space model.
arXiv preprint arXiv:2401.10166, 2024.

cumamba
Gu T. Deng R.
Cu-mamba: Selective state space models with channel learning for image restoration.
arXiv preprint arXiv:2404.11778, 2024.

desci
Yang Liu, Xin Yuan, Jinli Suo, David J Brady, and Qionghai Dai.
Rank minimization for snapshot compressive imaging.
IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(12):2990–3006, 2018.

hdnet
Xiaowan Hu, Yuanhao Cai, Jing Lin, Haoqian Wang, Xin Yuan, Yulun Zhang, Radu Timofte, and Luc Van Gool.
HDNet: High-resolution dual-domain learning for spectral compressive imaging.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 17542–17551, 2022.

self
Ziyi Meng, Zhenming Yu, Kun Xu, and Xin Yuan.
Self-supervised neural networks for spectral snapshot compressive imaging.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2622–2631, 2021.

diffsci
Zhenghao Pan, Haijin Zeng, Jiezhang Cao, Kai Zhang, and Yongyong Chen.
Diffsci: Zero-shot snapshot compressive imaging via iterative spectral diffusion model.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2024.

admm-net
Jiawei Ma, Xiao-Yang Liu, Zheng Shou, and Xin Yuan.
Deep tensor admm-net for snapshot compressive imaging.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10223–10232, 2019.

gsm
Tao Huang, Weisheng Dong, Xin Yuan, Jinjian Wu, and Guangming Shi.
Deep gaussian scale mixture prior for spectral compressive imaging.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16216–16225, 2021.

gapnet
Ziyi Meng, Xin Yuan, and Shirin Jalali.
Deep unfolding for snapshot compressive imaging.
International Journal of Computer Vision, 131(11):2933–2958, 2023.

localmamba
Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu.
Localmamba: Visual state space model with windowed selective scan.
arXiv preprint arXiv:2403.09338, 2024.

cave
Jong-Il Park, Moon-Hyun Lee, Michael D Grossberg, and Shree K Nayar.
Multispectral imaging using multiplexed illumination.
In 2007 IEEE 11th International Conference on Computer Vision, pages 1–8. IEEE, 2007.

kaist
Inchang Choi, MH Kim, D Gutierrez, DS Jeon, and G Nam.
High-quality hyperspectral reconstruction using a spectral prior.
In Technical report, 2017.

PSNR
Luen C. Chan and Peter Whiteman.
Hardware-constrained hybrid coding of video imagery.
IEEE Transactions on Aerospace and Electronic Systems, AES-19(1):71–84, 1983.

ssim
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.
Image quality assessment: from error visibility to structural similarity.
IEEE Transactions on Image Processing, 13(4):600–612, 2004.






 Paper Title: Identity-Consistent Diffusion Network for Grading Knee Osteoarthritis Progression in Radiographic Imaging.





 



IC-RDN for Grading Knee Osteoarthritis Progression






W. Wu et al.




The University of Sydney, Darlington NSW 2008, Australia 

{wenhua.wu, kun.hu, wenxi.yue, weiwilson.li, milena.simic, zhiyong.wang}@sydney.edu.au Sydney Polytechnic Institute Pty Ltd, Sydney, NSW, Australia

chris@ruddergroup.com.au School of Computing, Engineering & Mathematical Sciences, La Trobe University

w.xiang@latrobe.edu.au
Identity-Consistent Diffusion Network for Grading Knee Osteoarthritis Progression in Radiographic Imaging
    Wenhua Wu10000-0001-7488-5487 Kun Hu1,Corresponding author. 0000-0002-6891-8059 
Wenxi Yue 10000-0002-8270-417X Wei Li 10000-0003-4731-3226 Milena Simic 1 Changyang Li 20000-0001-9400-3998 Wei Xiang 30000-0001-9400-3998 Zhiyong Wang 10000-0002-8043-0312 

    August 5, 2024
===================================================================================================================================================================================================================================================================



§ ABSTRACT 
Knee osteoarthritis (KOA), a common form of arthritis that causes physical disability, has become increasingly prevalent in society. Employing computer-aided techniques to automatically assess the severity and progression of KOA can greatly benefit KOA treatment and disease management. Particularly, the advancement of X-ray technology in KOA demonstrates its potential for this purpose. Yet, existing X-ray prognosis research generally yields a singular progression severity grade, overlooking the potential visual changes for understanding and explaining the progression outcome. 
Therefore, in this study, a novel generative model is proposed, namely Identity-Consistent Radiographic Diffusion Network (IC-RDN), for multifaceted KOA prognosis encompassing a predicted future knee X-ray scan conditioned on the baseline scan. 
Specifically, an identity prior module for the diffusion and a downstream generation-guided progression prediction module are introduced. 
Compared to conventional image-to-image generative models, identity priors regularize and guide the diffusion to focus more on the clinical nuances of the prognosis based on a contrastive learning strategy.  
The progression prediction module utilizes both forecasted and baseline knee scans, and a more comprehensive formulation of KOA severity progression grading is expected.  
Extensive experiments on a widely used public dataset, OAI, demonstrate the effectiveness of the proposed method. 

  




§ INTRODUCTION





Knee osteoarthritis (KOA) is a common form of arthritis, mainly affecting knee joints with the symptoms of knee joint swelling, tenderness, and bone spurs. As the major factor of physical disability, KOA has recently shown an increasing prevalence in society, especially among the elders <cit.>. Currently, there are rare treatments for curing KOA without physical surgery (total knee replacement) <cit.>. Therefore, to relieve the pain and control the progression of KOA, the diagnosis and prognosis of the disease are significant, using different clinical assessment methods, such as medical imaging scans and patients' clinic questionnaires. In specific, the imaging scans are utilized to detect the structural changes within the knee joint areas, which indicates the characteristics of KOA. Typically, the severity of KOA is measured by the Kellgren and Lawrence (KL) grading scale system, which encompasses five different levels, from grade 0 (healthy cases) to grade 4 (highly severe cases) <cit.>. 

Compared to the extensive time and human resources required for specialists to detect changes, employing computer-aided techniques to automatically assess the severity and progression of KOA can significantly reduce time costs. This aids in the prompt diagnosis and prognosis of KOA at early stages <cit.>. 
Methods for grading the severity and progression of KOA have evolved from traditional machine learning approaches <cit.>  to advanced deep learning techniques <cit.>.  These modern approaches utilize various forms of input, including single images, multiple images, and images combined with demographic information (such as age, gender, BMI, and the WOMAC questionnaire) <cit.>.
Particularly, the advancement of X-ray technology and its application in KOA significantly improves the capability to detect and monitor the progression  <cit.>. Meanwhile, utilising the generative AI methods in the KOA analysis field focuses on augmenting the knee radiological dataset to improve the grading process <cit.>. 
Yet, existing X-ray prognosis research generally yields a singular progression severity grade using the baseline X-ray scan <cit.>, overlooking the potential visual changes for understanding and explaining the progression outcome. However, the visually predicted outcome is also crucial in clinical practice to explain the disease trajectory. 

In addition to enhancing these methods, we hope the longitude visual content will also aid clinicians in their decision-making by providing further explainability and as domain knowledge to guide the algorithms. For the methodology part, previous findings revealed that GANs yield poor bone structures and inconsistency in patients' identity information, which motivated us to devise a diffusion-based method with identity information. Identity patterns are crucial with a principle widely utilized in the forensic field. We anticipate that incorporating such information improves the consistency between baselines and synthesized longitude images.

Therefore, in this study, a novel deep generative model, namely Identity-Consistent Radiographic Diffusion Network (IC-RDN), is proposed for multifaceted KOA prognosis. IC-RDN creates a predicted future knee X-ray scan conditioned on the baseline scan and a future KOA severity grade as illustrated in Fig. <ref>. 
Specifically, an identity prior module for guiding the diffusion and a downstream generation-guided progression prediction module are introduced. 
Compared to a conventional image-to-image generative model,  identity priors regularize and guide the diffusion to focus more on the clinical nuances related to the prognosis based on a contrastive learning strategy.  
Using both forecasted and baseline knee scans, we anticipate a more comprehensive and enhanced formulation of KOA severity progression measured in grades with the progression prediction module.  
Extensive experiments on a widely used public dataset, OAI <cit.>, demonstrate the effectiveness of the proposed method. 




The key contributions of this study can be summarised in three-fold:


 
  * A novel deep generative model - IC-RDN is proposed for multifaceted KOA prognosis prediction for the first time, encompassing radiographic visual predictions and grading forecasting. 
 
  * An identity prior module for persisting patient's identity information on bone X-ray images is proposed to guide the diffusion process focusing on clinical nuances. 
 
  * Comprehensive experiments with the public OAI dataset demonstrate the effectiveness of the proposed method. 




§ RELATED WORKS




 §.§ Computer-Aid KOA Severity Grading



Existing studies in grading the severity of KOA can be categorised into KOA detection and KOA multi-class grading <cit.>. The KOA detection focuses on categorising several or combined severity grades from the whole grading system to achieve a coarse-grained classification. 

For instance, a graph neural network has been designed to utilize the shape boundary information of knee bones for KOA detection <cit.>. 
In addition to shape, texture features have been independently utilized for KOA detection to examine the internal pixel relationship <cit.>. 
To leverage both textural and shape information, the features extracted from the shallow and deeper layers of a DenseNet model are combined, where the shallow layers contain more shape patterns and the deeper layers focus on textural information <cit.>.

Multiclass grading, on the other hand, concentrates more on exploring each grading severity level in the system. 
As the KL grades can be considered a continuous measurement, with convolution neural networks, both regression and classification strategies are studied <cit.>. 
Likewise, an ordinal loss function is proposed to formulate the continuous nature of KL grades <cit.>. 
In parallel, the attention mechanisms identify the significant areas related to KOA multiclass grading <cit.>. 
In <cit.>, attention mechanisms are studied to combine and fuse multi-scale features from a network. 
Furthermore, to take advantage of the symmetrical structure of knee joints, a Siamese Network is employed to investigate the relationship between both knee joints. This approach facilitates the analysis of intricate KOA patterns in both fully supervised and semi-supervised manners, addressing the challenge posed by the scarcity of annotated datasets <cit.>. Likewise, a self-supervised method has been introduced. This method leverages both X-ray and MRI images to expand the pool of positive pairs in contrastive learning settings, offering a robust solution to dataset limitations <cit.>. 



 §.§ Computer-Aid KOA Longitude Progression Grading
 


Apart from KOA severity grading, predicting the progression of the disease is also significant for medical treatment. Existing works mainly focus on different aspects in predicting the progression of KOA: 1) Knee joint pain status; 2) KOA risk factor prediction; 3) Joint space narrowing (JSN) and cartilage structural change prediction; 4) KL grade prediction; and 5) total knee replacement surgery (TKR) prediction. 

Early studies on grading the progression of KOA mainly rely on conventional machine learning algorithms. In specific, the feature selection is leveraged in filtering the significant features contributing to the progression grades <cit.>. Parallelly, the comparison among different machine learning classifiers is conducted to select the algorithm with better progression grading performance <cit.>. <cit.> utilises the XGBoost model to predict the knee joint width as a KOA progression factor. To take advantage of various classification algorithms, the ensemble strategy is deployed among RandomForest, logistic regression, and support vector machine <cit.>. 
These machine learning-based methods show some encouraging results in feature engineering, where the extracted features are combined and utilized with methods such as multilayer perceptrons (MLPs) to predict the cartilage volume from MRI scans <cit.>. 

With the powered deep learning methods explored in both general and medical fields, end-to-end approaches have shown promising results in KOA progression grading. An MLP is deployed as a classifier to predict progression grades <cit.>. Neural Networks, otherwise, are utilised to predict pain status via extracting information from radiographics and demographics <cit.>. Likewise, CNN-based methods are also leveraged in analysing various KOA risk factors <cit.> and predicting the JSN progression <cit.>. To explore the KOA evolutionary trajectory, an Adversarial Evolving Neural Network is proposed to involve the estimated KOA patterns at different grades as auxiliary information <cit.>. 
Although predicting the direct KOA symptoms shows a clearer progression measurement, grading the progression KL grades and predicting TKR are also able to illustrate the progression in the clinical assessment domain. A binary classification model is designed based on whether the patient needs to take TKR as a future event <cit.>. 

Despite using medical imaging only, the information from multiple modalities, including imaging scans, clinical questionnaires, and demographics, are leveraged in a multimodal method to analyse the KOA progression comprehensively <cit.>. However, these studies only produce a direct outcome for progression grading ignoring the visual information from the given medical imaging techniques. Therefore, including visual information as an additional modality for analysing comprehensive multifaceted progression patterns is attractive to be explored. 






§ METHODOLOGY
 

As shown in Fig. <ref>, the proposed IC-RDN consists of three modules: 1) an identity prior module extracts patient's identity representations; 2) a diffusion-based generative module produces the future X-ray images based on the current-stage X-ray scans and the associated identity representations; and 3) a down-stream generation-guided progression prediction module that predicts the progression severity using both current and generated future X-ray images to provide comprehensive multifaceted progression patterns. 



 §.§ KOA Progression Radiographic Generation


This study incorporates X-ray medical imaging scans to predict the progression of KOA. The objective is to forecast the severity of KOA progression from two perspectives, encompassing severity grading score and visual information (i.e., a generated knee X-ray scan in the future - 12-month in this study).  
Initially, radiographic images for joints are captured during the baseline (first) X-ray scan, referred to as the baseline visit or 0-month visit. Patients are then asked to undergo a follow-up X-ray scan after 12 months for longitudinal tracking.  
Specifically, the i^th longitude X-ray joint pairs from the dataset is denoted as 𝐈_visit^i ∈{𝐈_baseline^i, 𝐈_12-month^i }, where 𝐈_visit^i ∈ℝ^C× W× H and C, W and H are the number of channels, the width and the height of the X-ray joint images. For notation simplicity,  𝐈_Baseline^i  and 𝐈_12 months^i can be shortened as 𝐈_v00m^i  and 𝐈_v12m^i. And the generated 12-month X-ray image for the i^th joint pair is denoted as 𝐈̂_12 months^i or 𝐈̂_v12m^i. The corresponding severity grade for the 12-month visit is indicated by y_KL^i and the grades produced by the prediction model are shown by ŷ_KL^i. 





 §.§ Diffusion-based Visual Prognosis


The diffusion-based generative module in the IC-RDN is designed to produce 12-month X-ray images from baseline visit images for understanding and explaining the potential visual changes in terms of KOA patterns. 
This incorporates a T-step Denoising Diffusion Probabilistic Model (DDPM) <cit.> under an image-to-image transition manner. 
The DDPM consists of two primary processes: the forward (or diffusion) process and the reverse process. 
Typically, in the forward process, the original data x_0 ∼ q_data(x_0) is transformed towards an isotropic Gaussian distribution. 
However, within our image-to-image diffusion model, the aim is to map the data x_0 ∼ q_data(x_0) to x_T, where, in our context, x_0 represents a 12-month visit X-ray image 𝐈_v12m^i, and  x_T corresponds to the baseline visit X-ray images 𝐈_v00m^i. 
This transformation is modelled as a fixed Markov Chain of length T. 
Subsequently, the reverse process mirrors the forward process, aiming to predict  
x_0 from x_T ∼ q_data(x_T). 

However, executing both processes directly in pixel space is highly resource-intensive, necessitating substantial memory and computational power. To mitigate this, a VQ-GAN-based encoder (see Fig. <ref> (b.1.)) encodes x_0 into a latent space z_0, wherein both the forward and reverse processes (see Fig. <ref> (b.3.)) are performed to generate z_T. A VQ-GAN-based decoder (depicted in Fig. <ref> (b.2.)) then decodes  z_T back from the latent space to the pixel space (x_T). The objective of this stable diffusion-based approach <cit.> is formulated as follows:

    ℒ_LDM := 𝔼_ε(x), ε∼𝒩(0,1), t[ ε - ε_θ(z_t, t) ^2_2 ],

where ε_θ(z_t, t) denotes a denoising autoencoders (time-conditional UNet) and θ contains the corresponding weights, where t = 1...T. 



 §.§ Knee Joint Identity Prior and Injection


To persist a subject's identity information in generating the corresponding 12-month X-ray scan, an identity prior module is devised to extract the identity-sensitive patterns. 
Specifically, we denote the identity representation for the i^th subject joint as 𝐅_ID^i ∈ℝ^d, where d indicates the dimension of the identity representation. A CNN-based network f(·) is designed to obtain the identity representation. Mathematically, we have:

    𝐅_ID^i = f(𝐈_v00m^i).

The identity information guides and regularizes the diffusion process to not only focus on the target visual domain distribution but also maintain necessary information from the subjects. 

For the supervision of f to formulate identity information, it is proposed to conduct training in a contrastive learning manner. A triplet loss is introduced for this purpose as follows:

    ℒ_Triplet = ∑[  f(𝐈_v00m^i) - f(𝐈_v12m^i) _2^2 -  f(𝐈_v00m^i) - f(𝐈_v00m^k) _2^2 + α]_+,

where a triplet is formulated as (𝐈_v00m^i, 𝐈_v12m^i, 𝐈_v00m^k): 𝐈_v00m^i indicates the anchor of current subject's joint; 𝐈_v12m^i refers to a positive sample from the same subject's joint; 𝐈_v00m^k is a negative sample with i≠ k; and α is a margin enforcing between positive and negative pairs. 
Note that the number of triplets M used in the training process is defined in the experiment implementations. 

The identity information acts as a guide and regulator for the diffusion process, ensuring that it not only aligns with the target visual domain distribution but also preserves essential information pertaining to the subjects. To inject the identity into the diffusion process, 

a cross-attention mechanism is applied between the features from the UNet layers in ϵ_θ and 𝐅_ID. For the t^th diffusion step and j^th layer, the attention can be formulated as follows: 

    𝐀_j = Softmax(𝐐𝐊^⊺√(d))𝐕,
    
    𝐐 = ϕ_j(z_t)𝐌^Q_j , 𝐊 = 𝐅_ID^i𝐌^K_j , and 𝐕 = 𝐅_ID^i𝐌^V_j ,





where 𝐌^∘_j represents a projection matrix containing trainable weights with a projection dimension d, and ϕ_j(z_t) indicates the (flattened) intermediate representation of the UNet with j^th layer. The UNet can be formulated by: 

    ϕ_j+1(z_t) = ϕ_j(z_t) + 𝐀_j,


To this end, the IC-RDN can be learnt via the following scheme:

    ℒ_LDM-Identity := 𝔼_ε(x), ε∼𝒩(0,1), t[ ε - ε_θ(z_t, t, f(𝐈_v00m^i)) ^2_2 ].




 §.§ Generation-Guided KOA Progression Prediction
 

In addition to generating 12-month X-ray image 𝐈̂_v12m^i for prognosis, a singular progression severity grade is predicted as well based on the baseline visit X-ray image 𝐈_v00m^i and the generated 12-month visit X-ray image 𝐈̂_v12m^i. 
Specifically, a two-stream CNN backbone with two branches h_v00m(·) and h_v12m(·) is applied for the two images to produce the associated representations:

    𝐏_v00m^i=h_v00m(𝐈_v00m^i),  𝐏_v12m^i=h_v12m(𝐈̂_v12m^i).

The two representations are then concatenated, denoted by ‖, and used in a multi-layer perceptron (MLP) classifier g(·) to produce the predicted progression severity grade (ŷ_KL^i), which can be formulated as:

    ŷ_KL^i = g( 𝐏_v00m^i ‖𝐏_v12m^i)


A cross-entropy loss is leveraged in the training process to guide the prediction training process, which can be formulated as: 


    ℒ_CE = -∑_i=1^N y_KL^i logŷ_KL^i

where N indicates the total number of joints in the dataset. 



§ EXPERIMENTAL RESULTS AND DISCUSSION
 



 §.§ Dataset and Pre-Processing


The longitudinal X-ray images used in this study were sourced from the Osteoarthritis Initiative (OAI) <cit.>, which is a multi-centre and longitude study containing chronological medical imaging scans and clinical assessment records for the subjects up to 96 months. This public dataset enables researchers to explore the natural progression of osteoarthritis, identify risk factors, and evaluate biomarkers for grading the KOA severity and progression. In this study, the patients with both baseline and 12-month visits and available 12-month severity KL grades are selected. Due to some participants opting out of the follow-up study and missing recording of clinical severity grading outcome, not all have recordings for the 12-month visit image and severity KL grades. Consequently, 7,923 joints have both the KL grade available and radiographic recordings from both the baseline and 12-month visits. As the demographic information shown in Table <ref>, the amount of joints with changed severity grades among 12 months is much smaller than that of without changes, where only 6.15% joints show a changed severity grade. However, across these joints without the severity grade changes, the visual structural details do not remain the same, which also shows the importance of the proposed multifaceted progression predicting method. The dataset was randomly divided among subjects with a distribution of 70%, 10%, and 20% for training, validation, and testing, respectively. 



As the raw data from the OAI contains bilateral knee joints and includes almost all knee bones around the joint area, knee joint localisation methods are leveraged to extract the knee joint area individually <cit.>. After detecting the knee joint edge, a 224×224 region surrounding the middle of the joint was extracted from the raw data for modelling. The remaining transformation approaches for the joint images are followed <cit.>, including adjusting the brightness, contrast and saturation of the images. 




 §.§ Implementation Details



A pretrained ResNet 18 on ImageNet <cit.> is adopted as the backbone classification model for identity detection, which extracts the identity features from the baseline X-ray scans. Within Identity Prior's contrastive learning framework, we have designated the number of triplets M as 7,000 for training and 1,000 for testing. 
The Identity Prior training phase is fine-tuned using the Adam optimizer with a learning rate of 0.01, and the batch size for this phase is established at 128. 
For the training process of the diffusion-based generative module,  we set the number of steps to 1,000,000 across 400 epochs. 
The Adam optimizer is employed here as well, with a learning rate set at 1× 10^-4. The batch sizes for training and testing are 64 and 8, respectively. 

For the backbone models in the generation-guided KOA progression prediction module, various CNN methods are adopted, such as VGG 16, VGG 19, ResNet 18, ResNet 50, and DenseNet 121, for a comprehensive comparison <cit.>. Specifically, the backbone prediction models for both branches, corresponding to 𝐈_v00m and 𝐈̂_v12m, share the same architecture. In addition, the output features dimension from these backbone models is set to 128. Then the MLP classifier g(·) takes the input dimension of 256, which contains two linear layers, an activation layer, and a dropout layer. 
The training process is conducted with a batch size of 16. An SGD optimizer is utilized with the learning rate of 5× 10^-3 and a weight decaying factor of 5× 10^-3 for every 5 epochs.
The entire training regimen is executed using PyTorch 2.0.1 on an NVIDIA RTX A6000 with 48 GB of video memory.





 §.§ Quantitative Performance on Visual Prognosis


The evaluation of the generated future X-ray images encompasses both quantitative and qualitative assessments. Quantitatively, we engage in analysis through the conventional generation task metrics, inception score and a downstream task: the prediction of KOA progression severity grade. For using the inception score in this study, we use the Inception-v3 network trained on the OAI 12-month visit X-ray dataset to replace the weights pre-trained on the ImageNet <cit.> as the knee X-ray images are different from the natural images at both category and structural level. As shown in Table <ref>, the X-ray images generated by our proposed IC-RDN illustrate a similar IS with the X-ray images generated by the image-to-image diffusion model and outperform the X-ray images generated by Cycle-GAN <cit.> and other SOTA generative models <cit.> significantly. As for the similar IS between the proposed IC-RDN and the image-to-image diffusion model, the results show similar performance on image quality and diversity. However, 
the images generated by the I2I diffusion model show a minor contribution to predicting the KOA progression severity grades, and the challenge of identity inconsistency remains in the generated X-ray images from the I2I diffusion model. The detailed comparison for identity consistency is shown on Sec. <ref>. 

Although the inception score is able to show the quality and diversity of the generated future images, the downstream task of predicting the KOA progression severity grade can be more relevant to the value of how the generated future X-ray images contribute to exploring the KOA progression patterns. We have introduced the approach to predict KOA severity progression (Sec. <ref>), leveraging both the generated future X-ray image and the baseline X-ray image to determine a progression severity grade. 

In addition, as stated in Sec. <ref>, various CNN-based backbones are deployed in the generation-guided KOA progression prediction module. The comprehensive results among these backbone models are shown in Table. <ref>. 

Specifically, the predicted 12-month progression X-ray images generated by Cycle GAN, the image-to-image diffusion model, the recent SOTA generative methods, and the proposed IC-RDN are used in the generation-guided KOA progression grade prediction module separately. Accuracy is used as the evaluation method to measure the performance of predicting the progression grades. Among all backbone models used in this module, predicting the progression severity grades using X-ray images from baseline visit and generated by the proposed IC-RDN shows significantly outperformed results. In detail, the image-to-image diffusion model shows a limited contribution to providing the multifaceted KOA progression patterns. And the Cycle GAN illustrates the minor ability to enhance the performance of predicting the KOA progression severity grades. Therefore, the quantity evaluation demonstrates the added value of incorporating generated future X-ray images in enhancing the accuracy of KOA progression severity grade predictions.




To illustrate the performance of generated X-ray images contributing to predicting the KOA progression grades, the t-Distributed Stochastic Neighbor Embedding (t-SNE) technique, a sophisticated machine learning algorithm designed for the task of high-dimensional data visualization, is utilised in this study. In detail, to reduce the complex, multi-dimensional data into a two-dimensional space, t-SNE converts the high-dimensional distances between data points into conditional probabilities that represent similarities. Specifically, given a predicted X-ray image generated by the proposed IC-RDN, the trained VGG 19 network is used to produce a feature map corresponding to the KOA progression grade patterns. Then, these representations are visualised using t-SNE to show the related grades spatially. As shown in Fig. <ref>, the nodes with different colours indicate the corresponding grades. In Fig. <ref>.a, clear clusters for the KL grades of 0, 3, and 4 are located on the left and right sides. The green and yellow nodes referred to KL grades 1 and 2, show that the module for predicting the KOA progression severity grade has a minor ability to categorise these two classes, which is the challenge among different KOA progression prediction and severity grading methods <cit.>. However, in Fig. <ref>.b, the nodes for different grades are mixed and the cluster centres are hard to identify. 








 §.§ Qualitative Performance on Visual Prognosis



In terms of qualitative analysis, the comparison between X-ray images synthesized by the Cycle GAN and those produced by the proposed IC-RDN offers a compelling visual assessment, as depicted in Fig. <ref>. The comparative study highlights notable distortions in bone structures emanating from the Cycle GAN approach, particularly in the articulation of joint spaces. In stark contrast, our proposed IC-RDN exhibits a pronounced fidelity in preserving the integrity of key bone edges and the intricacy of inner textures, thereby underscoring the superior capability of IC-RDN in maintaining structural information.

Further elucidation is provided through an expanded corpus of X-ray images generated by the IC-RDN, which is shown in Fig. <ref>. This compilation serves not only to reinforce the initial observations regarding the preservation of critical structural details but also to demonstrate the consistent performance of the IC-RDN across a diverse set of examples. The consistent replication of complex anatomical features by the IC-RDN model is indicative of its sophisticated understanding of underlying bone structures, which is paramount for illustrating the multifaceted KOA progression patterns.

Moreover, this comparison elucidates the inherent limitations of traditional GAN-based approaches in contexts where the preservation of exact structural features is paramount. The propensity for distortion within Cycle GAN-generated images can significantly impact the clinical utility of such synthesized X-ray images, particularly in monitoring the KOA progression. In contrast, the proposed IC-RDN shows its adeptness at capturing and reproducing the nuanced details of bone structure suggesting its potential as a multifaceted KOA progression prediction outcome.







 §.§ Analysis on Identity Consistency



To further evaluate the efficacy of the identity prior module, we delve into a comprehensive comparison involving the X-ray images generated by the I2I diffusion model and the proposed IC-RDN. The visual examples presented in Fig. <ref> meticulously outline the distinctions between these methods alongside the actual 12-month visit X-ray images, where the key differences are highlighted by the bounding boxes. Initially, we observe the generated X-ray images devoid of identity persistency, which starkly contrasts with the subsequent imagery. 

A noteworthy observation is the preservation of identity features in the X-rays produced by our IC-RDN, as evidenced by the distinct texture and precision of bone edge structures. This preservation extends to intricate details within the knee joint, such as the orientation of the menisci and the nuanced protrusions along the bone edges, which are critical for maintaining identity persistency. 




 §.§ Ablation Study


To better evaluate the contribution of the proposed Identity Prior module, an ablation study is conducted by removing the Identity Prior and only the Diffusion backbone remaining. The results in Table <ref> demonstrate the effectiveness of the Identity Prior module with a superior longitude prediction accuracy among all selected classification methods. 





 §.§ Limitations


The major limitation of this study regards the generation quality. Directly using X-ray scans generated by IC-RDN for a KOA severity grading showcases a suboptimal result compared with using the baseline visit X-ray to predict the KOA progression grade directly. Though the predicted 12-month scan can enhance the KOA severity prediction with using the baseline visit jointly. In addition, computer hallucinations <cit.> are a crucial issue among large model fields. Particularly, the analysis regarding the disease-specific changes is insufficient. Therefore, careful consideration of applying the generative method to actual clinical usage is needed. 



§ CONCLUSION
 

In this study, a generative-based method, IC-RDN is proposed for KOA prognosis, which forecasts future X-ray scans for joints and singular progression KOA severity grades. This obtains multifaceted KOA progression patterns for understanding and explaining KOA prognosis.
Specifically, an identity prior module for the diffusion and a downstream generation-guided progression prediction module are introduced for regularising and guiding the diffusion to focus more on the clinical nuances related to the prognosis and producing a more precise KOA severity progression grading. 

Extensive experiments on a widely used public dataset, OAI, demonstrate the effectiveness of the proposed method.

In our future studies, the disentanglement of the representations related to the identity and KOA severity will be explored for fine-grained guidance in the X-ray image generation process. Meanwhile, the multimodal patterns among different medical imaging scans (X-ray, MRI, and CT) and patients' demographic information will be involved in producing a fine-grained KOA progression assessment outcome. 





§ ACKNOWLEDGEMENTS

This study was partially supported by Australian Research Council (ARC) grant DP210102674. The OAI <cit.> is a public-private partnership comprised of ﬁve contracts (N01-AR-2-2258; N01-AR2-2259; N01-AR-2- 2260; N01-AR-2-2261; N01-AR-2-2262) funded by the National Institutes of Health, a branch of the Department of Health and Human Services, and conducted by the OAI Study Investigators.






splncs04



 Paper Title: Domain Shift Analysis in Chest Radiographs Classification in a Veterans Healthcare Administration Population.








    Curvature-Adapted Hypersurfaces of 2-type

in non-flat quaternionic Space Forms
    Ivko Dimitrić
    
===============================================================================


    ^1 Advanced Computing for Health Sciences, Computational Sciences and Engineering Division,
 Oak Ridge National Laboratory 

     ^2 IT Services Division, Oak Ridge National Laboratory 

     ^3 Veteran Affair's Association 

     ^4 Los Alamos National Laboratory 

     ^5 Yale University 

     ^6 Stanford University 

     ^7 Department of Biomedical Informatics, Vanderbilt University Medical Center 

     ^* Corresponding Author Email: chandrashekm@ornl.gov
     

§ ABSTRACT 

Objectives: This study aims to assess the impact of domain shift on chest X-ray classification accuracy and to analyze the influence of ground truth label quality and demographic factors such as age group, sex, and study year.

Materials and Methods: We used a DenseNet121 model pretrained MIMIC-CXR dataset for deep learning-based multilabel classification using ground truth labels from radiology reports extracted using the CheXpert and CheXbert Labeler. We compared the performance of the 14 chest X-ray labels on the MIMIC-CXR and Veterans Healthcare Administration chest X-ray dataset (VA-CXR). The VA-CXR dataset comprises over 259k chest X-ray images spanning between the years 2010 and 2022. 

Results:
The validation of ground truth and the assessment of multi-label classification performance across various NLP extraction tools revealed that the VA-CXR dataset exhibited lower disagreement rates than the MIMIC-CXR datasets. Additionally, there were notable differences in AUC scores between models utilizing CheXpert and CheXbert. When evaluating multi-label classification performance across different datasets, minimal domain shift was observed in unseen datasets, except for the label "Enlarged Cardiomediastinum." The study year's subgroup analyses exhibited the most significant variations in multi-label classification model performance. These findings underscore the importance of considering domain shift in chest X-ray classification tasks, particularly concerning study years.

Conclusion:
Our study reveals the significant impact of domain shift and demographic factors on chest X-ray classification, emphasizing the need for improved transfer learning and equitable model development. Addressing these challenges is crucial for advancing medical imaging and enhancing patient care.





 Notice: Office of Science of the U.S. Department of Energy. This manuscript has been authored by UT-Battelle, LLC, under contract DE-AC05-00OR22725 with the US Department of Energy (DOE). The US government retains and the publisher, by accepting the article for publication, acknowledges that the US government retains a nonexclusive, paid-up, irrevocable, worldwide license to publish or reproduce the published form of this manuscript, or allow others to do so, for US government purposes. DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan (http://energy.gov/downloads/doe-public-access-plan). 





§ INTRODUCTION



Chest radiography is the first-line imaging test for respiratory disease and some forms of cardiac disease. Chest X-ray abnormalities detection has been automated by recent advanced artificial intelligence techniques. Accurate chest X-ray classification has played an important role in many biomedical applications to accelerate the diagnosis and treatment of conditions like pneumonia, heart failure, rib trauma, pulmonary fibrosis, etc. Though there has been an increase in the integration of machine learning models to enhance diagnostic promise, the efficacy of these models hinges on the availability and quality of training data. 

Recent years have seen a wave of artificial intelligence (AI) and machine learning models for clinical applications from large research medical centers or based on limited de-identified datasets. The introduction of transfer learning has resulted in AI models being accessible to researchers with fewer computational resources and less data, allowing them to leverage existing knowledge. Transfer learning is the ability to apply a model trained on a dataset to another dataset of interest.  A major consideration for transfer learning is domain shift, the dissimilarity between data distributions from the source used for training and the population to which the models are applied. This divergence between publicly available datasets and private, institution-specific datasets can introduce substantial bias and hinder the generalization of machine learning models to real-world scenarios.

In this study, we quantify the domain shift between the public domain (MIMIC-CXR) and a private dataset (VA-CXR). The efficacy of the transfer learning approach for chest X-ray classification and the subsequent impact on classification accuracy when dealing with domain shift form the main focus of this study.  The domain shift is traditionally viewed only based on the model and its performance, ignoring the multi-fold causes leading to the shift. We comprehensively address the effects of domain shift in this three-stage study: 1) Compare the performance between the source domain and target domain accuracy on the chest X-ray classification. 2) Quantify the quality of the ground truth extracted from the radiology reports, as supervised learning models are heavily dependent on the quality of the labels.  3) Analyze the relationship between demographic factors and classification accuracy, as domain mismatches often originate in demographic mismatches. 

By addressing the critical interplay of domain shift and demographic factors, this study not only provides insights into the technical challenges of adapting machine learning models to private datasets but also underscores the broader implications of these challenges in the context of healthcare. Our research contributes to developing more accurate, robust, and generalizable chest X-ray classification models by creating a systematic approach for using an existing model and understanding the nuances of a model's performance before applying it to a new population. 




§ METHODS





 §.§ Dataset

MIMIC-CXR is a publicly available dataset of 377,110 chest X-rays associated with 227,827 imaging studies from 65,379 patients. MIMIC-CXR is collected from the inpatient setting of Beth Israel Deaconess, a Boston hospital. (refer Table <ref>) We also use the test-split of this dataset (Test-split MIMIC-CXR) as the hold-out set analysis, which consists of 5159 studies with 293 patients <cit.>. 

VA-CXR is a private dataset of 259,361 chest X-rays associated with 91,020 imaging studies from 35,771 patients. VA-CXR is collected in the outpatient setting of the Boston Veterans Healthcare Administration station (refer to Table <ref>). The ground truth labels were extracted from the VA's corporate data warehouse (CDW) and joined with images using patient information from DICOM headers as published in Knight et al. <cit.>.   






 §.§ Ground truth Label Extraction


We used the CheXbert labelers <cit.> to expertly assign labels to 14 specific labels (Atelectasis, Cardiomegaly, Consolidation, Edema, Enlarged Cardiomediastinum, Fracture, Lung Lesion, Lung Opacity, No Finding, Pleural Effusion, Pleural Other, Pneumonia, Pneumothorax, Support Devices) associated with different chest conditions from radiology reports. CheXbert is a BERT-based approach that automates the detection of these observations, effectively streamlining the process of annotating medical images and reports. The NLP label extraction outputs scores for four classes: positive, negative, blank, and uncertain, associated with each of the 14 labels. As the class names indicate, for example, for Pneumonia, positive class: radiology report indicates that the patient has Pneumonia; negative class: radiology report indicates that the patient doesn't have  Pneumonia; uncertain class: radiology report mentions  Pneumonia but NLP tool unable to determine if it is positive or negative; blank class: radiology report doesn't mention  Pneumonia.



 §.§ Label Validation

Because ground truth assignment ultimately determines the accuracy of the imaging classifiers, we developed an image evaluation procedure in two steps. First, we evaluated the agreement between the NLP label extraction tool, CheXbert, and its precursor, CheXpert <cit.>, a rule-based tool. We focused on positive class agreements for our evaluation, using only positive classes for classification, and have combined uncertain/negative class agreements. The disagreement between NLP label extraction tools indicates ambiguity/ less confidence on the labels, ultimately creating an unreliable ground truth. The agreement is measured for both MIMIC-CXR and VA-CXR datasets; Of note that neither of the datasets was used to train the NLP extraction tools. 




  §.§.§ Relation to Diagnoses Codes

To validate the ground extracted from chexpert-labeler, we analyzed the relationship of specific ground truth labels to ICD codes in the patient's electronic health record (EHR) extracted from the VA's Corporate Data Warehouse (CDW).The assignment of ICD-9 and ICD-10 diagnosis codes associated with each condition was exploratory and not extensively optimized. 

Starting concepts were retrieved from the Chexpert-labeler github repository , where phrases they used to search notes can be found at https://github.com/stanfordmlgroup.  These phrases, along with our own expertise, were used to identify diagnosis codes and cross-reference radiology reports with diagnoses.For example:   Pneumonia was identified by chexpert-labeler as indicated by pneumonia, infection, infected process, and infectious; 
Edema was indicated by terms edema, heart failure, chf, vascular congestion, pulmonary congestion, indistinctness, and vascular prominence; Fracture was indicated solely by the word fracture; and Pneumothorax was identified by either pneumothorax or pneumothoraces.

This method enabled us to validate the ground truth labels by correlating them with the relevant ICD codes in the patients' EHRs, ensuring accurate cross-referencing of radiology reports with diagnoses.

For instance, the ICD-9 codes we used to indicate a pneumonia diagnosis in the outpatient diagnosis tables ranged from 480 to 486 and included 487.0. These codes encompass viral, bacterial, and other types of pneumonia, as well as pneumonia caused by unspecified pathogens. A similar approach was applied to ICD-10 codes for pneumonia and other conditions. The specific ICD codes used for each condition are detailed in Appendix Table <ref>.



 §.§ Multi-label Image Classification


Using the 14 labels extracted from the corresponding radiology reports with CheXbert, we created a multi-label image classification model from X-ray images. We used a pre-trained DenseNet model<cit.> as the core framework, removed the top classification layer, and integrated a custom classification layer for multi-label output. 
Previous work shows the effectiveness of different resolutions of DenseNet-121 based multi-label classification chest X-ray model on MIMIC-CXR dataset  <cit.>. This work uses the MIMIC-CXR trained DenseNet-121 Model on Chest X-ray pre-processed into 256x256 JPG images <cit.>. 




 §.§ Metrics

We evaluated our models using the Area Under the Curve (AUC). The AUC score was calculated separately for each of the 14 labels, indicating the separability measure for a given chest X-ray label. 
We also analyzed the difference in AUC scores between MIMIC-CXR and VA-CXR, and the prevalence for each label was calculated as the number of studies with positive results for a given label divided by the total number of labels, indicating the label's presence in the given cohort. 




 §.§ Domain Shift Analysis

We compared our source and target datasets along different dimensions: 1) Demographic Details: Age At Time of Imaging Study, Sex; 2) Imaging Study Details: Study Year, View  Point (Lateral view (Lat), erect anteroposterior (AP), posteroanterior (PA)) ; 3) Ground truth Labels: 14 labels. All the factors were analyzed against the accuracy of multi-label image classification. 
The impact can be estimated by comparing the prevalence of the above-listed factors across source and target domains with the performance of the multi-label classification. 

The subgroup analysis is performed on unseen datasets: Test Split MIMIC-CXR and VA-CXR. 


Study Year Analysis: We obtain the study year from VA-CXR studies date reported in DICOM headers.  The MIMIC-CXR dataset is unidentified; study years are replaced by anchor years. The comparison of performance across study years between the datasets isn't possible, so we don't present a study year analysis on the MIMIC-CXR dataset. 


Performance Analysis based on Sex: We obtain the sex of the patients in VA-CXR from the VA's corporate data warehouse (CDW) and the patients in MIMIC-CXR from the metadata. Sex is only classified into binary classes of male and female. We analyze the label-wise performance for each sex across both datasets. 


View Position Analysis: We obtain the view position of the images from the DICOM metadata. We use four main viewpoints for analysis: Lateral view (Lat), erect anteroposterior (AP), posteroanterior (PA))  and Left Lateral view (LL). We compare the label-wise performance across both datasets


Age Group Analysis:
We define the age of the patient as the age when the imaging study was performed. For VA-CXR, this age was calculated from the date of birth obtained from VA CDW and the date of imaging study from DICOM metadata. For MIMIC-CXR, we calculated age using the anchored date of birth and anchored imaging study date. The anchored dates were amended in a manner that the difference in years is constant, so we were able to approximate the age of the patient. The label-wise performance for both datasets was compared across six age groups: (0-50], (50-60], (60,70], (70,80], (80,90], and (90,100]. 



§ RESULTS




 §.§ Ground truth Label Validation

Table <ref> presents a comprehensive analysis of agreement and disagreement rates between Chexpert and Chexbert on the MIMIC-CXR and VA-CXR datasets across the 14 labels. The results shed light on the performance and consistency of Chexpert and Chexbert and the uncertainty of the ground truth labels. Notably, in MIMIC-CXR, Atelectasis was identified in 19.5% of cases with a disagreement rate of 10.1%, whereas in VA-CXR, the identification rate was lower at 9.8% with a disagreement rate of only 0.7%. This discrepancy in positive identification and disagreement rates is further exemplified in conditions like Cardiomegaly, where MIMIC-CXR reported positive identification in 15.7% of cases with a significant disagreement rate of 39.1%, contrasting with VA-CXR's 9.3%  positive identification rate and 13.7% disagreement rate.









 §.§ Comparison with Diagnosis Codes


While the X-ray classification accuracy was evaluated against the label extracted using NLP of the appropriate radiology report, it is possible also to compare directly against diagnosis codes in the clinical record.  However, interpreting this comparison is challenging due to two main issues: first, diagnoses are based on a broader range of information beyond just the X-ray, and second, patients may have multiple conditions, making it difficult to determine the specific reason each X-ray was ordered. Despite these challenges, we have conducted two useful comparisons of X-ray findings against diagnosis codes for nine categories within our dataset.

The top of Figure <ref> shows the fraction where a diagnosis related to the label was recorded within a week (either before or after) of the X-ray, compared to the total number of patients, including those who never had the diagnosis. Patients who had the diagnosis but not within the one-week window were excluded from this calculation. This plot, labeled 'Sensitivity,' reflects the proportion of cases where a diagnosis occurs within a week of noting the condition in the X-ray.  
We observe a positive finding on an X-ray, as extracted from the radiology report, which is mostly associated with a specific diagnosis of pneumonia and edema and frequently associated with pleural effusion and pneumothorax.  We also see that the ChexBert model is more frequently associated with a corresponding diagnosis than the ChexPert model.



The bottom of Figure <ref> compares the factor by which the enrichment ratio of positive finding in the radiology report increases with a concurrent (within one week) diagnosis.  Pleural effusion, for example, is seen 10 times more frequently in the radiology report when a concurrent diagnosis code is noted.  It is also true, however, that 2/3 of the pleural effusion diagnoses are not accompanied by a positive finding in a radiology report, as shown in the top bar chart, and also that 90% of the time a negative finding is made, no diagnosis is found.  We see that the conditions with the highest enrichment ratios differ from those with the highest sensitivities, but the trend that the ChexBert model generally outperforms the ChexPert model by a few to 20% remains.

Several factors were observed to contribute to the quantitative comparison of labels extracted from the radiology reports to the observed diagnoses, including the extent to which diagnosis codes could be found that correspond well to the radiology report finding, whether the X-ray is a screening or confirmatory test, the prevalence of the condition, and the chronic vs. acute nature of the condition.  We provide this figure as a survey across these factors and present the complete set of counts in Appendix Table <ref>.  Specific ICD 9 and 10 codes used for each condition are provided in the Appendix Table <ref>.





 §.§ Multi-label Image Classification on VA-CXR across NLP tools


Table <ref> compares ground truth extraction techniques, specifically focusing on the label-wise performance metrics on the VA-CXR dataset. Two techniques, ChexPert and ChexBert, are evaluated based on their Area Under the Curve (AUC), prevalence, and count across 14 labels. Across the findings, both techniques generally demonstrate similar AUC values, indicating comparable performance in distinguishing positive cases. For instance, in identifying Atelectasis, both ChexPert and ChexBert exhibit AUC values around 0.80, with similar prevalence and count numbers. Notably, in the case of Cardiomegaly, ChexBert shows a notably higher AUC of 0.862 compared to ChexPert's 0.753, suggesting potentially superior performance in this specific finding. However, the prevalence and count metrics vary across the findings and between the two techniques. For example, ChexPert tends to have higher prevalence and count values for several findings like ECM (Enlarged Cardiomediastinum), while ChexBert shows higher values for others such as Pleural Effusion. 




 §.§ Multi-label Image Classification Performance across multiple datasets


Table <ref> shows the comparison of MIMIC-CXR (Source dataset), Test Split MIMIC-CXR (Hold-out Source Dataset), and VA-CXR based on AUC on the 14 labels.  The Hold-out Source dataset is the test split of MIMIC-CXR dataset,  DenseNet-121 Model. The Test Split of MIMIC-CXR gives us a fair comparison to VA-CXR, the unseen target dataset. This can be observed based on the AUC drop from the overall MIMIC-CXR to test-split. In Table <ref>, the difference in AUC between Hold-out and Target indicates the performance variation between VA-CXR and Test Split of MIMIC. The Negative value of the difference in AUC indicates that the VA-CXR performs better than the Test Split MIMIC-CXR, and the positive value indicates that the Test Split performs better. The Enlarged Cardiomediastinum (ECM) label has the highest difference in AUC, indicating a huge performance drop in VA-CXR. This could directly impact the lack of a large number of image studies in ECM in the source dataset compared to the target. 






 §.§ Study Year-wise Performance on VA-CXR 

Table <ref> shows the label-wise distribution of the study years. The AUC systematically drops as for study year 2020 to 2022 for all labels except Consolidation and Pleural Other, which peaks at 2020. The prevalence increases over the years in VA-CXR for Atelectasis, Enlarged Cardiomediastinum, and Pleural Effusion. 








 §.§ Performance across datasets based on Sex

Figure <ref> compares the label-wise AUC and prevalence of the two sexes across the Test Split MIMIC-CXR and VA-CXR. This comparison is essential as the female-male patient ratio in VA-CXR is higher than that of MIMIC-CXR; dashed lines in the figure <ref> can observe this. 







 §.§ View Position-wise across datasets

Figure <ref> compares labels across Test-Split MIMIC-CXR and VA-CXR. The VA-CXR doesn't contain any Lateral images; the figure shows only MIMIC-CXR performance on the Lateral. The prevalence and AUC of viewpoints vary based on the label, with ECM having the most difference between the datasets. Pleural Other has a drop in performance in VA-CXR in AP view position, potentially due to low prevalence in both datasets. 
 






 §.§ Age Group-wise comparison across datasets

Figure <ref> shows the VA-CXR prevalence increase as the age increases across all labels except No Finding, Lung Lesion, and Pneumothorax.  The highest performance drops of 0.15 to 0.2 AUC  in VA-CXR compared to the Test Split of MIMIC-CXR can be observed in Enlarged Cardiomediastinum and Support Devices. In VA-CXR, it is interesting that the AUC performance across the age groups is more stable, i.e., there is not much change in AUC, with the exception of Atelectasis.  








 §.§ Summary

As seen from the results, the ground truth validation and multi-label classification performance across the NLP extraction tools showed that though the VA-CXR dataset has lesser disagreement rates than the MIMIC-CXR datasets, and there were AUC differences between models when using ChexPert and ChexBert (potentially propagated by distribution differences). When comparing the multi-label classification performance on different datasets, the unseen datasets didn't show domain shift other than a few labels such as Enlarged Cardiomediastinum. Among the different subgroup analyses, the study year had the most drastic differences in the performance of the multi-label classification model. These differences indicate that the domain shift is definitely of concern in study years.




§ DISCUSSION


Our study quantified the domain shift between a publicly available critical care dataset (MIMIC-CXR) and an outpatient, private institutional dataset (VA-CXR), assessing the efficacy of transfer learning for chest X-ray classification.  Because domain shift encompasses multifaceted factors, our comprehensive study was structured into three interrelated parts: the quality of ground truth, performance comparison, and subgroup analysis. 

Quality of Ground Truth 
Supervised learning models heavily rely on the quality of labels for training. Therefore, we quantified the ground truth quality extracted from radiology reports in both datasets. Despite being a source dataset, MIMIC-CXR exhibited significantly higher disagreement rates between the two ground truth information extraction NLP algorithms than the target, VA-CXR dataset. The potential performance drift of the target can be attributed to the mismatch of ground truth extraction methods. Our analysis underscored the importance of high-quality annotations in mitigating the effects of domain shift <cit.> and improving model performance.

Performance Comparison: 
We compared the difference in classification performance between MIMIC-CXR and VA-CXR to understand the extent of domain shift and its implications for model generalization. Our findings revealed notable variability in classification accuracy, highlighting the challenges posed by domain shift. We observe that the prevalence and performance were directly associated; for example, MIMIC-CXR's enlarged cardio mediastinum low prevalence may be the reason for its low performance. 

Subgroup analysis 
Subgroup analysis using demographic factors is crucial as we expect our populations to have different demographic distributions <cit.>. 
We observed a decline in the performance of the VA-CXR dataset over time, particularly in the years following 2020. This drop may be due to differences in the study years between the source dataset (MIMIC-CXR) and the VA-CXR dataset, as well as potential impacts from the pandemic years (note: we did not evaluate additional labels related to the pandemic).
We observed that both groups, though with different distributions, have aging populations across conditions of interest. Still, the prevalence of most conditions with age increases in the VA-CXR dataset, but it is highly variable in MIMIC-CXR.  For sex subgroup analysis, the performance was similar across the female and male populations in the datasets. Though the VA-CXR population is skewed towards the male population, the model performed well on the female population. This can be attributed to the source domain (MIMIC-CXR) having a balanced male-to-female ratio. 
For the view-based subgroup, we observed high variability between the views. This behavior can be attributed to the necessity of specific views for accurately diagnosing certain diseases. For instance, conditions such as Pneumonia and ECM primarily rely on PA (Posteroanterior) or AP (Anteroposterior) views for diagnosis.

Chest radiography is a first-line imaging modality for assessing lung and heart conditions. With advanced artificial intelligence techniques, the automation of chest X-ray abnormality detection has seen significant progress <cit.>. This has accelerated the pace of diagnosis and treatment and opened avenues for exploring biomedical applications in various clinical settings. However, the efficacy of these machine learning models heavily relies on the quality and availability of training data <cit.>.


Clinical relevance
Despite the surge in artificial intelligence and machine learning applications in clinical settings, their adoption has been uneven, primarily due to disparities in funding<cit.>, expertise, and availability of computing resources. While academic medical centers have been at the forefront of AI adoption, the accessibility of such resources remains a challenge for smaller research institutions. Introducing transfer learning has mitigated some of these challenges by enabling researchers to leverage pre-existing models trained on publicly available datasets for their specific applications. However, transfer learning introduces the concept of domain shift, which poses a significant obstacle to achieving robust and generalizable models.

Potential Workflows
While we primarily focused on matching image classification to known findings typical of chest X-rays and extracted those from radiology reports using NLP, this may not be the best workflow with which image classification can be incorporated into a medical outcomes workflow.  Tiu, et al. have noted the difficulties and inaccuracies associated with training chest X-ray classifiers against such labels and demonstrated improved accuracy and efficiencies with a self-supervised workflow to train their image classification model.<cit.>.  Even more ambitiously, they argue in Reference <cit.> that all modalities of data should be co-optimized in a much broader framework, known as a foundation model.  Our observations in Figure <ref> that most of the X-ray findings do not map directly onto diagnoses support the idea that more flexible workflows may be required to optimally incorporate AI-based image classification into medical outcome prediction.

By systematically addressing the challenges of domain shift and demographic factors, our study contributes valuable insights to developing more accurate and robust chest X-ray classification models. Furthermore, our findings have broader implications for healthcare, emphasizing the importance of understanding and mitigating biases in machine learning applications to ensure equitable healthcare delivery. Our research serves as a foundation for future studies aimed at refining transfer learning techniques and enhancing the generalizability of machine learning models in clinical practice.

Limitations
It is crucial to acknowledge that the data for the two studies were collected under significantly different clinical settings, which may influence the outcomes. The MIMIC database primarily includes inpatient data, where patients may be more critically ill, as indicated by the prevalence of portable anteroposterior (AP) images that tend to be less precise. In contrast, the VA outpatient population likely represents individuals who are relatively less ill, with higher-quality image data collection due to the setting's focus on outpatient care. Although we were unable to assess these differences quantitatively, we hypothesize that the clinical setting could substantially impact the performance of our models, influenced by the severity of diseases and other contextual factors.





§ CONCLUSION

In conclusion, our study sheds light on the critical interplay between domain shift, demographic factors, and the efficacy of transfer learning in chest X-ray classification. By quantifying the domain shift between a publicly available dataset and a private institutional dataset, we have identified disparities in classification accuracy and highlighted the challenges posed by transferring models across domains. Moreover, our analysis of demographic factors underscores the importance of considering population diversity in model development to ensure equitable healthcare outcomes. Moving forward, addressing these challenges will require concerted efforts to improve data quality, develop robust transfer learning techniques, and enhance the generalizability of machine learning models in clinical practice. Ultimately, by fostering a deeper understanding of these complexities, our research paves the way for the development of more accurate, reliable, and accessible chest X-ray classification models with the potential to transform patient care and advance the field of medical imaging.



§ ACKNOWLEDGMENTS

This work is sponsored by the US Department of Veterans Affairs using resources from the Knowledge Discovery Infrastructure which is located at the Oak Ridge National Laboratory and supported by the Office of Science of the U.S. Department of Energy. This manuscript has been authored by UT-Battelle, LLC, under contract DE-AC05-00OR22725 with the US Department of Energy (DOE). The US government retains and the publisher, by accepting the article for publication, acknowledges that the US government retains a nonexclusive, paid-up, irrevocable, worldwide license to publish or reproduce the published form of this manuscript or allow others to do so for US government purposes. DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan (
http://energy.gov/downloads/doe-public-access-plan
). 






 
 



§ APPENDIX:




 §.§ Age When Image was Taken versus Study Year














 Paper Title: Learned Compression of Point Cloud Geometry and Attributes in a Single Model through Multimodal Rate-Control.


Ground-to-UAV and RIS-assisted UAV-to-Ground Communication Under Channel Aging: Statistical Characterization and Outage Performance
     
	Thanh Luan Nguyen, Georges Kaddoum, Tri Nhu Do, Zygmunt J. Haas,




    
    
	T.-L. Nguyen and G. Kaddoum are with the Department of Electrical Engineering, the École de Technologie Supérieure (ÉTS), Université du Québec, Montréal, QC H3C 1K3, Canada. 
    G. Kaddoum is also with Artificial Intelligence & Cyber Systems Research Center, Lebanese American University (emails: thanh-luan.nguyen.1@ens.etsmtl.ca, georges.kaddoum@etsmtl.ca).

	T. N. Do is with Electrical Engineering Department, Polytechnique Montréal, Montréal, QC, Canada (email: tri-nhu.do@polymtl.ca).
	
	Z. J. Haas is with Department of Computer Science, University of Texas at Dallas, TX 75080, USA, and also with the School of Electrical and Computer Engineering, Cornell University, Ithaca, NY 14853, USA (e-mail: haas@cornell.edu).

    Received August 5, 2024
==============================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Point cloud compression is essential to experience volumetric multimedia as it drastically reduces the required streaming data rates.
Point attributes, specifically colors, extend the challenge of lossy compression beyond geometric representation to achieving joint reconstruction of texture and geometry.
State-of-the-art methods separate geometry and attributes to compress them individually. 
This comes at a computational cost, requiring an encoder and a decoder for each modality.
Additionally, as attribute compression methods require the same geometry for encoding and decoding, the encoder emulates the decoder-side geometry reconstruction as an input step to project and compress the attributes.


In this work, we propose to learn joint compression of geometry and attributes using a single, adaptive autoencoder model, embedding both modalities into a unified latent space which is then entropy encoded.
Key to the technique is to replace the search for trade-offs between rate, attribute quality and geometry quality, through conditioning the model on the desired qualities of both modalities, bypassing the need for training model  ensembles.
To differentiate important point cloud regions during encoding or to allow view-dependent compression for user-centered streaming, conditioning is pointwise, which allows for local quality and rate variation.
Our evaluation shows comparable performance to state-of-the-art compression methods for geometry and attributes, while reducing complexity compared to related compression methods.









§ INTRODUCTION






Point clouds are a popular representation for immersive multimedia content <cit.>, benefiting from low processing overhead compared to other formats for three-dimensional data such as meshes.
The idea of using point clouds as a data structure for streaming dynamical volumetric multimedia experiences over the Internet leads to significant bandwidth requirements, which establishes the need for efficient compression approaches.

Standardization proposes solutions for point cloud compression through G-PCC <cit.>, which uses octree data structures for the geometry and tailored approaches for the attributes like Region Adaptive Hierarchical Transform (RAHT) <cit.>. When considering point cloud sequences, V-PCC <cit.> utilizes video codecs to compress projections of the point cloud, thus allowing to exploit inter-frame correlations in 2D.



Inspired by and going beyond classical compression methods, Deep Learning has shown promising results for the compression of images <cit.>, videos <cit.> and point clouds <cit.>. 
Similar to the traditional approach, learned point cloud compression mostly treats geometry and attribute compression separately.
A geometry codec handles the compression of the points, while an attribute compression model handles the compression of colors given the knowledge of the undistorted geometry.


When combining the individual geometry and attribute codecs, whether learned or traditional, while relying on lossy compression, the reconstructed geometry at the decoder diverges from the geometry necessary as a prior for attribute decoding. 
This impacts the quality of the attribute reconstruction.
To address this issue, the geometry is decoded and the attributes are projected on the reconstruction <cit.> before attribute encoding, ensuring they are compressed using the same geometric prior which would be available at the receiver side.
We observe, however, that this introduces additional distortions and adds to the computational complexity of the encoding, requiring not only two individual encoding steps but additional geometry decoding and attribute projection, mostly done through costly nearest-neighbor searches, at the sender.

As depicted in Fig. <ref>, we propose a joint variable-rate compression method for point clouds in a single, adaptive autoencoder that embeds both modalities into a unified latent space. Figure <ref> shows that using this joint model, we can individually gauge the quality of the geometry and the attributes.
Using our single model dispenses with training an ensemble of models to fit a fixed number of individual tradeoffs between the bitrate and the qualities of the geometry and attributes.

The contributions of this paper are:

    
  * We formulate the rate-distortion loss for multimodal compression of geometry and attributes to incorporate adaptive weighting of conditional loss components. This allows to train a single adaptive model covering a wide range of geometry and attribute qualities and thus rates.
    
  * We propose a joint compression model, handling the geometry and attributes in one architecture, avoiding the need for geometry decoding and reprojection on the sender side.
    
  * Utilizing conditional encoding, attribute and geometry quality can be modulated independently, allowing variable rate-encoding with a single model, as well as, region-of-interest encoding of point clouds.






§ RELATED WORK


Attribute compression is essential for point cloud multimedia applications in contrast to solely compressing geometry, e.g., in LiDAR applications <cit.>.
Merkuria et al. <cit.> developed a codec utilizing octree partitions and intra-frame prediction while projecting attributes to 2D grids and compressing them with legacy image codecs.
MPEG standards <cit.> distinguish between static point cloud compression using Geometry-based PCC (G-PCC) <cit.> and dynamic point cloud compression with Video-based PCC (V-PCC) <cit.>.
While G-PCC relies on octree partitions for geometry compression and methods like Region Adaptive Hierarchical Transform (RAHT) <cit.> for attributes, V-PCC projects geometry and attribute information into video frames and leverages video codecs for compression. 
Driven by promising results in learned image compression <cit.>, learning based solutions for point cloud geometry or attribute compression showed competitive results in both domains.
However, most of these studies <cit.> do not consider the case where both, geometry and attributes, are subjected to distortions, as it would be the case for a realistic transmission use case. 

We note that early work explored learned compression of attributes and geometry in a single model <cit.>, however, learned compression methods using separate models for both modalities <cit.> were shown to outperform the single model approach.
The latter further boosts performance through wrapping RAHT <cit.> implemented in G-PCC <cit.> in the attribute branch and applying adaptive filters <cit.>, which need to be optimized during each encoding, adding to the already large computational cost.



 §.§ Learned Geometry Compression

Learning-based methods demonstrated competitive performance for point cloud compression, where early work <cit.> proposed to use 3D convolutions on voxelized point clouds and leverage learned entropy models <cit.> to optimize a rate-distortion loss.
The reconstruction is formulated as a binary classification problem using the (weighted) Binary Cross Entropy or the Focal Loss <cit.> to handle class imbalance due to point clouds sparsity in 3D space.
The authors of <cit.> employed entropy models from <cit.>, allowing to model spatial correlations combined with auto-regressive prediction based on already decoded coefficients.
Utilizing the above reviewed contributions, the authors of <cit.> employ a Variational Autoencoder (VAE) architecture based on Voxception-ResNet <cit.>, allowing to show on-par performance with V-PCC <cit.> in terms of geometry compression. 
We note, however, as all previously mentioned methods utilize 3D convolutions on sparse data represented in dense data structures, processing the point cloud in cubes is necessary to handle the memory requirements at the cost of limiting throughput and performance.
To overcome this, sparse convolutions <cit.> have been employed in <cit.> to process large-scale point clouds in one pass and to exploit global context during hyper-encoding, while a multi-scale classification in each stage of the decoder prunes non-occupied voxels.
Instead of using an encoder-decoder approach, <cit.> propose to encode the upsampling for each stage individually while relying on group-based decoding to leverage correlations between upsampled voxels. The model in <cit.> is then extended for inter-coding of dynamic point cloud sequences <cit.>.



 §.§ Learned Attribute Compression

Attribute compression is still dominated by traditional approaches, relying on Graph Transform <cit.> or RAHT <cit.>, with the latter being part of G-PCC.
In the field of learning-based approaches, models using sparse convolutions achieved comparable performance <cit.>. 
Combining traditional methods with learning-based approaches, the authors of <cit.> estimate the probabilities of transform coefficients from RAHT with deep entropy models, further improving compression efficiency.
The attribute compression branch in YOGA <cit.> relies on a thumbnail point cloud compressed at low resolution through RAHT to then predict residuals of the attributes through a conditioned autoencoder similar to <cit.>.  
To further increase the quality of the reconstructed attributes, YOGA <cit.> employs adaptive filters <cit.>, which are optimized during encoding and the parameters are added as side information to the bitstream to be used at the receiver.




 §.§ Variable-Rate Control for Learned Compression

Numerous compression models for images <cit.> and point clouds <cit.> based on the non-linear transform coding framework <cit.> achieve remarkable rate-distortion performance.
A major drawback of these methods is that they are conditioned on a fixed tradeoff between compression and reconstruction at training time.
As a result, an ensemble of models needs to be trained when multiple qualities are required, as it is commonly done in adaptive streaming <cit.>. 
Additionally, this further reduces flexibility as encoding configurations need to be fixed through hyper-parameters at training time. 

Variable-rate control and progressive coding schemes are promising techniques to freely adapt the rate-distortion tradeoff using a single model for encoding and decoding.
For example, for learned image compression, progressive coding is realized by iteratively refining the quantization grids <cit.>. 
Trit-planes <cit.> or dead-zone quantizers <cit.> extend this idea to better fit the distribution of quantization errors.
Following a different approach, quantization or reconstruction results are compressed in cascading <cit.> or recurrent compression <cit.> models, allowing for progressive coding. 
In contrast to these architectural decisions, training with uniformly sampled tradeoff parameters λ while dropping a corresponding number of feature channels in the entropy bottlenecks allows progressive coding through sorting feature channels by importance <cit.> and selecting a number of coefficients to decode.

Compared to progressive methods, adaptive rate-control targets conditioning a single model on a rate variable to allow controlling the rate-distortion of the compressed representation.
Early work such as in <cit.> introduces a gain control layer, multiplying and shifting channels of the feature representation to control the quantization before encoding the transformed coefficients in the entropy bottleneck. 
By incorporating multiple gain control layers in the encoder and decoder, features are scaled hierarchically <cit.>. 
Additionally, invertible layers allow increasing the reconstruction quality <cit.>.
YOGA <cit.> adopts these ideas for variable-rate compression with the layers proposed in <cit.> using sinusoidal embeddings, but requires an ensemble of entropy models as well as selecting appropriate rate parameters for the wrapped color compression using G-PCC.

Finally, we note that works such as <cit.> emphasizes regions of interest in images by conditioning compression models on 2D quality maps instead of a single scalar quality parameter.
As the quality map is not available on the decoder side, it is inferred from the compressed representation for decoding.
Extending this idea from image compression to point cloud compression is promising as it allows view-dependent encoding of three-dimensional content as done with traditional codecs <cit.>.



§ CHALLENGE AND APPROACH





 §.§ Problem Statement



We consider the problem of compressing a point set 𝒫 = (𝐠, 𝐚) containing a set of point locations 𝐠, denoted geometry, and corresponding attributes 𝐚. 
Specifically, for the geometry, quantized coordinates 𝐠={ (x_i,y_i,z_i) : x_i,y_i,z_i ∈ℕ_0,x_i,y_i,z_i < R } with maximum value R∈ℕ are assumed[This allows using sparse data structures and thus memory and computational efficient convolutional neural networks.]. Note that for real-valued coordinates, a voxelized representation could be easily derived by rounding and partitioning the full 3-dimensional space into blocks with size R^3.
We refer to attributes (and features) of a sparse tensor at the geometrical location u ∈𝐠 as 𝐚_u.


We observe that state-of-the-art compression methods <cit.> handle both modalities, i.e., geometry and attributes, separately.
Specifically, methods such as <cit.> or  <cit.> solely compress geometry or attributes, respectively. 
Methods that compress both modalities such as <cit.> still resort to separating the individual compression procedures for geometry and attributes. 
In essence, they encode the geometry first, which is then decoded to project the attributes to the possibly distorted geometry after reconstruction.
Finally, they compress the projected attributes on the distorted geometry. 
On the decoder side, the geometry is decoded and then used as side-information to allow attribute decoding. 
While this process allows compression of point cloud attributes under geometric distortion, it has a number of disadvantages: i) Encoding and decoding requires a sequence of non-parallelizable steps, resulting in high latency and ii) reprojecting the attributes on distorted geometries at encoding introduces additional distortions before compression.

Now, we aim at training a model for joint attribute and geometry compression.
In a Lagrangian fashion, we seek to minimize the rate ℛ under the quality constraints on geometry 𝒟'_G and attributes 𝒟'_A given the  loss function 

    ℒ = ℛ + λ_A𝒟'_A + λ_G𝒟'_G  ,

with corresponding deterministic, positive scalars λ_A,λ_G.
This can also be understood as finding a tradeoff between geometry and attribute quality at a given rate.
Compared to compression models that handle only a single modality in form of ℒ = ℛ + λ𝒟, with fixed hyperparameter λ, the loss function above adds another dimension to the number of models required to train, tune and deploy, making ensembles of joint models inadequate for most use cases. 
Hence, a key to solving this challenge is providing a single, adaptive model. This stands in contrast to the image compression literature <cit.> that coined the rate-distortion tradeoff in a bijection sense between the rate R and the quality D. In this work, this bijection does not exist as the rate results from defining a tradeoff between geometry and attributes.

A seemingly promising approach is variable rate compression <cit.>, in the sense

    ℒ = ℛ + 𝒟_A + 𝒟_G  ,

where we condition a single model for compression at multiple rates, thus only requiring a single training run and reducing the memory requirements at deployment.
In contrast to (<ref>), 𝒟_A,𝒟_G now integrate locality weighting functions  λ_A, λ_G instead of scalar quality parameters λ_A,λ_G allowing to define a location-specific tradeoff between the geometry, attributes and the rate. 
This is especially suited for multimedia point clouds, as it allows to differentiate important regions such as faces, or to offer view-dependent encoding. 
Additionally, instead of searching for tradeoffs between both modalities and the rate in (<ref>) at training time, e.g. through grid search, attribute and geometric quality can be balanced post-training. 



 §.§ Intermezzo: From Non-Linear Transform Coding to Conditional Transform Coding


Next, we briefly review the non-linear transform coding framework from <cit.> in addition to its extension towards conditional variable rate coding <cit.>.
The following review, also depicted as building blocks in Fig. <ref>, considers an arbitrary data vector 𝐱. 
In a later section, we show our building blocks where some are based on this framework. 

Figure <ref> shows the operational diagram of the mean-scale hyperprior <cit.> where an input data vector 𝐱 is non-linearly transformed through an encoder g_a. 
The transform output, i.e. the coefficient vector 𝐲, is quantized as ŷ = Q(𝐲) and fed to the decoder g_s producing the reconstruction 𝐱̂ = g_s(ŷ).
Note that the quantized coefficients are usually entropy coded for transmission while the encoder and decoder non-linear transforms g_a, and g_s are learned. 

Entropy coding the quantized coefficients ŷ requires an estimate of their distribution. 
Given the observation that the coefficients vary locally in mean and scale, a hyperprior model allows to parameterize a local estimate in form of a Gaussian distribution <cit.>. 
A hyperprior 𝐳=h_s(𝐲) is inferred from the coefficients 𝐲 as depicted in  Figure <ref>.
The quantized hyperprior ẑ is then entropy coded using a factorized entropy model <cit.> and used to locally reconstruct the parameters of the Gaussian distributions necessary for encoding/decoding ŷ using the hyperdecoder h_s.
Since backpropagation through a quantization operation results in gradients becoming zero almost everywhere, quantization is replaced through additive uniform noise <cit.> for training, i.e. ỹ =𝒬̃(𝐲) = y + Δ with Δ being an independently and identically distributed uniform random variable, Δ∼𝒰_[-0.5, 0.5]. In the following, we refer to all variables affected by this quantization proxy through (̃·̃)̃.

Extending this approach to allow variable rate coding, the encoder g_a in Fig. <ref> is now conditioned on a variable 𝐜 which is inferred through the conditional encoder p_a from a quality variable 𝐪 <cit.>.
This quality variable allows to locally dictate the quality of corresponding elements in the data vector 𝐱.
As the quality variable is not available for decoding, a proxy quality variable 𝐪̂ is inferred from the hyperprior ẑ.
This allows to derive the condition variable ĉ required in the synthesis transform g_s.





 §.§ Multimodal Conditioning for Geometry and Attributes


Next, we show how to achieve variable rate compression of geometry and attributes in a single adaptive model through integrating the multimodal rate-distortion loss in (<ref>) into the conditioned transform coding (cf. Sec. <ref>). 
To this end, we provide a per-point quality map 𝐪=[𝐪^(G), 𝐪^(A) ] to dictate the target quality individually for both modalities, i.e. geometry and attributes, to the model.
The quality map instantiates the quality variable 𝐪 in Sec. <ref> residing on the object geometry 𝐠. 
During training, each element of the quality map, which is set to 𝐪^(A)_u, 𝐪^(G)_u ∈ [0, 1], with u ∈𝐠, is mapped to a corresponding weighting map element λ^(A)_u, λ^(G)_u through a monotonically increasing function <cit.>.
This allows to individually define the range of the weights in λ^(A),λ^(G) and thus the target distortions for the respective modality. 
We use a quadratic weighting function[We note that exponential functions can be used as in <cit.>, however, experiments show that a root function is more suitable.], i.e. T_A: [0, 1] ↦[λ_min, λ_max] with T_A(x) = (λ^A_max - λ^A_min) x^2 + λ^A_min, such that λ^(A)_u = T_A(𝐪^(A)_u), with u ∈𝐠.
The same mapping holds for λ^(G)_u.
The range [λ^A/G_min, λ^A/G_max] defines the minimum and maximum values in the form of distinct hyperparameters.
At training time, quality maps 𝐪 are either uniformly drawn or supplied as a gradient in a random direction, allowing the model to vary the compression ratio both globally and locally.



  §.§.§ Geometry Loss


For the geometry reconstruction, we use the Focal Loss <cit.>
at every level of the decoder as proposed by <cit.> to allow iterative upsampling of the features and geometry in K steps[corresponding to K resolutions also denoted as scales throughout this work.].
In a nutshell, iterative upsampling in K steps is required for constructing the point cloud geometry from the downsampled representation of the encoder. 
More specifically, in each upsampling step k ∈{0,…,K-1}, a set of possible point locations 𝐝̃^(k), is generated from the lower scale geometry 𝐠̃^(k+1) s.t. 𝐠̃^(k+1)⊂𝐝̃^(k) <cit.>[Later, we omit the scale superscript for k=0, i.e. the original resolution 𝐠̃^(0):=𝐠̃, 𝐠^(0):=𝐠.].
Finally, an occupancy probability 𝐩̃_u^(k) at location u and scale k is inferred for the possible point locations 𝐝̃^(k).
Incorporating the geometry loss map λ_u^(G, k) in each scale k, the distortion loss 𝒟_G in (<ref>) is

    𝒟_G = - ∑_k=0^K-1∑_u ∈𝐝̃^(k)λ_u^(G, k) (1 -𝐩̇_u^(k) ))^γlog(𝐩̇_u^(k)) 
        with 𝐩̇_u^(k) = 𝐩̃_u^(k)   if  u ∈𝐠^(k)
    
            1 - 𝐩̃_u^(k)   else  ,

where 𝐠^(k) is the ground truth geometry  at scale k and the hyperparameter γ known from the focal loss <cit.> reduces the gradients for well-classified samples, allowing the model to emphasize harder samples.
As the loss map λ_u^(G, k) is only given on ground truth voxel geometry, i.e. u ∈𝐠, substitute values λ_u^(G, k) are derived through downsampling and upsampling using average pooling and transposed average pooling, respectively, resulting on the same base geometry as in 𝐝̃^(k). 





  §.§.§ Attribute Loss


The attribute reconstructions is typically cast as an element-wise regression problem <cit.>.
Using the L_2-norm the attribute component 𝒟_A in  (<ref>) becomes

    𝒟_A = ∑_u ∈𝐠∩𝐠̃λ_u^(A)‖𝐚_u - ã_u ‖_2^2   ,

where  𝐚_u denotes the attributes at the geometrical location u ∈𝐠 and ã_u, 𝐠̃ are, as discussed in Sect. <ref>,  affected by the quantization proxy. Recall that the attribute loss map λ_u^(A) contains individual, i.e. local, weighting to gauge the attributes to target distortion. 

Note that the loss can only be computed on locations u ∈𝐠∩𝐠̃, i.e. the intersection of ground truth geometry 𝐠 and the reconstruction 𝐠̃, where we compare the ground truth attributes 𝐚_u with the reconstruction 𝐚̃_u.






  §.§.§ Rate Loss

The rate loss ℛ is computed according to <cit.>, aiming to minimize the cross-entropy as in

    ℛ =    - ∑_u ∈𝐠^(K)∑_i log_2(p_ỹ |  μ̃, σ̃(ỹ_u,i |  μ̃_u,i, σ̃_u,i))
       - ∑_u ∈𝐠^(H)∑_ilog_2(p_𝐳̃ | ϕ(𝐳̃_u,i  | ϕ^(i))),

where Φ={ϕ^(i)}_i are the learnable parameters of the factorized entropy model <cit.>. 
We use ỹ_u,i and 𝐳̃_u,i to denote the transformed coefficient of ỹ and 𝐳̃ at the i^th channel in the feature dimension at the geometry position u, respectively.

In addition, μ̃_u,i, σ̃_u,i denote the local estimates of the means and scales that are obtained from the hyperprior 𝐳̃ as in (μ̃, σ̃) = h_s(𝐳̃).
Hereby, the hyperprior 𝐳̃ capturing the estimates resides on lower scale geometry 𝐠^(H), i.e. H > K.




 §.§ Model Architecture



We propose a sparse neural network architecture as depicted in Fig. <ref>, following the framework proposed by <cit.> (cf. <ref>).
Recall the notation conventions from Sec. <ref>, with (̂·̂)̂ denoting variables affected by quantization (i.e. inference time), and their training time counterparts (quantization proxy) given as (̃·̃)̃. 

The model receives a point cloud 𝒫 = (𝐠, 𝐚).
The encoder g_a downsamples the point cloud 𝒫 to the (resolution) scale k ∈{0,…,K-1}, defined in Sect. <ref>, in combination with the quality map 𝐪 residing on the same geometry 𝐠, to obtain a transformed representation (ĝ^(K), ŷ).

The key to rate adaptive feature transformation at each scale k is the Conditional Feature Extraction (CFE) block (see Fig. <ref>) in the encoder that takes the quality map 𝐪 as input.
On the decoder side, the surrogate[we call estimates at the decoder side surrogate.] quality map 𝐪̂^(𝐊) required for adaptive decoding at the lowest scale[in this work we set the number of resolution scales to K=3.] is inferred from the hyperprior ẑ and the transformed coefficients ŷ. 
The second ingredient to our method is the Upsampling and Prunining (UP) block during decoding (see Fig. <ref>) which prunes the geometry at each scale k to maintain the sparsity of the input geometry, thus overcoming memory restrictions <cit.>.





  §.§.§ Conditional Feature Extraction on Geometry and Attribute Importance
 

We employ local conditioning of feature representations through a scale and shift operation, following common procedure to allow variable-rate control in image <cit.> and point cloud compression <cit.>.
Hereby, features 𝐭^(k) residing on the geometry 𝐠^(k) at a given scale k are transformed through the local scale and shift operation into 𝐭'^(k) as

    𝐭_u'^(k) = α_u^(k)𝐭_u^(k) + β_u^(k)  ∀ u ∈𝐠^(k).

using the scales α_u^(k) and shifts β_u^(k). 
In contrast to the related work <cit.>, which only predicts element-wise scales α_u^(k) and shifts β_u^(k) from a single quality map value 𝐪^(k)_u per geometry location u, joint adaptation of the geometry and attribute quality requires learning a mapping from geometry and attribute quality to the local transformation variables. 

In essence, the feature extraction layer learns the importance of features for optimizing a given task by varying its scale and shift.
This results in broader distributions at the bottleneck when aiming for higher rates, thus reducing quantization errors.
Vice versa, dictating low qualities through 𝐪^(k) results in more narrow distributions or in disabling certain feature channels completely through trivial distributions.
Additionally, the additive shifting operation allows to encode the quality map 𝐪 into the transformed coefficients 𝐲, embedding information required to infer the surrogate quality map from ŷ.



  §.§.§ Upsampling Sparse Features with Geometry Pruning
 

During encoding, lower scale geometry 𝐠^(k+1) is derived from 𝐠^(k) through rounding in a dyadic scheme. 
As a result, when upsampling a representation ĝ^(k+1) in the decoder, the correct voxels to be occupied in 𝐠̂^(k) are unknown.
Hence, to prevent information loss, voxels are generated at each possible location through a Generative Transposed Convolution <cit.>, resulting in a dense geometry 𝐝̂^(k), which is a superset of 𝐠^(k).
While the continuation of decoding with 𝐝̂^(k) is conceptually possible, it leads to significant increase in memory and computational requirements.
Geometry compression methods overcome this challenge by performing multiscale classification <cit.>, pruning the upsampled geometry 𝐝̂^(k) at each scale k by inferring per-point occupancy probabilities 𝐩̂_𝐮^(k) for all u ∈𝐝̂^(k).
Note that attribute decoders, such as <cit.>, do not face this challenge by assuming the availability of undistorted geometry 𝐠^(k) from a preceding step, which is in practice not true at the decoder. 

Now, when combining both modalities, i.e. geometry and attributes, we need to predict the occupancy and prune the feature representation for the next scale.
The challenge hereby is maintaining an expressive feature representation after pruning under potential geometry miss-predictions for attribute reconstruction.
To increase robustness against these miss-predictions, we directly transform the dense, upsampled features, arriving at a dense feature representation, as shown in Fig. <ref>.  
Then, an occupancy prediction branch infers occupancy probabilities 𝐩̂^k_u from the transformed, dense feature representation 𝐟^(k).
Finally, top-k selection is performed, keeping the k most probable voxel locations in a pruning step on the transformed, dense features. 










§ EVALUATION RESULTS






 §.§ Evaluation Setup



  §.§.§ Implementation


We implement our approach in PyTorch, using MinkowskiEngine <cit.> to leverage sparse convolutions. 
For training, we collect a variety of human point clouds from the UVG dataset <cit.>, select a number of frames for each point cloud sequence and crop them into cubes with edge length 128, resulting in 18816 distinct samples.
Opposed to projecting image textures on the samples during training <cit.>, we solely subject the attributes to color jitter and randomly rotate each sample. 
For each sample, a random quality map (𝐠, 𝐪) is derived according to Sec. <ref>. 
We train for 200 epochs using a batch size of 16.
Adam <cit.> is used for optimization with an initial learning rate of 10^-3, which we decrease with factor 0.1 every 50 epochs. 
To stabilize training, gradient norm clipping with threshold 1.0 is integrated.




  §.§.§ Baselines


We compare our method against YOGA <cit.>, which we identify as the current state-of-the-art method for learned compression of geometry and attributes. 
Additionally, we use the tmc2v24 test model for V-PCC <cit.>, using all-intra test conditions and the tmc13v23 implementation of G-PCC <cit.> employing octree compression for geometry and RAHT <cit.> for attributes. For both methods, we select the first 4 rate configurations specified in the respective test models, as we find that they result in a similar range of rates as our model.



  §.§.§ Evaluation Metrics


We measure the bitrate in bits per point (bpp), using the cardinality |𝐠| as divisor.
Geometric quality is assessed using the D1-PSNR <cit.> metric, computing the mean distances between points of the tested geometry 𝐠 and their nearest neighbor in the reference 𝐠̂. 
Utilizing the derived neighbor association, the Y-PSNR and YUV-PSNR metrics <cit.> are computed to assess the attribute quality, with the former only using the luminance channel and the later referring to the weighted average over all channels using weights (6/8, 1/8, 1/8), following evaluation protocols in related literature <cit.>. 
We consistently present symmetric metrics, i.e. computation is performed for swapped source and reference while reporting the stricter result.
Finally, PCQM <cit.> is used as a joint metric, incorporating geometric and attribute features. 
Note that for all presented results, the reported bitrates correspond to the compression of geometry and attributes together.

We select 4 frames each from the 8iVFBv2 <cit.> and MVUB <cit.> dataset for testing, which are specified in the supplementary material. 
If not specified otherwise, uniform quality maps 𝐪^(A) and 𝐪^(G) are used in our model to allow fair comparison.







 §.§ Perceptual Quality Optimized Pareto Front


Following the evaluation protocol specified in <cit.>, we traverse the cross product of quality configurations for our approach with step size 0.05, resulting in 400 configurations for the point clouds in 8iVFBv2 <cit.>. 
Then, the Pareto front for the tradeoff between PCQM <cit.> and rate is determined and depicted in Fig. <ref>. 
Two examples of the chosen configuration pairs can be found in Fig. <ref>, which shows the pairs of attribute quality q^(A) and geometric quality q^(G) contributing to the Pareto front as a path over all possible combinations. 
The underlying contour in each plot shows the bitrate, attribute quality and geometric quality, resulting from all possible combinations.
Similarly, Pareto fronts for V-PCC <cit.> and G-PCC <cit.> are determined by traversing the full range of attribute QPs for each geometry setting.
We also compare against the results reported in the original paper of YOGA <cit.> for the 8iVFBv2 <cit.> dataset.

From Fig. <ref> we find that our model shows improved PCQM performance compared to all baselines for medium to high rates, indicating comparable or better perceived quality of the reconstructions at this compression ratio. 
However, when aiming for lower rates, the quality deteriorates faster compared to related methods.
We attribute this to the fact that our approach is the only method solely relying on learned compression for the point clouds, while, e.g. YOGA <cit.> not only inherits low rate compression capabilities by compressing a thumbnail attribute representation using RAHT <cit.> but also employs adaptive filters <cit.> which are optimized during encoding on the specific content to counteract quality degradation.

We note that the shown results present optimized configurations for the respective point clouds.
This is also apparent when comparing the selected configuration pairs for the point clouds redandblack and loot contributing to the Pareto fronts in Fig. <ref>, which are visualized as paths in Fig. <ref>.
Depending on the individual attribute and geometric properties of the content, we notice that the paths and thus the chosen set of configurations to arrive at the Pareto fronts, differ significantly.
Recall that this evaluation required traversing a large number of possible configurations to select the best tradeoff points to arrive at the pareto fronts and is thus not practical.
Nevertheless, this evaluation provides insights into the capabilities of the tested compression methods.





 §.§ Fixed-Configuration Encoding Performance



Now, we set the conditioning of our model to a fixed set of configurations, i.e. (𝐪^(G)_u, 𝐪^(A)_u) ∈{(0.05, 0.1), (0.1, 0.2),  (0.2, 0.4), (0.4, 0.8)} ∀ u ∈𝐠, independent of the input point cloud.
This allows to compare against G-PCC <cit.> and V-PCC <cit.> using their respective standard configurations. 
We exclude YOGA <cit.>, as such a set of configurations is not publicly available.

Selected rate-distortion curves are presented in Fig. <ref> and corresponding Bjønte­gaard-Deltas to G-PCC <cit.> are shown in Tab. <ref>. 
We provide further rate-distortion curves in the supplementary material.
Note that our approach allows for significantly higher geometric quality according to D1-PSNR, thus small overlap of the resulting rate-distortion curves in the y-axis. This renders the Bjønte­gaard-rate Delta Δ_r for the D1-PSNR metric less expressive on a subset of the tested point clouds. 


Over all tested point clouds, our approach shows significantly higher geometric quality compared to G-PCC <cit.> and V-PCC <cit.>. 
The corresponding attribute rate-distortion curves according to the Y-PSNR metric show notable gains over G-PCC, but underperform when compared against V-PCC. 
Considering attribute quality according to the YUV-PSNR metric, Tab. <ref> shows similar results, with the exception of sequences with complex textures, i.e longdress, redandblack and andrew9.
For the Bjønte­gaard-rate metrics Δ_r and Δ_D of the joint metric PCQM <cit.>, we see high potential for rate-savings in the fixed configuration range compared to G-PCC <cit.>.








 §.§ Latency Analysis


Using a single, joint model in our approach allows to significantly simplify the encoding process, only requiring a single pass through the encoder. 
This stands in contrast to the standard procedure, which requires projection of the attributes to the reconstructed geometry during attribute encoding, as it is done in G-PCC <cit.>, V-PCC <cit.> and YOGA <cit.>.
Consequently, the encoding latency of our model (0.75s) in Table <ref> is more then one order of magnitude lower then the one for the implementations of traditional codecs <cit.>. 
When comparing to the learned approach in YOGA <cit.>, we also notice a lower encoding latency by one order of magnitude compared to the results reported in the paper.
We note that the original paper utilizes different hardware, however, the encoding speed advantage of our approach, due to the single pass, is apparent.

Considering the decoding latency, we find that our method is competitive with respect to the related approaches, although all methods may require further optimization to allow application in practical use cases.
Note, however, that the latency for all approaches significantly depends on the content (i.e. number of points) and the quality configuration. 
A more detailed, per-content and per-configuration analysis is given in the appendix.





 §.§ Showcase: View-dependent Compression


Recall that conditioning our model on a quality map (𝐠, 𝐪) instead of a scalar quality parameter allows to perform region-of-interest rate control.
To showcase this capability, we derive a view-dependent compression of a point cloud given some arbitrary user viewing angle.
This is achieved by instantiating quality maps using a gradient along the viewing direction, assigning higher qualities to the front of the point cloud which decreases to the back.
For comparison, we consistently apply the same procedure using a uniform quality map in Fig. <ref> and a quality map according to the step function in Fig. <ref>.
For all compressed point clouds in Fig. <ref>, the geometry quality map 𝐪^(G) is instantiated following an identical trend. 
The reconstructions are rendered from the given viewing-angle and standard image metrics, i.e. SSIM and the PSNR of the MSE in the YUV color space, are computed.
Considering the image metrics of the renders, we find some potential for reducing the rate by locally dictating the quality 𝐪 in both image metrics, i.e., PSNR and SSIM. 













§ CONCLUSION AND OUTLOOK

We presented a method for training a single, joint compression model for point cloud geometry and attributes.
Key to our approach is extending variable-rate compression techniques to multimodal compression, allowing to gauge the model on a tradeoff between the triplet of the quality of both modalities and the rate.
Equipped with this, we proposed a joint compression architecture operating on a sparse tensor representation. Our technique achieves comparable compression performance to the state of the art in learned point cloud compression <cit.> which, however, has significantly higher encoding latency as it requires geometry decoding at the sender to encode the attributes, while outperforming G-PCC <cit.> and further closing the gap to V-PCC <cit.>. 
Our adaptive, variable-rate model illustrates the local degrees of freedom and flexibility in the joint compression of geometry and attributes. 

However, compressing point clouds according to content specific configurations  allows to fully utilize our model capabilities, but requires computationally expensive grid search in the space of possible configuration pairs.
In some use cases, this may result in resorting to fixed configuration sets as evaluated.
Tackling a similar issue in video compression, per-title encoding <cit.> offers more efficient methods for optimizing a configuration set on specific content.

Deriving meaningful quality maps to locally condition the point cloud geometry and attribute quality using our method is a promising extension of this work. 
Our evaluation gives a perspective  on the capabilities of view-dependent and region-of-interest based compression.
Possible directions are incorporating per-point importance, e.g.  through harnessing fixation maps of the content <cit.> or incorporating task-specific optimization <cit.>. 





unsrtnat
 








§ ARCHITECTURE

We extend the visualization in Fig. <ref>, presenting a more detailed version in Fig. <ref> with implementation channel sizes for reproducability. 
All convolution layers, denoted Conv are three-dimensional sparse convolutions, for which we use Minkowski­Engine <cit.>.
For layers with the extension ↑, Generative Transposed Convolutions <cit.> are used, allowing to generate geometry locations and corresponding features during upsampling.
Note that the hyperdecoder h_s is prone to non-deterministic behaviour <cit.>, which arises when using Minkowski­Engine <cit.>.
To this end, we reduce the kernel size for generative convolutions in the hyperdecoder to 2, as we found that this allows to counteract non-determinism in combination with sorting point locations before and after the convolution.

The geometry representation 𝐠^(3) at the entropy bottleneck is encoded using G-PCC in lossless geometry mode, which is common practice <cit.>. 
The base geometry 𝐠^(5) required to decode the hyperpriors 𝐳 is derived through dyadic downsampling of 𝐠^(3).

In addition, the number of voxels to keep at each pruning stage in the UP block (cf. <ref>) is appended as side-information in the header of the bitstream.






§ ADDITIONS TO THE EVALUATION



 §.§ Test Data

The point cloud frames select for testing are presented in Tab. <ref>, consisting of the 8iVFBv2 <cit.> and MVUB dataset <cit.>.
While point clouds in the 8iVFBv2 dataset <cit.> consist of watertight surfaces and high quality geometry and textures at voxel resolution 1024, point clouds in MVUB dataset <cit.> are frontal scans of human upper bodies containing visible (geometry) noise at voxel resolution 512. 
Incorporating both datasets in the evaluation allowed us to assess the performance of our model over broad input conditions.





 §.§ Extended Results for the Latency

Complementary to the results presented in Table <ref> in Sec. <ref>, we showcase per content results of the encoding latency in Table <ref> and decoding latency in Table <ref>.
Since YOGA <cit.> is not open-sourced at the time of our evaluation, we can only report the averaged latencies from their results, which were measured on different hardware.
For all other approaches, we use the configurations for the generation of the pareto-fronts in Sec. <ref> and measure the times for each encoding and decoding.
We notice that the encoding and decoding latency of G-PCC <cit.> heavily relies on the quality configuration for the geometry and thus split the results for each tested geometry setting. 
It becomes apparent that the latency highly depends on the content, more specifically on the number of points. 

Overall, we observe similar behaviour as presented in Sec. <ref>: Our method allows significantly faster encoding times as other methods, which can be mainly attributed to the fact that the geometry does not need to be reconstructed for attribute re-projection during the encoding process. 
For decoding, we find our method performing comparable to related work. 

 
 

  



 §.§ Complementary Rate Distortion Results for Fixed-Configuration Encoding

We showcase the rate-distortion results on all point clouds according to the evaluation procedure described in Sec. <ref>, as only selected curves were shown in the main document. 
Fig. <ref> shows results for the 8iVFBv2 dataset <cit.> and Fig. <ref> shows results for the MVUB dataset <cit.>.
Overall, we find increased performance for the geometry metric D1-PSNR over V-PCC <cit.> and G-PCC <cit.>. 
For attribute quality our approach achieves a better rate-distortion tradeoff then G-PCC but falls short of the attribute qualities achieved by V-PCC. 
Especially the point cloud andrew9, which shows a checkerboard texture pattern on the shirt, is hard to compress using our learned approach. 
Finally, for the joint metric PCQM <cit.>, we find that our approach consistently achieves higher quality at the given rates compared to G-PCC <cit.>. 
For higher rate configurations, our method can deliver comparable and sometimes better quality as V-PCC <cit.>, but falls of when increasing the compression ratio.

Additionally, we notice that the range of compressed bitrates for our model is small compared to the rate-distortion rates achievable by both traditional methods. 












 Paper Title: 3D U-KAN Implementation for Multi-modal MRI Brain Tumor Segmentation.












Tang et al.





Department of Biostatistics, School of Global Public Health, New York University
 Equal contributions
* Correspondence: hs120@nyu.edu
 








3D U-KAN Implementation for 
Multi-modal MRI Brain Tumor Segmentation
    Tianze Tang*0009-0007-6244-5813 Yanbing Chen*0009-0009-1218-0129 Hai Shu**0000-0002-6968-4063
    
=================================================================================================




§ ABSTRACT 
We explore the application of U-KAN, a U-Net based network enhanced with Kolmogorov-Arnold Network (KAN) layers, for 3D brain tumor segmentation using multi-modal MRI data. 
We adapt the original 2D U-KAN model to the 3D task,  and introduce a variant called UKAN-SE, which incorporates  Squeeze-and-Excitation modules for global attention. We compare the performance of U-KAN and UKAN-SE against existing methods such as U-Net, Attention U-Net, and Swin UNETR, using the BraTS 2024 dataset. Our results show that U-KAN and UKAN-SE, with approximately 10.6 million parameters, achieve exceptional efficiency, requiring  only about 1/4 of the training time of U-Net and Attention U-Net, and 1/6 that of Swin UNETR, while surpassing these models across most evaluation metrics.
Notably,
UKAN-SE slightly outperforms  U-KAN.







§ INTRODUCTION


Gliomas are a common type of malignant brain tumors and  a leading cause of cancer-related deaths among adults. Diagnosing gliomas is particularly challenging due to their invasive nature and their ability to occur in any part of the brain <cit.>. Accurate detection of these tumors requires multi-protocol magnetic resonance imaging (MRI), which remains the gold standard for post-treatment imaging of gliomas. MRI provides crucial information on tumor size, location, and morphological changes. The most commonly-used MRI modalities include T1-weighted (T1), post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (FLAIR) <cit.>.


The Brain Tumor Segmentation (BraTS) Challenge, an annual event hosted by the International Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI) <cit.>, has been held since 2012. 
In the BraTS 2024 Challenge, we participated in Task 1: Segmentation–Adult Glioma Post Treatment <cit.>,  which focuses on evaluating advanced techniques for segmenting brain diffuse gliomas. 
Participants are provided with a  multi-modal MRI dataset to develop automated algorithms to accurately segment the  tumors. This initiative is valuable for creating tools to objectively assess residual tumor volumes, which is crucial for treatment planning and predicting patient outcomes.



 U-Net <cit.> has become 
 the benchmark network
for medical image segmentation, and U-Net based models have successfully demonstrated their efficacy in past BraTS challenges <cit.>. Originating from a straightforward design that features convolutional layers and optionally normalization layers, U-Net has been significantly enhanced through various modifications. For example, Attention U-Net <cit.> incorporates attention blocks to improve focus on relevant features,  while Swin UNETR <cit.> integrates a Swin Transformer as the encoder and connects it to a CNN-based decoder at various resolutions through skip connections.

While previous U-Net based models have made progress in medical image segmentation, they face significant challenges with conventional kernels. These kernels often struggle to capture complex nonlinear patterns and rely on empirical designs that lack interpretability, limiting their clinical trustworthiness. 
Recently, Kolmogorov-Arnold Networks (KANs) <cit.> have emerged as a promising alternative, using learnable activation functions on edges instead of fixed ones on nodes. This approach enhances the network's ability to model complex functions more accurately, improves interpretability, and more effectively links the network's structure to its performance.

Given the potential of KANs, U-KAN <cit.> has been developed by integrating KAN layers into the U-Net architecture. This network design includes a tokenized KAN block that captures complex patterns more effectively. Empirical evaluations of U-KAN have shown its superior performance in medical segmentation tasks and its potential as an advanced alternative in various visual applications.



In this paper, we compare five segmentation models
on the BraTS 2024 Challenge's Task 1,
including 
U-Net <cit.>, Attention U-Net <cit.>, Swin UNETR <cit.>, U-KAN <cit.>, and 
our  UKAN-SE,
a new U-KAN variant with
  Squeeze-and-Excitation
modules <cit.> for global attention. We evaluate these models based on total parameter size, training time, 
and segmentation accuracy. 




§ METHODS





 §.§ Dataset

The dataset for Task 1 of the 2024 BraTS Challenge <cit.> comprises 1,350 labeled training samples and 188 unlabeled validation samples. Each sample includes four modalities (T1, T1Gd, T2, and FLAIR), as illustrated in Figure <ref>. The data is derived from multi-institutional, routine post-treatment clinically-acquired multi-parametric MRI (mpMRI) scans of glioma patients. The MRI scans undergo preprocessing to align them to a standard anatomical template, normalize the resolution to 1  mm^3, and remove the skull. The final dimensions of the processed scans are 218 × 182 × 182.





The ground-truth segmentation labels for the validation set are hidden from the Challenge participating teams, who can only access  the Dice scores and the 95% Hausdorff distance of their segmentation results via the online participation platform. The labels for the MRI scans include non-enhancing tumor core (NETC), surrounding non-enhancing FLAIR hyperintensity (SNFH), enhancing tissue (ET), and resection cavity (RC), as shown in Figure <ref>.





 §.§ Data Augmentation


Data Processing. Our data augmentation pipeline for MRI images incorporates several transformations to enhance the diversity of the training data. We combine the four modalities into a single 4D image (C × H × W × D, where C = 4) and save it, along with the segmentation labels, in the  file format. Each image is stored in 32-bit floating-point format, while the labels are stored as 8-bit unsigned integers. The grayscale values of each image are normalized, ensuring a background value of zero. This process helps reduce the file size.



Data Enhancing. 
The augmentation process begins with a fixed crop of 192 × 160 × 160 to ensure a consistent input size, followed by random flips along different axes to introduce spatial variability. Gaussian noise is added to simulate variations in image quality, and random rotations are applied to account for potential misalignments. Additionally, the contrast of the images is adjusted to vary the intensity distribution. Finally, the images and labels are converted into tensor format, preparing them for further processing in the neural network. This augmentation strategy enhances the model's robustness.





 §.§ Network Architectures




Loss Function.
The loss function 
is given in equation (<ref>),
which combines the soft Dice score and Cross Entropy (CE) loss, with uniform weighting applied across all labels in the CE loss. We experimented with different weights for the CE loss, but no improvement in training performance was observed. The loss is calculated on a batch-by-batch basis. Specifically,

    ℒ_total = (1 - α) ·ℒ_CE + α· (1 - Soft Dice),

where α is a weight factor, and we use 0.5. The softmax output excluding the background class is used as the input for the soft Dice loss, and the one-hot encoded target excluding the background class is used as the target.



Previous U-Net based models.
We consider
three U-Net based models, including the classic U-Net <cit.>, Attention U-Net (Att-Unet) <cit.>, and Swin UNETR <cit.>, 
which have demonstrated strong performance in previous challenges. The U-Net architecture captures features at multiple resolutions through its downscaling and upscaling layers (encoding and decoding), making it particularly suitable for tasks like brain tumor segmentation. Att-Unet  incorporates attention gates, enhancing segmentation accuracy. Swin UNETR combines a U-shaped network design with a Swin Transformer encoder, utilizing multi-scale attention mechanisms to capture both global and local contextual information, making it effective for complex pattern recognition in medical images.



U-KAN.
The U-KAN model <cit.> integrates KAN layers <cit.> into the traditional U-Net structure. In our experiment, we use the default U-KAN configuration, with channel sizes D_1 = 128, D_2 = 160, and D_3 = 256 (see Figure <ref>). The architecture employs a two-phase design: a Convolution Phase for initial feature extraction, followed by a Tokenized KAN (Tok-KAN) Phase where the KAN layers refine the feature representations. This design enables U-KAN to effectively capture complex non-linear patterns, making it particularly suitable for medical image segmentation tasks. The KAN layers are mathematically described as

    KAN(𝐙) = Φ_K-1∘Φ_K-2∘⋯∘Φ_0(𝐙),

where Φ_i 
is the i-th KAN layer
consisting of learnable activation functions, enhancing the model's ability to capture intricate patterns in medical images.

Given the strong performance of U-KAN  in 2D image segmentation, we hypothesize that it can also excel in 3D tasks like brain tumor segmentation. Thus, we adapt it to a 3D version by replacing the 2D convolutional and batch normalization blocks with their 3D counterparts, as shown in Figure <ref>. Notably, the KAN itself does not impose any dimensional constraints, as all input data are transformed into 1D channels, facilitating a straightforward transition to 3D.






UKAN-SE.
We introduce a new U-KAN variant,
called
UKAN-SE,
which incorporates
 a  Squeeze-and-Excitation (SE) module <cit.>  after each convolutional block to provide global attention and enhance feature representations.
Figure <ref> illustrates the network structure.

















 §.§ Evaluation  Metrics


Besides the 
Dice score  
and 95% Hausdorff distance (HD95)
for the whole image segmentation, 
the lesion-wise Dice and lesion-wise HD95 are used in the BraTS 2024 Challenge <cit.> to evaluate the performance of segmentation models. The
two lesion-wise metrics are defined as


    Lesion-wise Dice = ∑_i=1^LDice(l_i)/TP + FN + FP,



    Lesion-wise HD95 = ∑_i=1^LHD95(l_i)/TP + FN + FP,


where L is the number of ground-truth lesions, l_i denotes the i-th lesion, TP is the number of true positives, FN is the number of false negatives, and FP is the number of false positives.

The lesion-wise Dice measures the overlap between predicted and ground-truth segmentations, excluding true negative voxels. The lesion-wise HD95  evaluates the distance between the closest points of the predicted and ground-truth segmentations, emphasizing the most significant segmentation errors. The lesion-wise metrics provide a detailed understanding of a model's ability to accurately segment complex and multi-focal lesions, and help prevent bias toward models that excel only in detecting larger lesions.

The Dice and HD95 metrics are computed 
for each of the six tumor sub-regions:
non-enhancing tumor core (NETC), surrounding non-enhancing FLAIR hyperintensity (SNFH), enhancing tissue (ET), resection cavity (RC),
 tumor core (TC = ET + NETC),
 and  whole tumor (WT = TC +
SNFH).
 
 To compare the efficiency of the five segmentation models, we report the total number of  parameters and calculate the average training time per epoch. 
All models are implemented on 
 a NVIDIA RTX8000 GPU (48GB memory).




 §.§ Training Details


To find the optimal hyperparameters for U-KAN, we aim to achieve the highest lesion-wise Dice and smallest lesion-wise HD95 scores to ensure the best performance. The best-performing set of hyperparameters includes a combination of batch size, epochs, optimizer, weight decay and learning rate.  Due to the large-scale BraTS images, we set the batch size to 2. Model training is conducted for 100 epochs, beginning with a 10-epoch warm-up period. The AdamW optimizer <cit.> is utilized, with a weight decay of 0.0001. The learning rate starts from 0.005, reaches a maximum value of 0.01 after 10 warm-up epochs, and then gradually decreases to 0.005 through cosine annealing <cit.>. For consistency, this set of hyperparameters is also applied to the other four models in our study.












§ RESULTS






The lesion-wise Dice and whole-image Dice results in Table <ref> indicate that UKAN-SE  achieves the best performance across various lesion types, particularly excelling in ET and RC. It also shows consistent high scores in WT segmentation.  Att-Unet  performs well in some areas, notably in SNFH and WT, but does not surpass  UKAN-SE in overall performance.  U-KAN  also demonstrates strong performance, with competitive scores across multiple categories. 


In Table <ref>, which  presents lesion-wise HD95 and whole-image HD95 metrics, the UKAN-SE method continues to lead, showcasing lower values in most lesion types, indicative of better boundary delineation and contour accuracy. This method is particularly strong in ET, RC, and WT segmentation. The U-KAN method also shows commendable performance, particularly in RC and SNFH. 


Table <ref> shows
the parameter size and training time
for each
of the five models.
Notably,
while U-KAN and U-Net have comparable parameter sizes, with 10.612 million and 10.743 million parameters, respectively, U-KAN achieves significantly faster training time, averaging only 803 seconds per epoch  compared to U-Net's 3322 seconds per epoch. Although UKAN-SE has slightly more parameters and requires slightly more training time than U-KAN, it generally performs better in segmentation accuracy 
due to its global attention mechanism. An example of segmentation results is presented in Figure <ref>.


























§ DISCUSSION


This study evaluated the performance of the latest KAN-enhanced U-Net model, U-KAN, on  3D multi-modal MRI brain tumor segmentation using the BraTS 2024 dataset. The results indicate that the U-KAN structure achieves strong performance and exceptional efficiency in the brain tumor segmentation. Compared to traditional models like U-Net and Swin UNETR, the 3D U-KAN model delivers comparable performance  with approximately only 1/4 to 1/6 of the computational cost, highlighting the significant potential of the KAN structure.

In this study, we implemented only simple data augmentation techniques and did not utilize ensemble learning, unlike the more complex approaches described by the past challenge's winner <cit.>. A medium-sized KAN layer configuration ([D_1,D_2,D_3] = [128, 160, 256]) was used. Future research may explore varying these dimensions to assess their impact. Besides, the KAN structure's inherent advantages <cit.> in interpretability suggest promising avenues for further investigation in image segmentation tasks, potentially enhancing the understanding of model predictions.















§ ACKNOWLEDGEMENTS

The data used in this publication were obtained as part of the Challenge project through Synapse ID (syn51156910). This work was partially supported by the NYU IT High Performance Computing resources, services, and staff expertise.

















































































































































 Paper Title: Sentinel-1 SAR Based Weakly Supervised Learning For Tropical Forest Mapping.



Femtosecond switching of strong light-matter interactions in microcavities with two-dimensional semiconductors
    Giulio Cerullo
    August 5, 2024
==============================================================================================================


§ ABSTRACT 
Tropical forests play an important role in regulating the global carbon cycle and are crucial for maintaining the tropical forest biodiversity. Therefore, there is an urgent need to map the extent of tropical forest ecosystems. Recently, deep learning has come out as a powerful tool to map these ecosystems with the caveat of curating high quality reference datasets. Since, manually annotating high quality reference datasets is time consuming and expensive, weakly supervised learning techniques offer the potential to train high quality models without the need for manually annotating large quantities of reference datasets. In this manuscript, we propose two weakly supervised approaches that are based on Sentinel-1 SAR images, sparsely distributed pixel-wise high quality reference labels and densely distributed noisy reference labels. The proposed approaches were tested in a tropical setting in the Brazilian amazon. The results demonstrate that high quality tropical forest maps can be derived from weakly supervised learning without the need for manually annotated labels.  


Sentinel-1 SAR, Weakly supervised Learning, Unet, Deep Learning




§ INTRODUCTION



Tropical forests play an important role in regulating the global carbon and water cycle and provides bio-diversity and ecosystem services <cit.>. Therefore, mapping tropical forest helps in the quantification of carbon stocks to mitigate climate change and also enables timely interventions such as policy or law enforcement response to illegal deforestation. For this purpose, optical satellite images were widely used to map tropical forests but were often limited by pervasive cloud covers existing in these regions. In this regard, the advent of Synthetic Aperture Radar (SAR) sensors was a game changer as it allowed the acquisition of images in almost all weather, enabling the mapping of forests in any locations irrespective of cloud cover.




The recent success of deep learning in computer vision proliferated its application in the field of remote sensing for the purpose of landcover mapping <cit.>. In particular, deep fully convolutional neural networks such as Unet <cit.> has been ubiquitous in remote sensing based semantic segmentation tasks. However, one of the main challenges in applying these deep learning architectures was the lack of high quality reference data that are required for training. In their absence, users often resort to manual annotation to synthesize high quality reference data. Recently,  weakly supervised learning (WSL) approaches were adopted to train a high quality model without the need to annotate large quantities of reference datasets. Weakly supervised machine learning approaches are techniques that use higher level and often noisier labels to synthesize a large quantity of reference labels by iterative refinement <cit.>. WSL are often categorized into three groups as incomplete supervision, inexact supervision or inaccurate supervision <cit.>. Incomplete supervision refers to the case when the available reference labels are precise but few in number to train a high quality model and there exists sufficient supply of unlabeled data. On the other hand, inexact supervision happens when the labelling is coarse or high level such as image level annotations whereas the required level is pixel level annotations. Inaccurate supervision indicates the reference labels are crowd-sourced datasets that are often noisy, thereby requiring iterative refinement before application. 


A few notable works applying WSL to remote sensing data include the work by the authors in <cit.>  to apply class activation maps and masked Unet approach to segment croplands. The authors in <cit.> used image level labels in the inexact supervision category for the multi-class landcover segmentation of aerial and satellite images. Whereas, the authors in <cit.> demonstrated the effectiveness of WSL using high resolution satellite images for building extraction. 


In this manuscript, we demonstrate and compare two weakly supervised approaches based on the incomplete supervision and inaccurate supervision to derive a forest non-forest cover map in the Brazilian Amazon using Sentinel-1 SAR images.
  



§ METHODOLOGY





 §.§ Incomplete supervision


We use a sparse forest/non-forest pixels that were extracted from the intersection between the JAXA ALOS-PALSAR2 based global four class forest/non-forest map <cit.>, Hansen global forest cover/ forest change dataset <cit.> and the Land Use Change Alerts <cit.> recently released to the public through an interactive app. We apply a masked approach to randomly extract 2% of pixels in the reference label patch to synthesize the reference labels. The extracted labels were visually verified with a Planet scope image before training a Unet-like model for semantic segmentation of Sentinel-1 SAR images  (Figure <ref>b) to map forest cover.  In this way, we force the network to learn from the 2% of pixels available within the patch.  Therefore, the deep learning objective can be formulated as:


    argmin_θ∑  L(f_θ(S_1), R).




Here, S_1 are the input Sentinel-1 SAR image,  R is the available 2% binary forest non-forest pixel labels we used as reference data  and f_θ is the deep learning network parameterized by the learned parameters θ.  The Loss function (L) in (1) is a binary cross-entropy loss function. 




 §.§ Inaccurate supervision


For the inaccurate supervision approach, we use a globally available coarse resolution and noisy forest/non-forest map provided by <cit.>. We use these complete albeit inaccurate reference maps as pseudo-labels to train an initial model. We used the same Unet-like architecture as the incomplete supervision to train the network. We use the prediction from the initial model to use as a pseudo-label to train a subsequent Unet-like model (Figure <ref>a).  We continue this sequence until number of pixels in the newly predicted forest/non-forest maps class changes is below 10% compared to a previous iteration.  
 



 §.§ Architecture


The Unet-like architecture consisted of an encoder block that downsamples the feature maps to its latent space and a decoder block that reconstructed the feature maps to their original sizes. The architecture consisted of 5 encoder blocks each consisting of of 2 convolution layers interleaved with batch normalization, rectified linear unit activation (ReLU), pooling layer and 50% dropout layer. Whereas the decoder block consists of transpose convolutions interleaved with batch normalization and ReLU activation layers. Residual connections connect the encoder and decoder layers at each level. Finally a segmentation head with 1×1 convolution and a Sigmoid activation provides the probability of forest (Figure <ref>). The architecture used in this manuscript is implemented using the TensorFlow deep learning library in a Python 3.10 environment and trained on the Amazon Web Services (AWS) platform using 4 Tesla T4 GPU's.


 



 §.§ Dataset


The proposed weakly supervised approaches were trained and tested using an analysis ready Sentinel-1 SAR ground range detected images in both the VV and VH polarizations <cit.> in Mato grosso state in Brazil that is characterized by a multi-layer evergreen forest structure at a nominal pixel size of 10m. The VV and VH channels were synthesized by taking the median value of 144 images acquired from January 1, 2019 to December 31, 2019.  The training area's upper left coordinates are -53.83^0 lon, -9.70^0 lat and the lower right coordinates are -53.42^0 lon, -10.04^0 lat. Whereas the testing area's upper left coordinates are -52.59^0 lon, -11.93^0 lat and the lower right coordinates are -52.23^0 lon, -12.25^0 lat. The Sentinel-1 SAR sensor acquires data in C band for both the VV and VH polarizations. The reference pixels for the incomplete WSL was extracted from the intersection between the JAXA ALOS-PALSAR2 based global four class forest/non-forest map <cit.>, Hansen global forest cover and forest cover/ forest loss datasets <cit.> and the Land Use Change Alerts <cit.> and was visually checked using Planet scope high resolution images whereas for the inaccurate WSL, we synthesized the forest cover for 2020 by subtracting the globally available Hansen forest loss dataset from 2000-2019 from the 2000 forest cover dataset without quality control <cit.> <cit.>. 




 §.§ Network training

We synthesized the training tensor by randomly selecting 11,995 Sentinel-1 SAR image patches of size 64×64×2 and a corresponding 11,995 label patches with size 64×64×1. The two channels in the input data patch correspond with the the VV and VH polarizations. To validate the model, we extracted 1990 patches from a geographically disjointed area to validate the model during training. 


 The network was trained using Adam optimization method. We initialized the parameters randomly using the improved Xavier initialization. To diversify the training data, we applied data augmentation consisting of random vertical flips and 20^0 rotations. In both WSL approaches we applied 32 filters in the first convolution block and progressively increased the filters by a factor of 2, 4, 8, 16, 32 times in each convolution blocks respectively (Fig. 1). We also applied a learning rate of 0.001 in both WSL architectures.  During training a mini-batches of 16 samples and a weight decay factor of 5× 10^-4 was used.   





§ RESULTS

  To evaluate the performance of the proposed two WSL approaches,  we applied the trained models to a test image that was extracted from an area that was geographically disjointed from the training data region. We evaluated the accuracy of the trained models by deriving precision, recall and F-1 score by using the test dataset. We also compared the proposed approaches against a Unet-like architecture with identical parameters trained on accurately annotated dense labels. In this regard, the proposed WSL approaches achieved high accuracy that was comparable to the baseline Unet-like model that was trained on dense labels (Table <ref>). This indicated the proposed approaches are able to achieve high accuracy in mapping forest covers in a tropical setting using a small number annotated data. This was also confirmed when visually comparing the predicted forest/non-forest maps with that of a high resolution Planet scope image (Figure <ref>). 






 Overall, the proposed approaches achieved high accuracies in detecting forest cover from two weakly supervised approaches using SAR images.
 
 


§ CONCLUSIONS


In this abstract, we demonstrated the use of SAR images and two weakly supervised learning approaches for the mapping forests and non-forest regions in a tropical setting. We demonstrated a masked Unet-like architecture that trained a model from sparsely annotated reference labels and iterative refinement of a pseudo-label synthesized from globally available noisy reference dataset.


In future work, we will further improve the efficacy of the method by minimizing overfitting by using a diverse set of training dataset acquired from different eco-regions. Furthermore, we will expand the prediction capability of the model from binary segmentation to a multi-class landcover segmentation problem. Finally, we'll explore different fusion approaches between multi-modal datasets to further improve the robustness of the proposed method.





IEEEbib




 Paper Title: Outlier Detection in Large Radiological Datasets using UMAP.

















Outlier Detection in Large Radiological Datasets using UMAP
     Mohammad Tariqul Islam and Jason W. FleischerCorresponding Author 

	Department of Electrical and Computer Engineering

	Princeton University

	Princeton, NJ 08544, USA 

	 

    July 2024
===================================================================================================================================================================================



§ ABSTRACT 
The success of machine learning algorithms heavily relies on the quality of samples and the accuracy of their corresponding labels. However, building and maintaining large, high-quality datasets is an enormous task. This is especially true for biomedical data and for meta-sets that are compiled from smaller ones, as variations in image quality, labeling, reports, and archiving can lead to errors, inconsistencies, and repeated samples. Here, we show that the uniform manifold approximation and projection (UMAP) algorithm can find these anomalies essentially by forming independent clusters that are distinct from the main (“good”) data but similar to other points with the same error type. As a representative example, we apply UMAP to discover outliers in the publicly available ChestX-ray14, CheXpert, and MURA datasets. While the results are archival and retrospective and focus on radiological images, the graph-based methods work for any data type and will prove equally beneficial for curation at the time of dataset creation. 







§ INTRODUCTION


Accepted in MICCAI-2024 Workshop on Topology- and Graph-Informed Imaging Informatics (TGI3)
A prominent reason behind the current success of machine learning-based disease detection is the availability of large medical datasets. However, for the machine learning models to be reliable, quality datasets representative of the target population need to be ensured <cit.>. 
The labels in these datasets are often generated from human annotations using automated extraction or entity detection tools. 
However, these annotations (and their structured archiving) from automated tools can have errors due to faulty perceptions, interpretations, and human errors <cit.>. 

Even if the error rate of the annotator is less than 4%, this can lead to millions of annotation errors per year <cit.>. 
This parallels the 3.3% error rate in large computer vision dataset <cit.>.
Thus, there needs to be a better way to identify such errors before they are included in a dataset.




For images, the search can be performed visually. However, examining individual images is a daunting task that requires many human hours. 
A popular alternative is neighbor embedding <cit.>, which can produce a two-dimensional (2-D) cluster plot that can be analyzed visually quickly. 
(This class of methods is also known as nonlinear dimensionality reduction as the 2-D plot preserves the pairwise similarity, i.e., graph structure, of the original high-dimensional space.)
Widely used neighbor embedding algorithms are t-distributed stochastic neighbor embedding (t-SNE) <cit.> and uniform manifold approximation and projection (UMAP) <cit.>. 
UMAP was introduced relatively recently and has become very popular, as this method draws concepts from rich algebraic and topological structures and is computationally fast.

In this paper, we design a UMAP-based visual analytic method for extracting outlier images from large x-ray datasets. 
We validate our method by analyzing three publicly available and widely used medical image datasets. 
We show that the method can successfully cluster image features and produce interpretable visualization. 
We also discover labeling errors and erroneous images that have slipped through the verification process done prior to dissemination. Codes to reproduce the results are available at <https://github.com/tariqul-islam/Outlier_Detection_UMAP>.




§ RELATED WORKS


In the literature, the term outlier is often used interchangeably with abnormality and anomaly <cit.>.
Here, we define outliers as images that do not have sufficient signal for final decision-making or do not belong in the dataset due to specification. 
Generally, outlier detection methods assume an underlying distribution and often model it as normal distribution <cit.>, i.e., a data point is an outlier if it is far away from the mean of the fitted distribution.
Fritsch et al. <cit.> used the minimum covariance determinant estimator and its extensions to find outliers (due to motion or registration issues) in neuroimaging data by analyzing principal components.
Gang et al. <cit.> used a t-SNE plot to find outliers from binary lung masks in terms of size variation and segmentation error. Fleischer and Islam <cit.> employed UMAP on chest x-rays for phenotyping COVID-19 response.




§ SYSTEM OVERVIEW



Following preprocessing (discussed in <ref>), the major parts of the framework are feature extraction and dimensionality reduction (Fig. <ref>). To extract features from these images, we employed DenseNet-121 <cit.> trained on ImageNet <cit.>, a widely used deep neural network architecture designed to efficiently propagate features from earlier layers of a network to deeper layers. 
Importantly, neural algorithms are usually robust to many variabilities in images by design, and thus can accommodate standard images and outliers on equal terms.

Medical images usually vary in resolution, have different contrast, brightness, and alignment, and often suffer from registration issues.
In our framework, the features have been extracted from the final layer (before the softmax layer) of the network, where the features are generally most discriminating. 
Since we are not using a radiologically pre-trained model, these features generally will not be able to identify individual diseases. Rather, we employ other related labels (e.g., x-ray views, study labels) to examine the datasets. After extracting the features, we employ UMAP <cit.>, to obtain a 2-D approximation of the high-dimensional features.  



§ DATASETS


We evaluate our approach on three publicly available datasets: ChestX-ray14 <cit.>, CheXpert <cit.>, and Musculoskeletal Radiographs (MURA) <cit.>. 
ChestX-ray14 contains 112,120 frontal chest x-ray images from 30,805 unique patients. Images are from posterior-anterior (PA) and anterior-posterior (AP) views.
CheXpert dataset contains 224,316 chest x-rays (PA, AP, and Lateral) from 65,240 patients. We used 223,414 JPEG formatted x-rays from the training set of the dataset. 
MURA dataset contains 40,561 musculoskeletal x-rays from 14,863 studies. 
Like CheXpert, we used 36,808 x-rays from the training set and utilized study labels, e.g., finger, wrist, hand, forearm, elbow, humerus, and shoulder, for analysis.




§ EXPERIMENTS


In this section, we start by analyzing embeddings of chest x-rays (ChestX-ray14 and CheXpert datasets) and find different types of outliers. 
Then, we expand on variations of our method by altering the neural architecture, pre-training dataset, and embedding algorithms. Finally, we discuss outliers in musculoskeletal x-rays of the MURA dataset.





 §.§ Analyzing Chest X-ray Datasets




  §.§.§ Lateral X-rays in ChestX-ray14:








Our 2-D embedding (Fig. <ref> (a)) shows two large clusters of PA and AP views, corresponding to the primary topology of the DenseNet features. 
The embedding also includes a few satellite clusters around these large ones. 
These satellite clusters around the larger ones occur because the nearest neighbor graph creates a loop (or isolated sub-graph) of common features that are distinct from the rest of the data. 
In most cases, each of the satellite clusters of x-rays is from a single patient with a unique signature.
However, if any specific image features (such as similar artifacts in multiple images) are present in x-rays of different patients, these can create satellite clusters as well. Another interesting structure in Fig. <ref> (a) is the protruding region from the AP cluster.


Representative examples of anomalous clusters are shown in Figs. <ref> (b)-(e). 
The most surprising finding is the existence of some lateral x-rays in the dataset (Fig. <ref> (b)), as this dataset is supposed to be composed of frontal chest x-rays only. 
We found 92 lateral x-rays using our method. 

The protruding region from the AP cluster marked c and d 
(in Fig. <ref> (a)), consisting of x-rays with dark borders of PA and AP views, respectively. Finally, cluster (e) shown in Fig. <ref> (e) groups 46 x-rays from patient ID 9845, and a single x-ray from patient ID 12562.



  §.§.§ Corrupted Images in CheXpert:


Figure <ref> (a) shows the 2-D embedding of CheXpert dataset. 
As before, the large PA and AP clusters form the bulk of the mapping. The lateral x-rays also form a separate large cluster. 
A few of the large satellite clusters (b-f) have been marked by red circles in Fig. <ref> (a). Four images from each of the clusters are plotted in Fig. <ref> (b)-(e). 
Figure <ref> (b) depicts images with block artifacts, e.g., from poor JPEG compression or accidental splicing. 
We found 107 such images in this cluster. 
Figure <ref> (c) depicts images that are just noise (19 images). 
Figure <ref> (d) shows images with block artifacts and dynamic range issues (53 images). 
Figure <ref> (e) shows x-rays with vertical artifacts (88 images).
Finally, Fig. <ref> (f) shows rotated images. 
This cluster is placed near the large cluster of lateral (L) x-rays. 
Thus, DenseNet considers rotated x-rays to be more similar to lateral images than upright frontal x-rays.








 §.§ Effect of Different Pre-trained Networks


In the previous section, we used DenseNet-121 trained on ImageNet data for feature extraction (Fig. <ref>). 
However, using ImageNet may introduce domain shift, since the dataset is not specific to medical images. 
To assess, we looked at embeddings from different models trained on ImageNet and chest x-rays (Fig. <ref>). 

ImageNet models (Fig. <ref> (a,b)) create similar clusters of PA, AP, and L x-rays, whereas Chest-Xray14 models (Fig. <ref> (c,d)) merge the PA and AP clusters. 
The x-ray model embeddings also contain fewer satellite clusters, and their distinctness is lost. Thus, outlier detection fails. 
Additionally, we highlight the placement of images with vertical artifacts from the cluster (e) of Fig. <ref> in the alternate models (yellow points in Fig. <ref>). 
While both ImageNet models successfully separate these images (with a few scattered around), all the chest x-ray models fail (where the outliers are inside the large cluster and scattered throughout). 

A major difference between models is that ImageNet models use softmax (n-ary) cross entropy loss (for non-overlapping labels), but chest-x-ray models use binary cross entropy loss (for overlapping labels). 
To resolve this discrepancy, we trained chest x-ray models using non-overlapping labels of the ChestX-ray-14 dataset (for details, see Appendix <ref>). 
The resulting embeddings (Fig. <ref>) show a similar characteristic of the chest x-ray models with overlapping labels (Fig. <ref> (c,d)), in that the individual views are weakly separated while the specific outlier images fail to form the satellite clusters. 
For example, the outliers with vertical artifacts scatter within the 2D mapping. The networks consider them as any other x-ray image and ignore the artifacts. 
This result strengthens the idea that training on ImageNet (or a more general computer vision task) that has broader exposure benefits the discovery of outlier images. 



 §.§ Comparing Various Embedding Algorithms


To assess the effectiveness of UMAP, we compared it with a few other dimensionality reduction techniques, specifically principal component analysis (PCA), t-SNE, TriMap <cit.>, and PaCMAP <cit.>.

The PCA embedding in Fig. <ref> (a) shows the top two directions of largest variances. There are two clusters but the AP and PA views overlap. The nonlinear dimensionality reduction algorithms - t-SNE, UMAP, and variants - discover more features and make the clusters more distinct. 

The default t-SNE is tuned to preserve the neighborhood as best as possible. This often causes the individual clusters to spread out and be less compact (Fig. <ref> (b)). This behavior is apparent in the PA, AP, and lateral x-ray clusters: the separation among them is minimal, and there is little room for the satellite clusters. 
However, t-SNE can be tuned to produce a more UMAP-like output. 
Following the findings of Linderman et al. <cit.> and Bohm et al. <cit.>, we used an exaggeration factor of 4, which means we applied four times more attractive force than the repulsive force throughout the optimization procedure. 
The standard early exaggeration factor of 12 was also applied at the start of the optimization. The resulting plot in Fig. <ref> (c) largely resolves the PA and AP clusters, but there are fewer satellite clusters than the UMAP output.

Both PaCMAP <cit.> and TriMap <cit.> aim to preserve the global structure of the data while keeping the clustering properties of UMAP. PaCMAP modifies the pairwise relation of UMAP by considering different neighborhoods at different scales (near and far), while TriMap achieves this with the triplet constraint. These methods obtain the clustering of PA, AP, and lateral views (Figs. <ref> (d,e)), but, similar to exaggerated t-SNE, the satellite clusters are largely absent.

The yellow dots show the placement of x-rays with vertical artifacts from cluster (e) of Fig. <ref> for the different embedding algorithms. In all the alternate embedding algorithms, these images are scattered within the other x-ray images. The outlier images fails to be identified and labeled separately, which demonstrates a superior performance of UMAP.



 §.§ Extracting Mislabeled X-rays from MURA




Since the MURA dataset consists of x-rays from different parts of the arm and the shoulder, there is a natural ambiguity in the labels, e.g., both wrist and hand x-rays may contain the hand of a person, and shoulder x-rays may contain part of the chest. 
In such cases, finding mislabeled x-rays by embedding all the images may be sub-optimal.
To find outliers more directly, we searched for misclassified images by explicitly using labels of the dataset. The method has two parts: 1) introduce target images with a specific label (preferably from a different dataset than the MURA one); and 2) perform neighbor embedding on the joint dataset.

For example, to look for possible chest x-rays that are falsely classified as finger x-rays, we added 100 chest x-rays from the CheXpert dataset to the 5,106 finger x-rays of MURA. 
We then applied the UMAP to the composite set (Fig. <ref> (a)). As shown in Figs. <ref> (b-d), the seeded chest x-rays acted as an attractor for mislabeled images in MURA, with x-rays labeled `finger' now appearing in the (new) chest cluster. 
Interestingly, both of these x-rays were from patient 04547 (another 3 from this patient were labeled correctly). We describe an additional experiment in Appendix <ref>.



§ CONCLUSION


Neighbor embedding algorithms can be an effective tool for summarizing datasets and identifying outlier images. The principle of the method is that the outliers are different from the main data but can have similarities among themselves. 
Thus, the different outlier types form distinct clusters in the embeddings. 
Our experiments, using a DenseNet-121 feature extractor and UMAP neighbor embedding method on the ChestX-ray14, CheXpert, and MURA datasets, distinguished different radiological views of chest x-rays, classified differences, and identified wrongly labeled or corrupted images.  We further found specific types of outliers by seeding the dataset with target images and performing neighbor embedding.


While this study performed retrospective analysis of large x-ray datasets, outlier curation can be achieved during the initial assembly of the dataset as well. 
For suspected outliers, the method of seeding data with known labels can be applied. To streamline the process, appropriate reference datasets may be created beforehand. Undoubtedly, cleaner input data will result in cleaner output data. For larger datasets, more accurate results and faster embedding may be achieved by dividing them into smaller subsets and applying better alignment techniques <cit.>. Finally, since the methods are graph-based and agnostic to the underlying data type, all of the methods here can be applied to arbitrary datasets, including and especially those that are mixed modality.




  §.§.§ Acknowledgement
 The authors gratefully acknowledge financial support from the Schmidt DataX Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.



  §.§.§ Disclosure

The authors have no competing interests to declare that are relevant to the content of this article.










 unsrt
 



§ IMPLEMENTATION DETAILS




 §.§ Image Pre-processing

We apply the following transformations: histogram equalization <cit.>, resizing the images, center cropping, and normalization. Images are resized such that the lowest dimension contains 256 pixels and then center-cropped to a 224×224 dimensional image for feature extraction. 
Then the images are normalized according to the specification of ImageNet: mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225) of red, green, and blue channels, respectively.



 §.§ Feature Extraction

We use a DenseNet-121 architecture <cit.> pre-trained on the ImageNet dataset <cit.> from the PyTorch deep learning library <cit.>. (DenseNet is a deep neural network with many inter-layer connections designed to reduce the numerical instabilities that originate due to the depth of the network. The usage of neural network outputs as features is an effective baseline in machine learning algorithms <cit.> and is extensively used in medical image analysis <cit.>.)  Here, we remove the classification (softmax) layer from the neural network and use the output as features. 



 §.§ Parameter Settings of Embeddings




In general, we kept the number of nearest neighbors k to be low, as increasing k increases the computational budget exponentially.

ChestX-ray14 - Fig. <ref>: For ChestX-ray14 we used k=50. The minimum distance parameter m_d was set to 0.1. We experimented with smaller m_d values to increase the separation of the AP and PA x-rays, but it had little effect. The embedding was optimized for 200 epochs.

CheXpert - Figs. <ref>, <ref>, and <ref>: The embedding was obtained by using k=10. We used a smaller minimum distance m_d=0.001, as we found that this value provided a better separability of the large clusters. We ran the optimization for 300 epochs.
Figure <ref>: (b) t-SNE: early exaggerated for 250 steps, then 500 regular steps, (c) t-SNE: early exaggerated for 250 steps, then 500 exaggerated steps, (d) TriMap: 1000 steps with k=10, compactness parameter =0.1, and (e) PaCMAP: 450 steps with k=10.

MURA - Figs. <ref>, <ref>, and <ref>: Similar to CheXpert, we used k=10, m_d=0.001, and ran the optimization for 300 epochs.



§ PRE-TRAINING MODELS USING CHESTX-RAY14 DATASET


We trained several chest x-ray models using the ChestX-ray14 dataset. We trained traditional x-ray models using binary cross entropy loss for every disease (one vs all model) by following <cit.> which matches with the settings of the ImageNet pre-training. Additionally, to test whether the loss function and multilabel training have major effects on the resulting embeddings, we constructed a dataset by removing overlapping labels (Table <ref>) and trained two models using binary and softmax cross-entropy loss, respectively. For each mode, we optimized for 400,000 iterations using Adam (learning rate =0.001) with a batch size of 30 (2 images per disease label).





 §.§ Comparison to Publicly Available Chest X-ray Models



We compare our DenseNet-121 model to that from TorchXRayVision library <cit.>. It differs from ours in that the image pixels are normalized to [-1024,1024], images are converted to (or read as) grayscale, and `No Finding' label is untrained. 
This model (Fig. <ref>) agrees with our results and ignores the artifacts of the outlier images, showing embeddings are indifferent to alternate preprocessing scheme.



§ ADDITIONAL DISCUSSION ON MURA






Fig. <ref> shows the embedding of 36,808 musculoskeletal radiographs from the MURA dataset. There is decent separation among the x-rays in terms of the labels, but there is also a considerable overlap. For example, finger, wrist, and hand x-rays overlap, as finger x-rays include parts of the wrist and hand, and vice versa. Similarly, wrist and forearm clusters are often merged, since x-ray of the forearm tend to capture a portion of the wrist, and vice versa. The same happens for humerus and shoulder. The images within each cluster thus share similar acquisition views, aspect ratios, and specific features (e.g., circular window function, stitching of multiple x-rays in one image). 
Unlike ChexPert case, analysis of this mapping did not reveal any satellite clusters with corrupted images.

Based on this, we decided to focus on individual labels and extract images with specific labels. The chest x-rays from the CheXpert dataset used to produce Fig. <ref> are shown in Fig. <ref>.



§ EXTRACTING MISLABELED X-RAYS FROM MURA





In a different experiment on the MURA dataset, we used `finger' and `shoulder' x-rays (Fig. <ref>). 
The broad features of the finger and shoulder are easily separable with a few misclassified points. 
Analyzing `finger' x-rays misclassified in `shoulder' clusters, we can find the two chest x-rays labeled as finger (which we found in previous experiment as well) and two images that are just noise/non-x-ray images (Fig. <ref> (b)). 
The latter belongs to patient ID 04687, which uncovers two more outliers (one is an x-ray of keys). The misclassified `shoulder' labels in the `finger' cluster reveal three leg x-rays (Fig. <ref> (c)).



 Paper Title: Securing the Diagnosis of Medical Imaging: An In-depth Analysis of AI-Resistant Attacks.



Article Title]Securing the Diagnosis of Medical Imaging: An In-depth
Analysis of AI-Resistant Attacks













1]Angona Biswasangonabiswas28@gmail.com

[2]MD Abdullah Al Nasimnasim.abdullah@ieee.org
These authors contributed equally to this work.

3]Kishor Datta Guptakgupta@cau.edu
These authors contributed equally to this work.

4]Roy George george@cau.edu
These authors contributed equally to this work.

5]Abdur Rashid A.Rashid.653@westcliff.edu
These authors contributed equally to this work.

[1]Research and Development Department, Pioneer Alpha,  Dhaka, Bangladesh

[2]Research and Development Department, Pioneer Alpha,Dhaka,Bangladesh

[3, 4]Department of Computer and Information Science, Clark Atlanta University, Atlanta, GA, USA

[5]Master’s of Engineering Management (MSEM), Westcliff University, Irvine, California, USA







Machine learning (ML) is a rapidly developing area of medicine that uses significant resources to apply computer science and statistics to medical issues. ML's proponents laud its capacity to handle vast, complicated, and erratic medical data. It's common knowledge that attackers might cause misclassification by deliberately creating inputs for machine learning classifiers. Research on adversarial examples has been extensively conducted in the field of computer vision applications. Healthcare systems are thought to be highly difficult because of the security and life-or-death considerations they include, and performance accuracy is very important. Recent arguments have suggested that adversarial attacks could be made against medical image analysis (MedIA) technologies because of the accompanying technology infrastructure and powerful financial incentives. Since the diagnosis will be the basis for important decisions, it is essential to assess how strong medical DNN tasks are against adversarial attacks. Simple adversarial attacks have been taken into account in several earlier studies. However, DNNs are susceptible to more risky and realistic attacks. The present paper covers recent proposed adversarial attack strategies against DNNs for medical imaging as well as countermeasures. In this study, we review current techniques for adversarial imaging attacks, detections. It also encompasses various facets of these techniques and offers suggestions for the robustness of neural networks to be improved in the future.



















[
    [
    August 5, 2024
==================





§ INTRODUCTION




 §.§ Adversarial Attack


Adversarial attacks encompass a set of techniques aimed at manipulating machine learning models by introducing well-crafted, often imperceptible alterations to input data. These modifications aim to deceive the model, leading to misclassifications or inaccurate predictions <cit.>. Adversarial attacks have the potential to compromise machine learning systems' dependability and security, which is concerning for crucial applications such as cybersecurity, autonomous vehicles, and medical diagnosis. Of all the adversarial attack types, "White-box attacks" and "Black-box attacks" are the most common.

Ongoing research in machine learning security centers on the analysis of adversarial attacks and the development of defense strategies. Researchers continually explore novel attack tactics to identify potential weaknesses in AI systems while concurrently striving to construct robust models less susceptible to adversarial perturbations. Artificial intelligence (AI) modern technology is widely successful in fields like computer vision, natural language processing, and automated driving; however, its application in vital safety sectors is hindered by its susceptibility to adversarial attacks. Consequently, enhancing the resilience of AI systems against such attacks has become paramount for the progress of AI <cit.>.

The rapid advancement of AI technologies has found diverse applications in numerous domains, from machine translation, speech recognition, and object identification to more intricate tasks like drug composition analysis <cit.>. Noteworthy applications include brain circuit construction <cit.>, particle accelerator data analysis <cit.>, <cit.>, and DNA mutation impact analysis <cit.>. Since Szegedy et al.'s seminal work <cit.> highlighting neural networks' vulnerability to adversarial attacks, research on adversarial technologies for artificial intelligence has steadily expanded. New techniques for adversarial attacks and strategies for mitigating them are consistently emerging.

The term "adversarial attacks at the training stage" refers to actions taken by adversaries to manipulate the training dataset, input characteristics, or data labels during the training phase of the target model. This manipulation of the training dataset involves actions such as adding or deleting training data, as demonstrated by Barreno et al.'s approach <cit.>. During the testing phase, adversarial assaults are categorized into white-box attacks and black-box attacks <cit.>. In white-box situations, attackers have access to the parameters, techniques, and structure of the target model, enabling them to craft adversarial samples based on this information.



 §.§ The Adversary's Objective: Evasion attack versus poisoning assault


"Poisoning attacks" are assault methods that allow an attacker to add or change a large number of fake samples to the training set of a DNN algorithm. The trained classifier may perform poorly as a result of these bogus data. They may have poor accuracy <cit.> conversely make imprecise predictions on a few samples for analysis <cit.>.
The classifiers used in evasion attacks are fixed and often work well on safe testing samples. The adversaries are unable to reform the classifier's settings or parameters, but they do generate some fictitious instances that the classifier is unable to distinguish.



 §.§ Deep Neural Networks: Adversarial Attacks

Regarding medical image processing use cases such as diagnosing cancer and lesion identification, deep neural networks, also known as DNNs, are rapidly gaining popularity <cit.>. Nonetheless, a new study suggests that adversarial examples/attacks with tiny, barely noticeable changes can weaken medical deep learning systems. There are presently safety concerns associated with the usage of these devices in clinical settings <cit.>.  Deep neural networks (DNN) are increasingly well-liked and effective at a variety of machine learning requisitions. They have been employed with surprising effectiveness in a variety of recognition issues in the fields of pictures, graphs, text, and voice <cit.>. They are able to identify things in images with accuracy that is almost human <cit.>, <cit.>. Additionally, they are employed in speech recognition <cit.>, natural language processing <cit.>, and gaming <cit.>. 

   
Regarding the data set provided by MNIST, Biggio et al.<cit.> produce adversarial instances first, focusing on traditional machine learning classifiers like SVMs and 3-layer neural networks with full connectivity which is shown in Figure <ref>. In order to trick the classifier, it optimizes the discriminant function. As an illustration, consider a linear SVM classifier on the MNIST dataset. 



 §.§ Examining examples of opposition in the real world


The research <cit.> put decals to traffic signs that pose a serious threat to autonomous vehicles' sign recognition technology. These hostile objects are especially detrimental to deep learning systems since they can interfere directly with numerous real-world DNN usages, such facial recognition and autonomous cars. By assessing the adversarial images (FGSM, BIM) produced to see if they are "robust" against changes in natural circumstances (e.g., shifting viewpoint, illumination, etc.), the authors of the work <cit.> investigate the viability of creating tangible adversarial objects. Robust in this context means that even after transformation, the produced pictures are still antagonistic. The experiment's findings show that many of these adversarial examples—particularly those produced by FGSM—remain antagonistic to the classifier even after transformation. The findings imply that antagonistic real-world objects could trick the sensor in many situations.



 §.§ Medical Image Under Artificial Intelligence Attack


Researchers have access to strong models of developing science and technology thanks to deep learning. Convolutional neural networks, also known as CNNs, are among the most significant categories of deep learning (DL) frameworks for the analysis and processing of images because of their exceptional ability to learn valuable information.  The processing of the human organism utilizing different picture modalities for therapeutic, diagnostic, and health surveillance purposes is known as medical image analysis <cit.>. The advancement of computer vision through the use of deep neural networks addresses issues that were not well-solved by traditional image processing methods. Beinfeld et al. <cit.> asserted that $385 spent on medical imaging results in a savings of almost $3000. MRI, CT scans, ultrasound (US), and X-rays are the most frequently used image modalities.  Figure <ref> shows how adversarial assaults may be used to arbitrary alter diagnosis outcomes in three different medical picture datasets: Dermoscopy <cit.>, X-ray of the chest <cit.>, and Fundoscopy <cit.>. 

   


Hu et al. <cit.> assert that in order to develop more reliable models, it is crucial to investigate the reasons why adversarial cases emerge and to comprehend deep learning models better. Depending on the knowledge of the enemy, attacks can be categorized into three kinds.   The saliency (or attention) map of a picture that was input shows the regions that significantly change the model's output based on the gradients of the categorization loss with regard to the input <cit.>. We can see that some medical photos have highly concentrated regions that are noticeably larger. This could indicate that the rich biological textures included in medical photos occasionally entice the DNN model to focus more on aspects unrelated to the diagnosis. Tiny adjustments in these high-focus areas can significantly affect the model's output. 

Due to several factors, including High-Dimensional Data, High-Dimensional Data, Transferability of Adversarial Examples, Complex Decision Boundaries, Black-Box Attacks, Safety-Critical Applications, and Safety-Critical Applications, Medical Image Deep Neural Network (DNN) models can be relatively easier to attack when compared to some other domains. Despite these difficulties, experts in the field are actively attempting to create more reliable and safe DNN models for medical images. To make these models more resistant to hostile attacks, strategies like adversarial training, input preprocessing, and defensive distillation are being investigated. To further enhance the security and dependability of AI systems in healthcare applications, consistent evaluation frameworks for adversarial robustness in medical imaging must be developed.

Byra et al. <cit.> proposed an attack strategy on ultrasonography (US) imaging for liver fatty tissue. Radio-frequency signals are used to rebuild US pictures, and the reconstruction technique was subjected to a zeroth-order optimization attack <cit.>. The InceptionResNetV2 model was used in the studies, and the assault resulted in a 48% loss in model accuracy. Ozbulak et al. <cit.> proposed an attack specifically designed for medical picture segmentation called the adaptive segmentation masking attack (ASMA). The suggested attack offers significant intersection-over-union (IoU) degradation and produces nearly undetectable samples for most portions. Because the U-Net framework is among the most well-known models for clinical picture segmentation, they employed it in the trials. The datasets used were for segmentation of glaucoma optic disk <cit.> and ISIC skin lesion segmentation <cit.>.  Chen et al. published a method for fabricating hostile cases to thwart medical picture segmentation <cit.>. In order to simulate anatomical and intensity fluctuations, geometrical deformations are used to create the adversarial examples. By attempting to partition organs from abdominal CT scans using a U-Net model, they tested the effectiveness of these examples. With reference to the Dice score measure, they successfully reduced it significantly across all organs. The kidneys and pancreas, however, require a higher level of disturbance and are more arduous to assault than the liver as well as spleen.

   


Tian et al. <cit.> examined the phenomena of bias field, which can arise from incorrectly capturing a medical image and jeopardize a DNN's efficacy, as shown in Figure <ref>. Motivated by adversarial attacks, the authors created an adversarial-smooth bias discipline approach to trick a model. The ResNet50, MobileNet, among others and DenseNet121 models were employed to fine-tune the chest X-ray dataset that was used in their study. They looked at this attack's transferability and white-box attacks. Comparing the suggested attack to other cutting-edge white-box attacks, it exhibited a greater attack accuracy on transferability.

This review's major organization is as follows: We offer some significant adversarial notions on medical images in Section 2. Additionally, it provides a detailed overview of earlier research. We explore a few strategies in Section 3 with regard to the picture classification scenario. We utilize Section 4 to quickly summarize some findings from earlier research that tries to explain the phenomena of antagonistic examples on medical data. The overview is concluded in Section 5.



§  BACKGROUND LITERATURE REVIEW


Adversarial instances are intentionally created data inputs that are intended to degrade a machine learning model's performance. When researchers examined strategies used by spammers to evade spam filters in 2004, they unofficially coined the term "adversarial inputs" <cit.>. These misleading examples are usually produced by deliberately faking real data, such spam advertising messages, in order to trick the computer system that analyzes it. Alterations can be applied to text data, like spam, by introducing innocuous text or modifying frequently used phrases in malicious communications with synonyms. To fortify algorithms against adversarial attacks, researchers have explored various strategies, including training algorithms on adversarial samples and employing sophisticated data processing techniques to minimize the susceptibility to manipulation. The quest for fully robust models in machine learning aims to accelerate the development of algorithms capable of making decisions based on consistent explanations, with promising early efforts in this direction <cit.>.

In the realm of healthcare, medical claims codes play a crucial role in determining the amount spent on a patient visit post-payer approval. Payers often assess these claims using automated fraud detectors increasingly driven by machine learning. Historically, healthcare providers have shaped payer records of patient visits, including the associated codes, to influence the algorithmic outputs of payers <cit.>. Medical fraud, a market valued at $250 billion, exemplifies the extreme end of this strategic tailoring of patient presentations. While some practitioners may overtly fabricate medical claims, patient data falsification often takes more covert forms. For instance, consistently submitting codes for billing services identical to, but more expensive than, those actually provided is known as intentional upcoding. Here, Table <ref> shows past exclusive past research works where authors did some fantastic work for medical images to defend adversarial attacks. Different researcher proposed different frameworks to overcome this problem and to suggest a new thought for revolution. 



 §.§ Elements of a Cyber Adversary's Assault


   


In some instances, subtle adjustments to billing codes blur the boundary between fraudulent activities and conscientious best practices. One notable illustration may be found in the guidelines on the Endocrine Society website, which advise medical professionals not to bill for metabolic syndrome, or International Classification of Diseases (ICD) code 277.77, in situations of obesity. Coverage denial is likely to occur with this particular code/condition combo in Figure <ref>. The Society recommends an alternative approach, suggesting billing for the individual codes related to specific disorders comprising the metabolic syndrome, such as hypertension. 










 §.§ The components of an adversarial attack


   


Presenting a thorough overview of the various adversarial attack strategies and defense techniques is the aim of the study <cit.>. We first discuss the theoretical underpinnings, practices, and practical uses of adversarial attack strategies. Then, a brief discussion of the research on defense strategies obscuring the boundary of the large field follows. A few selected articles are published year-by-year from 2012 to 2021, as seen in Figure <ref>. This work aims to provide thorough taxonomies, evaluations of harmful assaults, and protections against the full DL pipeline. In this context, the numerous attack and defense tactics that have been developed over the previous two years have been categorized, with a focus on the clinical deep learning systems that are susceptible to adversarial attacks. 

Several defensive strategies have been put forth in an effort to neutralize possible threats. A popular method in natural imaging called adversarial training adds adversarial images to the training dataset in order to make Convolutional neural network, or CNN, models more resilient. Nevertheless, this method is not the best for datasets pertaining to medical imaging, since the addition of different adversarial images could considerably reduce classification accuracy. A robust detection technique for malicious images is presented in the work of Li et al. <cit.>, successfully defending against attacks on deep learning-driven medical image labeling systems. In a different study <cit.>, scientists compared CNN performance with Vision Transformers (ViTs) to determine how resilient CNNs are to different types of attacks in computational pathology. The authors developed robust neural network models, evaluating their efficacy against both white- and black-box attacks. The structures of attacks for both models were scrutinized, with an exploration of the underlying factors influencing their performance. The study's findings were validated through two clinically relevant classification tasks involving distinct patient groups. Preceding Ma et al.'s research <cit.>, ambiguity surrounded medical image adversarial attacks (AAs), limiting adversarial machine learning analyses to natural images. Unlike natural images, medical images may contain domain-specific elements, impacting medical deep-learning systems susceptible to AAs. AAs possess the potential to manipulate diagnoses and outcomes, emphasizing the critical need for a robust healthcare infrastructure to mitigate potential risks from malicious attacks.






§ MATERIALS AND METHODS



 §.§ Medical Adversarial Defensive Strategies

Various defense models, including techniques such as input denoising, input gradients regularization, and adversarial training, have been devised. However, recent attacks often manage to evade or partially circumvent these protective measures. This paper undertakes a comprehensive exploration of adversarial attack and defense methods within the domain of medical image analysis. It introduces a novel taxonomy based on the application context for a family of techniques, as outlined in Dong et al.'s work <cit.>. The author establishes a unified theoretical framework encompassing diverse adversarial attack and defense strategies tailored for medical imaging applications.

The primary emphasis in most research efforts on adversarial attacks in medical image processing is on the white-box scenario. These studies, characterized by a comprehensive understanding of medical Deep Neural Networks (DNNs), concentrate on the susceptibility of computer-aided diagnosis models across various medical imaging applications. In the execution of adversarial attacks, attackers may utilize the target diagnosis DNN as a locally deployed model. In addition to white-box adversarial attacks for medical classification tasks, scholars have delved into vulnerabilities associated with various medical imaging tasks. Notably, these publications focus predominantly on medical segmentation, as evident in works such as Chen et al.'s <cit.> and Ozbulak et al.'s <cit.>.

For natural photos, semi-white-box (Gray-box) attacks have been extensively researched <cit.>. However, there are only a few academic publications <cit.>, <cit.> that address this attack scenario for medical image processing. The semi-white-box adversarial approach typically consists of two stages: 1) To create adversarial instances against the target DNN model, the attacker trains a generative model. The attacker has complete access to the target model during training, including backward propagation gradients. 2) Instead of needing to know anything about the target model, as would be the case in a completely black-box scenario, the adversary generator can directly obtain adversarial examples against the target model with the input of authentic photos during the application stage.

Presently utilized white-box adversarial attacks often require multiple backward gradients of the target model. In simpler terms, attackers generate comparable adversarial samples by treating the target Deep Neural Network (DNN) as if it were a locally deployed model. However, due to its reliance on an in-depth understanding of the DNN model to execute an attack, the white-box approach may be unreliable in real-world scenarios. Conversely, the general black-box scenario may offer a more suitable environment for simulating real-world adversarial attacks, with numerous proposals focusing on investigating black-box assaults for natural images.

The outlined solutions usually involve multiple queries to the black-box model or depend on having a comprehensive understanding of the desired diagnosis model. However, in many real-world situations, attackers may lack direct access to the target diagnosis model. The restricted black-box (no-box) configuration, specifically representing the most challenging scenario for real-world adversarial assaults, may be more covertly dangerous even without querying the target black-box DNN. The transferability <cit.> of adversarial instances across different DNN models fundamentally determines the success of no-box attacks. For example, a no-box attacker can generate adversarial images based on a locally deployed surrogate model, transferring them directly to target medical diagnosis systems. Restricted black-box adversarial attacks are deemed more subtly perilous for natural vision tasks according to a renowned study.

Considering the substantial risks to the healthcare sector, numerous defense strategies have been proposed to counter medical adversarial attacks. In order to overcome the weaknesses shown by adversarial examples, a number of methods have been developed to provide trustworthy deep learning-based systems for natural pictures <cit.>. The development of accurate computer-aided diagnosis models for clinical applications not only contributes to trustworthy healthcare services for millions but also underscores the importance of investigating adversarially robust models within the realm of medical image analysis.

In their work <cit.>, the authors succinctly summarize studies on adversarial defense in the context of medical image analysis. Various attack methods in the realm of medical images serve as robustness evaluation criteria for adversarial defense, complementing the extensive range of adversarial attack methods mentioned. The minimal separation between adversarial examples and the model's decision boundary signifies the model's resistance to disturbances. The goal of adversarial training is to mitigate subsequent losses. 

    
Minimize the 


    L(f(x_adv, y) + λ· L(x, x_adv))


where L(f(x_adv, y) is the classification loss, L(x, x_adv) is the perturbation size, and λ is a trade-off hyperparameter. 



Transforming the input data into an alternative feature space that diminishes the impact of adversarial perturbations on the model. A technique aimed at enhancing resilience against adversarial attacks involves introducing noise to the features or intermediate representations that characterize the input data.




 §.§ Adversarial Training

Most medical adversarial defense techniques focus on using adversarial training to create reliable diagnoses systems. A significant fraction of the works among them go beyond current techniques for natural image adversarial training to tasks relating to medical classification <cit.>, <cit.>, Vatian et al. <cit.> looked into opposing examples for medical imaging and tried a number of defense strategies to oppose these nefarious representations.



 §.§ Adversarial Detection

Adversary detection seeks to identify adversary cases from input examples during the application stage, as opposed to developing robustness during the training stage of computer-aided diagnosis models.Many adversarial detection techniques have been put forth in the field of medical image analysis to stop additional misdiagnosis caused by adversarial samples <cit.>. In particular, it is possible to think about medical adversarial detection as an anomaly detection issue that can be resolved by combining explainability approaches <cit.>.



 §.§ Image-level Pre-processing

A clean image and the related adversarial perturbation make up an adversarial image in general. Meanwhile, it has been shown that DNNs can perform well on clean images while still being vulnerable to adversarial examples <cit.>. Denoising the adversarial example to remove the perturbation component can therefore help make the subsequent network diagnostic easier. Image-level pre-processing can be useful and secure in the context of biomedical image analysis because it does not require re-training or modifying medical models. 



 §.§ Improvement of Features


The difference in robustness between human and machine vision is attributed to adversarial examples linked to non-robust features extracted from specific patterns in the data distribution <cit.>. Consequently, enhancing feature representation is crucial for the development of robust inference systems. In this investigation, we characterize feature augmentation as the modification of architectures or mapping functions. Various techniques aimed at improving features have significantly enhanced the robustness of medical classification models <cit.>.



 §.§ Distillation of Knowledge

In the field of machine learning, knowledge distillation is a helpful technique for transferring learned information from a complex (teacher) model to a simple (student) model. Therefore, the situation where the network structures for the teacher and learners are identical is specifically referred to as self-distillation. Furthermore, a great deal of research has been done on the natural imaging domain using adversarial knowledge distillation, which moves the adversarial resilience from a heavy teacher model to a view student model <cit.>.



§ RESULT AND DISCUSSION


This study utilizes four publicly accessible benchmark datasets, as detailed in <cit.>, to investigate adversarial attack and defense in the context of medical image processing (illustrated in Figure <ref>). Firstly, the Messidor1 dataset comprises 1,200 eye fundus color images for detecting diabetic retinopathy across four classes based on retinopathy grade. Secondly, the International Skin Imaging Collaboration (ISIC) dataset consists of 2,750 dermoscopic images categorized into three classes for skin lesion classification and segmentation. Thirdly, the ChestX-ray 14 dataset comprises 112,120 frontal-view X-ray images representing 14 thorax diseases. Lastly, the COVID-19 database incorporates 21,165 chest X-ray images accompanied by segmentation-capable lung masks.

   


The study conducted by the authors <cit.> employs experiments to demonstrate the following: 1) The SSAT module significantly enhances the adversarial robustness of the model without compromising the classification accuracy of clean images. 2) The UAD module effectively identifies and excludes a majority of successful adversarial examples. 3) In comparison to other existing AI systems, their medical imaging AI solution (UAD + SSAT) minimizes the risk of adversarial attacks. The research utilizes a publicly available dataset of retinal OCT images, employing the "white-box" setting as the most challenging threat for assessing class prediction performance. Across all white-box attack conditions, the authors establish that SSAT consistently outperforms alternative baselines while maintaining comparable or superior performance for clean image classification. Furthermore, the authors demonstrate that the combined approach of UAD complementing SSAT results in the lowest adversarial risk according to the new measure presented. Regardless of the training techniques employed, systems based on UAD consistently exhibit reduced risks compared to those that do not incorporate this defense mechanism.

   
An AI-CAD model depicted in Figure <ref> was initially trained to manipulate diagnosis-sensitive elements within images, such as adding or removing malignant tissue. The model's effectiveness was then assessed using adversarial images produced by a GAN model. To determine how well human specialists could visually distinguish images produced by the GAN, a reader study was carried out.

In a study by Zhou et al. <cit.>, the responses of an AI-CAD model to adversarial attacks on GAN-generated mammography images were investigated. This involved introducing malignant tissue into healthy images and removing cancerous tissue from images affected by cancer. The study also evaluated the proficiency of experienced radiologists in visually identifying such adversarial images, both before and after instructional intervention. The study cohort, provided by the University of Pittsburgh Medical Center, included 1284 women, and 4346 mammography images were collected. Among these, 366 patients had biopsy-proven malignant breast cancer, while 918 patients had breast cancer assessed as benign (including benign signs).

Following training, the researchers assessed the model's classification accuracy concerning both the original genuine test data and its equivalent GAN-generated adversarial counterparts. Two GAN-trained U-Net23 models were employed to generate adversarial images, flipping labels in the test set. The high-resolution (1728 x 1408) photographs indicated the categorization impact of the AI-CAD model on the test data. The AI-CAD model achieved an AUC of 0.82 on the test set, consisting of 364 genuine negative samples and 74 real positive samples. The adversarial GAN-generated images (with flipped labels) in the test set yielded an AUC of 0.94. These AUC values suggested that the adversarial collection of images effectively deceived the AI-CAD model. Additionally, 59.5% (44 out of 74 cases) of genuine positive images with a classification accuracy threshold of 0.5 were correctly labeled as positive, whereas 95.5% (42 out of 44 cases) of adversarial samples created by GANs successfully tricked the system by being classified as negative. 

   


Figure <ref> displays the mean SSIM values for FGSM and PGD assaults across all photos. The Supplementary Material includes the SSIM data for every model. As can be shown, the effects of the identical disturbance applied to various imaging modalities on human visual perceptibility with the observed SSIM varies. The radiology images most clearly showed adversarial disturbances, with = 0.02 producing an already apparent, but very modest perturbation. At the same perturbation level for the ophthalmology and pathology photos, perturbations were nearly undetectable and became apparent with larger epsilon values. The authors examine black-box adversarial attacks on deep learning in healthcare imaging. <cit.>. They investigate three medical imaging sectors where deep learning algorithms are vulnerable. In all three datasets, perturbations calculated by FGSM showed lower SSIM than those calculated using PGD. This is a predicted outcome given that PGD optimizes perturbations based on their size and influence on model predictions.We opted to document attacks in our subsequent analyses employing 0.02, as this represented the maximum perturbation level that remained effective across all applications and attack strategies. Additionally, it exhibited greater transferability than an epsilon value of 0.01 in the majority of the examined applications.

   


Figure <ref> shows examples taken from the ophthalmology collection that demonstrate the transferability of assaults when the target and surrogate are randomly started and when both are pre-trained on ImageNet. Whether the target and surrogate models had the same architecture or a different one did not affect the findings. 



§ CONCLUSION

Diagnostic imaging AI systems based on computer vision are increasingly being used as models for classifying and segmenting diseases. The scaled-up use of medical imaging AI systems has given rise to serious safety issues because to DNNs' susceptibility to adversarial samples. Recently, a number of ways have been put out to increase the efficacy of medical image defense tactics. Although numerous protection strategies have been put forth, there are still reservations over the use of medical deep learning techniques. This outcome arises from certain constraints within medical imaging, including a scarcity of high-quality image datasets and labeled data in comparison to more abundant datasets available for natural images. The efficacy of adversarial defenses may not be universally applicable in this context. The findings underscore the importance of exercising caution in the design of Deep Neural Networks (DNNs) for medical imaging and their subsequent real-world applications.


























§ DECLARATIONS






 §.§ Funding
 A portion of the funding for this research comes from DOEd Grant P116Z220008 (1) and NSF Grant No. 2306109. The author(s) expresses all opinions, findings, and conclusions; the sponsor(s) does not necessarily agree with them.


 §.§ Conflict of interest
 The authors do not declare any conflicts of interest.



 §.§ Ethics approval
 There is no original research involving human subjects, animal subjects, or sensitive data in this review publication. 



 §.§ Consent to participate
 The studies and data included in this review were sourced from publicly available and previously published literature. 


 §.§ Consent for publication
 All authors listed on this review paper have reviewed and approved the final manuscript for submission to Machine Learning. Each author has contributed significantly to the research, writing, and revision of the paper. All authors have read and agreed to the content presented in this review, and their consent for publication is hereby provided.


 §.§ Availability of data and materials
 The data and materials used can be found in the references of this work.


 §.§ Code availability
 The review is based on a comprehensive analysis of publicly available literature, and any references to specific methodologies or algorithms are attributed to the respective original sources. We have strived to provide clear citations and references to enable readers to locate the original works for further examination.


 §.§ Authors' contributions

 Conceptualization: A.B., M.A.A.N.; Methodology: A.B., M.A.A.N.; Formal analysis and investigation: A.B., M.A.A.N.; Writing - original draft preparation: A.B., M.A.A.N., R.G., A.R.; Writing - review and editing: M.A.A.N., A.R.; Funding acquisition: K.D.G; Resources: R.G.; Supervision: K.D.G., R.G. All authors have read and agreed to the published version of the manuscript.

















 Paper Title: MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework.




MIST



A. Celaya et al.


The University of Texas MD Anderson Cancer Center, Houston TX 77030 Rice University, Houston TX 77005


MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework
    
Adrian Celaya1,2 Evan Lim1 Rachel Glenn1 Brayden Mi1 Alex Balsells1,2 Tucker Netherton1 Caroline Chung1 Beatrice Riviere2 David Fuentes1

    August 5, 2024
==============================================================================================================================================


§ ABSTRACT 
Medical imaging segmentation is a highly active area of research, with deep learning-based methods achieving state-of-the-art results in several benchmarks. However, the lack of standardized tools for training, testing, and evaluating new methods makes the comparison of methods difficult. To address this, we introduce the Medical Imaging Segmentation Toolkit (MIST), a simple, modular, and end-to-end medical imaging segmentation framework designed to facilitate consistent training, testing, and evaluation of deep learning-based medical imaging segmentation methods. MIST standardizes data analysis, preprocessing, and evaluation pipelines, accommodating multiple architectures and loss functions. This standardization ensures reproducible and fair comparisons across different methods. We detail MIST's data format requirements, pipelines, and auxiliary features and demonstrate its efficacy using the BraTS Adult Glioma Post-Treatment Challenge dataset. Our results highlight MIST's ability to produce accurate segmentation masks and its scalability across multiple GPUs, showcasing its potential as a powerful tool for future medical imaging research and development.





§ INTRODUCTION

Medical imaging segmentation is a highly active area of research. Since its introduction in 2015, the U-Net architecture has received nearly 90,000 citations on Google Scholar <cit.>. Other architectures and frameworks like nnUNet have achieved state-of-the-art accuracy in several medical imaging benchmarks like the Brain and Tumor Segmentation (BraTS) and Medical Segmentation Decathlon (MSD) <cit.>. Since nnUNet's introduction in 2018, several innovative approaches for medical imaging segmentation have emerged. These innovations include new architectures like vision transformers <cit.> and loss functions like the boundary and generalized surface losses <cit.>. Despite these recent advances, there remains a lack of standardized tools for testing and evaluating different approaches for medical imaging segmentation. For example, several works like <cit.> report superior performance to nnUNet while <cit.> refutes these results. This inconsistency makes it difficult to evaluate and assess claims of state-of-the-art performance for new research in deep learning-based medical imaging segmentation.


To address this inconsistency, we propose the Medical Imaging Segmentation Toolkit or MIST, a simple, modular, and end-to-end deep learning-based medical imaging segmentation framework that allows researchers to seamlessly train, test, and evaluate existing and new methods on their data in a consistent and reproducible way. MIST has a rule-based data analysis and preprocessing pipeline that it makes available to multiple architectures, loss functions, and training parameters. Additionally, MIST has a standardized and flexible evaluation pipeline that computes multiple metrics like the Dice, Haudorff distance, and surface Dice for any user-defined segmentation class. This standardization of several pipelines allows for fair comparisons between methods and their effect on segmentation accuracy for any given dataset.



§ METHODS
 
In this section, we describe the required data format for MIST, the different pipelines in MIST, and the auxiliary features, such as evaluation, postprocessing, and test-time inference. We also provide the specific settings we use to train on the BraTS challenge data.



 §.§ Data Format
 
MIST uses the same file structure as the BraTS challenge datasets <cit.>. That is, MIST assumes that the train and test (optional) data directories are formatted so that each sub-directory corresponds to one patient and that each imaging modality (i.e., T1, T2, or CT) is its own NIfTI file. The labels should all be in a single NIfTI file in the same sub-directory for each patient. Once the dataset is in the correct format, the next step is to prepare a small JSON file containing the details of the dataset. These include information like the imaging modality, naming conventions for the images and masks, labels, and the final classes that we use for evaluation (i.e., labels 1, 2, and 3 for whole tumor). An example of this JSON file is provided in the documentation page https://mist-medical.readthedocs.io/en/latest/getting_started/here. MIST also provides support for converting datasets that are in CSV or MSD format. 



 §.§ Pipelines
 
MIST contains three main pipelines: Data Analysis, Preprocessing, and Training. These pipelines can all run at once with a single command or individually. We also provide several auxiliary pipelines for evaluation, postprocessing, and test-time inference.



  §.§.§ Data Analysis
 
The data analysis pipeline gathers parameters about the dataset including cropping to foreground, target spacing, patch size selection, normalization parameters, and is described further below. The results of this analysis are saved in a  file.



  
Cropping to Foreground Often, zeros or useless background information surround the region of interest. For example, zeros surround the region outside of the brain in the BraTS datasets, and the table and air surrounding the body in CT imaging do not provide relevant information. From a computational perspective, it is desirable to crop the images to the foreground to make them smaller. Additionally, cropping to the foreground also reduces the sparsity of the dataset. MIST crops each image by first windowing it with its 33rd and 99.5 percentile values, then using an Otsu filter to produce a foreground threshold. We then crop the image according to the binary mask resulting from this threshold (i.e., one if the voxel intensity is greater or equal to the threshold). To determine if we need to use cropping, we check to see if, on average, the resulting volume reduction is at least 20%.



  
Target Spacing To determine the target spacing for a dataset, we first look at the distribution of the voxel spacings along each axis. Our initial target spacing is the median spacing in each direction. However, suppose the resulting target spacing is anisotropic (i.e., the maximum and minimum spacing ratio is greater than three), then we take the 10th percentile spacing along the lowest resolution axis in that case. However, if resampling with this selected target spacing results in any image taking up more than 2 GB of memory, then we increase all components of our target spacing by 25%. We repeat this process until none of the resampled images take up more than 2 GB of memory. This last step is important because it conserves computational resources. Larger images take longer to load from disk and require significantly more resources to manipulate (i.e., patch extraction and augmentation).



  
Patch Size Selection We set the patch size for training from the median resampled image size by taking the nearest power of two less than or equal to each dimension in the median image size up to a maximum patch size. The default value for this maximum patch size is 256×256×256. However, users can adjust this maximum patch size based on their knowledge of their computing resources. 



  
Normalization Parameters If the dataset consists of MR or other non-CT images, then our windowing parameters are the 0.5 and 99.5 percentile values of either the entire image or the non-zero values of the image if, on average, the ratio of the number of non-zero to zero-valued voxels is less than 0.2. We use the mean and standard deviation with the same rules regarding the non-zero values for normalization. On the other hand, if we are dealing with CT images, then we compute the 0.5 and 99.5 percentile values, mean, and standard deviation of all of the voxels over the entire dataset corresponding to regions labeled as non-zero in the ground truth segmentation masks.




  
Other Checks In addition to the previous analysis steps, MIST checks to see that the header information in the images and masks agree and, if there are multiple modalities, that the headers in each image agree. If there is a disagreement between the headers, MIST prints a warning with the patient ID to the console and excludes the example from training. MIST also computes equal-weighted class weights for each class based on the amount of voxels belonging to each class over the entire dataset. These class weights are primarily informational but can be used for weighted loss functions. 



  §.§.§ Preprocessing
 
The preprocessing pipeline takes the information from the analysis pipeline to crop the images based on their foreground mask (if called for), reorient to right-anterior-inferior (RAI), resample to the target spacing, and window and normalize the intensity values. This pipeline can also compute the distance transform maps (DTMs) for each class in the ground truth masks or apply bias correction for MR images. Users also have the option to skip this pipeline if, for example, their NIfTI files contain preprocessed images already. Figure <ref> illustrates the workflow for this pipeline.





  
Reorient & Resample After optionally cropping or bias correcting each image, the preprocessing pipeline will reorient each image and mask to the RAI orientation. For resampling images to the target spacing, we use third-order spline interpolation. We apply one-hot encoding and use linear interpolation to each channel for masks. If the current image spacing is anisotropic, we first resample along the low-resolution axis using nearest-neighbor interpolation and then apply the appropriate interpolation along the other two axes <cit.>. We implement these operations using SimpleITK to optimize the speed and computational performance of these operations <cit.>. 



  
Window & Normalize Before normalizing the voxel intensities in each image, the preprocessing pipeline windows each image by either the image's 0.5 and 99.5 percentile or, if we are using CT images, by the precomputed windowing parameters. After windowing, MIST applies z-score normalization with either the image's mean and standard deviation or the precomputed normalization parameters for CT images. Again, in the case that the ratio of non-zero to zero-valued voxels is less than 0.2, then the mean and standard deviation are computed from the non-zero voxels only, and we multiply the resulting images by the non-zero mask (different from foreground mask) to preserve the zero-valued voxels.



  
Distance Transform Maps A DTM represents the distance between each voxel and an object's closest boundary or edge. The values in a DTM are positive on the exterior, zero on the boundary, or negative in the object's interior. As a result, they can provide helpful information during training. Several loss functions and shape-aware networks use DTMs to improve the accuracy of medical imaging segmentation methods <cit.>. MIST allows users to pre-compute the DTMs for each label in the ground truth segmentation masks. Each channel in the resulting DTM represents the DTM for that label (i.e., channel one is the DTM for label one). If a label does not exist in the image, the DTM will be a 3D array with all values being the diagonal distance of the image. Additionally, MIST allows users to output normalized DTMs whose values range from -1 to 1. 



  §.§.§ Training

Once preprocessing is complete, MIST uses a five-fold cross-validation to train and evaluate a model on the dataset. In other words, for each fold we set aside 20% of the data as an independent test set and use the other 80% of the data for training. By default, MIST sets aside 10% of the training data as a validation set, but users can adjust this percentage. Additionally, users have the option to adjust the number of folds and can even specify custom folds if, for example, they are running leave-one-institution-out cross-validation.


MIST's default behavior is to use all available GPUs. When MIST runs multi-GPU training, it uses data parallelism via PyTorch's DistributedDataParallel package <cit.>. Additionally, MIST uses the Nvidia Data Loading Library (DALI) to handle data loading, patch extraction, and random augmentation. This feature further optimizes its computational performance by running most of the operations on the GPUs. 



  
Network Architectures MIST supports six different network architectures. These architectures include nnUNet (default) <cit.>, U-Net <cit.>, Swin UNETR <cit.>, and PocketNet <cit.>. MIST also provides users with options like adding deep supervision <cit.> or variational autoencoder regularization <cit.> to most of these architectures. Other regularization features like L2 and L1 regularization and gradient clipping are available for all architectures. For more details about the supported architectures (and their differences), please refer to our https://mist-medical.readthedocs.io/en/latest/advanced_topics/#network-architecturesdocumentation page. Finally, MIST's modular design allows users to implement their own architectures and integrate them into MIST with minimal effort. 




  
Loss Functions Like with the choice of network architecture, MIST allows users to choose from a variety of loss functions like Dice with Cross Entropy (default) <cit.>, the Generalized Dice Loss <cit.>, and boundary-based loss functions like the Boundary, Hausdorff, and Generalized Surface losses <cit.>. Again, MIST's modular design allows researchers to implement custom loss functions and seamlessly use them within the MIST framework. 



  
Other Training Options Finally, MIST gives users the flexibility to choose different optimizers like Adam (default) <cit.>, stochastic gradient descent, and AdamW <cit.>. MIST also supports several learning rate schedulers, with a constant learning rate of 0.0003 being the default. Additional features including transfer learning are included in our documentation page https://mist-medical.readthedocs.io/en/latest/here.



  §.§.§ Auxiliary Pipelines




  
Evaluation Once training is complete, the evaluation pipeline launches to assess the accuracy of the predictions with respect to the final classes defined by the user in their dataset JSON file. By default, these metrics are the Dice similarity coefficient and the 95th percentile Hausdorff distance (HD95). Other metrics like the surface Dice and average surface distance are also available. The output of the evaluation pipeline is a CSV file called  where each row shows the accuracy for each class and metric for one patient. The last five rows give the mean, standard deviation, median, and 25th and 75th percentiles for all of the metrics and classes. The evaluation pipeline is also available as a stand-alone pipeline that can run on any set of predictions. 



  
Postprocessing Postprocessing is available as an optional stand-alone pipeline. This pipeline allows users to specify which labels to apply a user-specified set of operations. These operations include the removal of small objects, only keeping the top-k components (adjustable value of k), morphological cleaning, and filling holes with a specified value. Once the postprocessing pipeline finishes, the evaluation pipeline launches and compares the change in accuracy to the baseline results. MIST computes this change in accuracy as a weighted average of the specified metrics. 



  
Test-Time Inference Finally, MIST provides a stand-alone inference pipeline that takes as input one or more MIST models and a CSV or JSON file with the paths to test-time data. This pipeline writes test-time predictions to a user-specified output directory. 



 §.§ BraTS 2024 Training Protocols
 
Our choice of architecture is the Pocket nnUNet with deep supervision (two supervision heads) and with residual convolution blocks <cit.>. MIST automatically selects a patch size of 128×128×128. We use eight NVIDIA H100 GPUs with a batch size of 32 uniformly distributed across the GPUs. Additionally, we use L2 regularization with a penalty parameter equal to 1e-5. Our choice of loss function is the Dice with Cross Entropy loss. While the default number of epochs per fold for MIST is 1,000, we train for 15,000 epochs per fold. We use a cosine learning rate schedule with an initial learning rate of 0.001 <cit.>. Within each fold, we set aside 2.5% of the training data (27 patients) as a validation set to select the best model. Once training is complete, MIST runs inference on the challenge validation data using test time augmentation (flipping along each axis and averaging each prediction) for each model. Inference uses sliding windows with an overlap of 0.5 with Gaussian blending (σ=0.125). The predictions from all five models are averaged to produce a final prediction. All other options and hyperparameters are left at their default values. Please refer to MIST's https://mist-medical.readthedocs.io/en/latest/documentation for these values.



§ RESULTS
 


 §.§ Accuracy

We use MIST and the parameters described in Section <ref> to perform a five-fold cross-validation (CV) with the BraTS Adult Glioma Post-Treatment Challenge data <cit.>. Table <ref> gives the results of this five-fold CV in terms of the Dice score and HD95. Tables <ref> and <ref> give the accuracy of the predictions in the validation set in terms of the same metrics and their lesion-wise versions.









 §.§ Computational Performance

In addition to the accuracy of our predictions, we also report the scalability of MIST for different numbers of A100 and H100 GPUs for multiple batch sizes. Figures <ref> and <ref> show these results for A100 and H100 GPUs, respectively. In these figures, we see that MIST effectively scales for multiple GPUs. This scaling is especially evident for the H100 GPUs. It is also important to note that we observe a roughly two times speed up with the H100 GPUs, which is in line with what Nvidia reports in their comparisons between the two types of graphics cards. 







§ DISCUSSION

Our results indicate that models trained using MIST can produce accurate adult post-operative gliomas, achieving median Dice scores of at least 0.9 for all of the segmentation classes in our five-fold CV. We choose to submit the raw output of our MIST pipeline to the BraTS 2024 competition, omitting the use of post-processing. However, future submissions to the continuous evaluation and competitions will explore different post-processing strategies with MIST by using its built-in post-processing pipeline. Indeed, previous BraTS competitions cite replacing small objects belonging to the ET class with the SNFH label as beneficial <cit.>. Additionally, we may see improvements in accuracy by using information from DTMs using boundary-based losses like the Generalized Surface Loss. Finally, different architectures like the multi-grid inspired FMG-Net or vision transformer-based architectures like Swin UNETR are also available on MIST <cit.>. These architectures may also improve our results.

MIST is able to utilize multiple GPUs to speed up training efficiently. This efficient use of multiple GPUs is especially true for H100s, where we see a nearly six times speed up when going from one to eight GPUs. The ability to take advantage of multiple GPUs and compute nodes effectively will be essential to keep training times manageable as larger datasets like TotalSegmentor become available <cit.>. As the size and quality of publicly available 3D medical imaging datasets increase, the creation of foundation models for medical imaging segmentation will be beneficial to speed up training through transfer learning <cit.>. MIST already supports transfer learning with other MIST models. Given this capability and its scalability, MIST can be instrumental in the creation of future foundation models.

In conclusion, our results show that MIST is capable of achieving accurate segmentation results while also scaling well to multiple GPUs by effectively taking advantage of technologies like PyTorch's DDP and Nvidia's DALI loader. MIST's simplicity, accuracy, and scalability have the potential to make MIST a valuable tool in the development of foundation models for medical imaging tasks like segmentation and radiation treatment planning. MIST's modular design makes it easy for researchers to add new loss functions or architectures to its pipelines. This modularity and easy integration into MIST's existing pipelines allow for new research to be evaluated and compared to existing methods fairly and consistently. MIST is still under active development. We encourage researchers and the broader medical imaging community to contribute to its development by using it, providing feedback, or contributing new features. 



§ AVAILABILITY & ATTRIBUTION

MIST is an open-source package (Apache 2.0 license) and is available on https://github.com/mist-medical/MISTGitHub or https://pypi.org/project/mist-medical/PyPI. Please cite this manuscript, <cit.>, and <cit.> if you use MIST for your own publications.



§ ACKNOWLEDGMENTS

The Department of Defense supports Adrian Celaya through the National Defense Science & Engineering Graduate Fellowship Program. This research was partially supported by the Tumor Measurement Initiative through the MD Anderson Strategic Research Initiative Development (STRIDE), NSF-2111147, NSF-2111459, and NIH R01CA195524


splncs04




 Paper Title: Knowledge-Guided Prompt Learning for Lifespan Brain MR Image Segmentation.



 
Brain MRI Segmentation with Knowledge-Guided Prompt Learning 



L. Teng et al. 


School of Biomedical Engineering & State Key Laboratory of Advanced Medical Materials and Devices, ShanghaiTech University, Shanghai 201210, China 
Shanghai United Imaging Intelligence Co., Ltd. Shanghai 200230, China 
feng.shi@uii-ai.com 
 
Shanghai Clinical Research and Trial Center, Shanghai 201210, China 
 
dgshen@shanghaitech.edu.cn 

Knowledge-Guided Prompt Learning for Lifespan Brain MR Image Segmentation
    Lin Teng 1 Zihao Zhao 1 Jiawei Huang 1 Zehong Cao 2 Runqi Meng 1 
Feng Shi 2^() 
Dinggang Shen 1,2,3^()
    August 5, 2024
===========================================================================================================



§ ABSTRACT  
Automatic and accurate segmentation of brain MR images throughout the human lifespan into tissue and structure is crucial for understanding brain development and diagnosing diseases. However, challenges arise from the intricate variations in brain appearance due to rapid early brain development, aging, and disorders, compounded by the limited availability of manually-labeled datasets. 
In response, we present a two-step segmentation framework employing Knowledge-Guided Prompt Learning (KGPL) for brain MRI.  
Specifically, we first pre-train segmentation models on large-scale datasets with sub-optimal labels, followed by the incorporation of knowledge-driven embeddings learned from image-text alignment into the models. The introduction of knowledge-wise prompts captures semantic relationships between anatomical variability and biological processes, enabling models to learn structural feature embeddings across diverse age groups. 
Experimental findings demonstrate the superiority and robustness of our proposed method, particularly noticeable when employing Swin UNETR as the backbone. Our approach achieves average DSC values of 95.17% and 94.19% for brain tissue and structure segmentation, respectively. 
Our code is available at https://github.com/TL9792/KGPLhttps://github.com/TL9792/KGPL.

 
 




§ INTRODUCTION
  
Dense segmentation of brain MR images involves voxel-wise labeling from different brain tissues, i.e., white matter (WM), gray matter (GM), cerebrospinal fluid (CSF), into fine-grained cortical and subcortical subregions, such as 106 structures according to Desikan-Killiany (DK) atlas <cit.> and Schaefer 400-parcels as described in <cit.>. 
Lifespan[Lifespan refers to the period between birth and death, emphasizing the wide age range covered in our study.] brain segmentation is essential for quantitative brain development analysis and neurodegenerative disorders diagnosis <cit.>. Manual annotation is time-consuming and labor-intensive. Hence, an automated brain image segmentation method is urgently needed. 

Several software packages, such as FMRIB Software Library (FSL) <cit.> and FreeSurfer <cit.>, have been developed to automatically segment brain images. However, processing each subject using these tools costs several hours, and the obtained results often require further manual corrections by experienced experts. Recent advancements in deep learning, such as convolutional neural network (CNN) based U-Net <cit.>, Vision Transformer (ViT) based methods <cit.> like UNETR <cit.>, and Swin Transformer based <cit.> Swin UNETR <cit.>, have shown promise in medical image segmentation. 

Brain morphology is dynamic in the human lifespan, leading to considerable appearance variations in brain images <cit.>. For instance, CSF exhibits gradual expansion, while GM experiences progressive reduction with aging, particularly among elderly individuals with Alzheimer's disease (AD).
However, most previous efforts are established on a single time point, and primarily focus on image-based information while overlooking the complementary textual information that provides essential contextual details and insights into the pathology or anatomy, which often results in lacking generalizability over a range of ages <cit.>. 
Additionally, the performance of learning-based methods is often restricted by the scarcity of annotated datasets, especially when labeled training samples for specific time points are limited in practice. To solve this problem, transfer learning <cit.> and fine-tuning <cit.> strategies have been widely adopted. 
Jia et al. <cit.> proposed an efficient fine-tuning method called Visual Prompt Tuning (VPT), which incorporates a small set of random learnable parameters into the input space. 
Directly applying random parameters representing unknown information may hinder model training <cit.>. 

To tackle the challenges mentioned above, we present a two-step framework with knowledge-guided prompt learning (KGPL) tailored for the dense segmentation of brain MR images across diverse age groups, which are the first to incorporate knowledge-wise prompts into brain segmentation models. Specifically, drawing inspiration from VPT, we replace random embeddings with knowledge-driven embeddings, learned from the pre-trained text encoder in BiomedCLIP[BiomedCLIP indicates Biomedical Contrastive Language-Image Pre-training (CLIP) <cit.>, which is pre-trained on large-scale figure-caption pairs by jointly training an image encoder and a text encoder, being able to well establish the semantic connection between images and language <cit.>.], into the input space of the encoder. The knowledge-wise learnable embeddings exploit the relationship between images and brain structural and developmental-specific biomedical characteristics, aiming to empower models to learn distinguishing representations of brain data acquired from different age phases. 
Experimental results demonstrate that our proposed method outperforms state-of-the-art (SOTA) baselines, and performs robustly in different backbone networks. 




§ METHODOLOGY
  
We aim to develop an automatic and robust approach to segment brain structure, based on brain MRI and specific biomedical attribute information. 
Fig. <ref> illustrates the overall framework. Specifically, we first pre-train segmentation models on large-scale datasets with sub-optimal annotation (Sec. <ref>).  
Subsequently, we refine the pre-trained models on relatively small-scale manually labeled datasets by introducing a few knowledge-driven learnable parameters into the input space, aiming to guide models to extract more crucial features of brain data spanning a large range of age. (Sec. <ref>). 


 




 §.§ Vision Pre-training
  
In this section, we utilize datasets with sub-optimal labels to establish models for segmenting brain MR images into tissue and structure, respectively. Specifically, we select one cutting-edge segmentation model as the backbone. Using the T1-weighted MR image as input, we first train the model to predict the brain tissue segmentation mask, which is supervised by the Dice loss function formulated as Eq. (<ref>). Subsequently, the predicted tissue segmentation results as input to train another model for brain structure parcellation, which is supervised by both Dice and Focal loss functions
as expressed in Eq. (<ref>). 


    Dice Loss = 1 - 2 × |P ∩ G|/|P|+|G|,
 
where P is the prediction result, and G denotes the ground truth.


    Loss = Dice Loss - α× (1 - p_t)^γ×log(p_t),
 
where α and γ are set to 100 and 0.2 based on experience.  p_t is the probability of the predicted result.   




 §.§ Fine-tuning with Knowledge-wise Prompt
   
To utilize specific textual information reflecting brain properties to enhance the discriminative learning of fine-grained brain anatomical structures, we introduce a novel KGPL approach. The procedure is illustrated at the bottom of Fig. <ref>.  




  §.§.§ Knowledge Prompt Generation.
  
Brain morphology, such as cortical thickness, shows a strong correlation with biological attributes like brain age, sex, and disease status <cit.>. This textual information comes from the corresponding data acquisition centers and official websites. 
To leverage this prior knowledge for guiding feature extraction, we initially categorize subjects into age groups every decade, considering that brain morphology remains relatively similar at close ages. 
Then individual attribute details are embedded into a template to generate a descriptive sentence (e.g., "This is a brain magnetic resonance image acquired from a male with mild cognitive impairment at fifty years old"). Finally, this sentence is encoded by a pre-trained text encoder from BiomedCLIP to obtain knowledge embeddings with dimensions (B, N, D), where B is the batch size, N is identical to the number of word tokens, and D is a fixed hidden dimension of 768.   



  §.§.§ Learnable Prompt Pre-initialized by Knowledge Prompt.
 
To effectively integrate prior knowledge into models, we first pre-define multiple learnable tokens (i.e., learnable embeddings) with zero values, matching with the dimension of knowledge embeddings. Then, the learnable embeddings are pre-initialized by knowledge embeddings through addition operation, referring to knowledge-wise learnable embeddings. 



  §.§.§ Model Fine-tuned Guided by Knowledge-wise Prompt.
 
Inspired by VPT, we serve knowledge-wise learnable embeddings as input to deep-level encoder layers, considering knowledge features are considered high-level. 
Specifically, we denote the collection of p knowledge-wise learnable tokens as P = { P_k ∈ℝ^M || k ∈ℕ, 1 ≤ k ≤ p}, where p is the number of word tokens. 
To ensure that they can interact with image embeddings, we reshape them on the hidden dimension. 
For U-Net or Swin UNETR as the backbone, knowledge-wise learnable embeddings are reshaped via the Adaptive Average Pooling (AAP) operator, transitioning from (B, N, D) to (B, N, 1). Following this, they pass through a linear layer with shape (B, C, N), where C is the number of channels; 
For UNETR, knowledge-wise learnable embeddings are first transposed from (B, N, D) to (D, N, B), then processed by a linear layer to match image embeddings' shape, except for the sequence length dimension. Next, they are inverse-transposed to restore the original dimension order. Finally, they are concatenated with image embeddings to calculate information interaction. The shape is (B, C, (N+S)), in which S represents the dot product of the feature map's length, width, and height. 
The procedure can be formulated as:  


    [X_i, _] = L_i-1 ([P_i-1,X_i-1]),
 


    [X_i+1, _] = L_i ([P_i, X_i]),
  
where X_i∈ℝ^B × C × S is image embeddings extracted by the i-1-th encoder layer L_i-1. 
P_i denotes knowledge-wise learnable embeddings at L_i's input space.
During the fine-tuning step, only the image embeddings are directed to the decoder.
Notably, we initialize the model of the target domain using weights learned from the source domain, by updating only the learnable parameters and the decoder while freezing the encoder.




§ EXPERIMENTS AND RESULTS
  


 §.§ Datasets and Metrics
  
Our datasets covering various age stages are collected from multiple repositories, including CBMFM (643 subjects, aged 18-81), Alzheimer’s Disease Neuroimaging Initiative (ADNI) (1285 subjects, aged 20-97) <cit.>, Autism Brain Imaging Data Exchange (ABIDE) (1062 subjects, aged 9-64) <cit.>, and Adolescent Brain Cognitive Development (ABCD) dataset (458 subjects, aged 9-11) <cit.>. Detailed analysis of demographic information can be found in the supplementary material.
Most datasets processed by FreeSurfer have sub-optimal labels. To obtain accurate labels for fine-tuning and model inference, FreeSurfer is initially applied, followed by manual annotations. Each dataset is randomly partitioned into training, validation, and testing sets in an 8:1:1 ratio. 
Each subject has a T1-weighted MR image, brain tissue segmentation mask, and structure segmentation mask (106 regions). T1-weighted MR images are pre-processed by skull stripping, bias field correction <cit.>, and intensity normalization. 
Performance evaluation utilizes two quantitative metrics: Dice Similarity Coefficient (DSC) and Average Surface Distance (ASD). 




 §.§ Implementation Details
 
Experiments are conducted using MONAI <cit.> on a single NVIDIA TITAN RTX GPU with 24GB memory. Training spans 1000 epochs with early stopping, utilizing an AdamW optimizer with a learning rate of 1e-4 and regularization weight of 1e-5. We adopt a linear warmup cosine annealing learning rate scheduler. Data augmentation involves randomly flipping along three axes with a probability of 0.5. Images are cropped to (128,128,128) based on foreground regions and a resolution of 1×1×1 mm^3 using linear interpolation. 




 §.§ Quantitative Comparison Analysis
  


  §.§.§ Effectiveness analysis of backbone architecture.
 
We validate the expansibility of our method based on three representative networks: (1) U-Net, a pure CNN; (2) UNETR, combining CNN and ViT; and (3) Swin UNETR, blending CNN and Swin Transformer. Their original architectures are used in our study, i.e., encoders have four CNN-based or twelve transformer-based blocks, decoders have four or five CNN-based blocks, and skip connections between the encoder and the decoder are included. 

 


Tables <ref> and <ref> present quantitative comparison results of brain tissue and structure segmentation using three backbones, respectively. Due to page constraints, we show partial brain structure segmentation results. The complete comparison results are shown in the supplementary material.
It is evident that our approach can significantly improve the performance using different backbones, especially with Swin UNETR obtaining a high prediction accuracy compared with other backbones. 
This indicates that our method is robust in applying to various network architecture types. 


 



  §.§.§ Comparison with full fine-tuning methods.
 
Quantitative analysis of Tables <ref> and <ref> shows that KGPL using Swin UNETR as the backbone improves average DSC and ASD values by 1.36% and 6.99%, 1.15% and 6.78% for brain tissue and structure segmentation, respectively.
Statistical significance of this improvement is confirmed by a paired t-test (p-value  0.05). Notably, our method achieves optimal performance with minimal trainable parameters, particularly in UNETR, where the parameter count is reduced from 92.78M to 4.46M without sacrificing performance. 
This suggests that KGPL aids models to prioritize essential information for brain image segmentation, thereby improving efficiency and adaptability. 



  §.§.§ Comparison with random-prompts-only methods.
 
Table <ref> of brain tissue segmentation shows that our proposed KGPL using UNETR as the backbone can offer 1.87% and 3.08% improvements of average DSC and ASD, respectively. In Table <ref>, KGPL using Swin UNETR shows improvements of 1.71% and 7.47% in average DSC and ASD for right hippocampus segmentation, respectively. Importantly, our proposed method is computationally faster (3×), i.e., extending the time from 3 days to 1 day. 
This emphasizes that knowledge-wise prompts can direct models to better comprehend and process brain image data with different characteristics and variations, thus promoting efficient convergence.


 



 §.§ Qualitative Comparison Analysis
  
To more intuitively illustrate the superiority of our method, we display qualitative comparison results of brain tissue and structure segmentation. 

 

In Figs. <ref> and <ref>, our KGPL method yields more accurate segmentation across different backbones, with predictions closer to ground truths compared to full or random prompt-only fine-tuning.
For example, Fig. <ref> demonstrates our brain tissue segmentation closely matching the ground truth, particularly in complex brain regions like sulci and gyri. In Fig. <ref>, full fine-tuning backbones exhibit significant errors in regions like the hippocampus and cerebral cortex, while our predictions display greater precision, suggesting potential for aiding in brain disease diagnosis, such as left/right hippocampal sclerosis detection.




§ CONCLUSION
 
This work presents a segmentation framework employing knowledge-guided prompt learning for brain MRI analysis throughout the lifespan. 
It stands as a pioneering effort in utilizing biomedical attribute knowledge as guidance to achieve precise segmentation for brain data. Experimental results indicate that knowledge benefits models in better learning structural variations of brain data across different age stages. 
The introduction of knowledge-wise learnable embeddings facilitates the extraction of fine-grained anatomical features and enhances training efficiency, even with small-scale datasets. 
Hence, the proposed method holds promise for clinical applications such as brain disease diagnosis. 
Moving forward, we plan to extend our methodology to infant phase datasets, enabling segmentation across the entire lifespan continuum. 





  §.§.§ 
 This study was partly funded by National Natural Science Foundation of China (grant numbers 62131015, 62250710165, U23A20295), the STI 2030-Major Projects (No.2022ZD0209000), Shanghai Municipal Central Guided Local Science and Technology Development Fund (grant number YDZX20233100001001), Science and Technology Commission of Shanghai Municipality (STCSM) (grant number 21010502600), and The Key R&D Program of Guangdong Province, China (grant numbers 2023B0303040001, 2021B0101420006).



  §.§.§ 

The authors have no competing interests to declare that are relevant to the content of this article. 



splncs04 
 



 Paper Title: Exploiting Change Blindness for Video Coding: Perspectives from a Less Promising User Study.







Exploiting Change Blindness for Video Coding: Perspectives from a Less Promising User Study 



 
    Mitra Amiri
ColorLab, dept. of Computer Science 

Norwegian University of Science and Technology (NTNU)

Gjøvik, Norway 

mitra.amiri@ntnu.no


Steven Le Moan
ColorLab, dept. of Computer Science 

Norwegian University of Science and Technology (NTNU)

Gjøvik, Norway 

steven.lemoan@ntnu.no
Christian Herglotz
Chair of Computer Engineering 

Brandenburgisch-Technische Universität Cottbus-Senftenberg

Cottbus, Germany 

christian.herglotz@b-tu.de
 
    Received —-; accepted —-
=====================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
What the human visual system can perceive is strongly limited by the capacity of our working memory and attention. Such limitations result in the human observer’s inability to perceive large-scale changes in a stimulus, a phenomenon known as change blindness. In this paper, we started with the premise that this phenomenon can be exploited in video coding, especially HDR-video compression where the bitrate is high. We designed an HDR-video encoding approach that relies on spatially and temporally varying quantization parameters within the framework of HEVC video encoding. In the absence of a reliable change blindness prediction model, to extract compression candidate regions (CCR) we used an existing saliency prediction algorithm. We explored different configurations and carried out a subjective study to test our hypothesis. While our methodology did not lead to significantly superior performance in terms of the ratio between perceived quality and bitrate, we were able to determine potential flaws in our methodology, such as the employed saliency model for CCR prediction (chosen for computational efficiency, but eventually not sufficiently accurate), as well as a very strong subjective bias due to observers priming themselves early on in the experiment about the type of artifacts they should look for, thus creating a scenario with little ecological validity.       



Video Compression, Change Blindness, HEVC (H.265)




§ INTRODUCTION

Although one might think that we see the surrounding world in great detail, it has been proved the human visual system (HVS) limits what we see due to limitations in our physiology as well as our working memory <cit.>. As a result, human observers tend to miss rather big and visible changes in their visual field. This phenomenon is usually referred to as change blindness (CB) <cit.>. Change blindness is dependent on individual parameters such as age, alertness, and even familiarity of the individual with the scene as well as the scene complexity <cit.>. The observer's failure to pick up information from some regions of a scene (Image or video) can be translated into the presence of irrelevancies. These irrelevancies are what lossy compressions take advantage of. 

Among the common media sources, raw videos have some of the largest files. During the last decade, the growth of streaming services <cit.> has led to increased demand for efficient codecs with high video-quality output. In addition, the growth in the consumer's interest in high dynamic range (HDR) technology has pushed streaming services to provide HDR video streams <cit.>. However, HDR video files are potentially larger than standard dynamic range (SDR) videos and would require heavier compressions in order to be streamed. The present compression methods such as advanced video coding (AVC)<cit.> and high-efficiency video coding (HEVC)<cit.> which are extensively used, use block-based motion compensation and content-adaptive entropy coding to encode the video. However, in none of the cases, the quantization parameter (QP) is spatially or temporally alternated based on saliency and attention.    
One of the earliest works done in order to include attention in video compression was done by Li et al. <cit.> where they suggested the modification of the bit allocation strategy in AVC encoding, based on the output of a saliency detection algorithm. In another work, Hadizadeh et al. <cit.> introduced the region-of-interest (ROI)-based method that uses the Itti-Koch-Niebur saliency model <cit.> to incorporate the saliency into the direct cosine transform (DCT) domain in the AVC encoding. In a similar work, Barua et al. <cit.> suggested using saliency maps to create high-fidelity regions incorporated into JPEG and MPEG-4 encoding for low-bitrate surveillance images and videos. However, despite the effectiveness of the discussed approaches, none of them has tried to incorporate CB. Saliency as an output of attention, although relevant to CB is not the only parameter affecting the observer's perception of the visual stimuli. Several studies have demonstrated the contribution of the limitations of visual working memory to CB <cit.>. In a recent work, Le Moan and Farup <cit.> took a further step and suggested a specifically defined change blindness map and used it to reduce the redundancies in regions with a high chance of CB occurrence. To our understanding, CB-based encoding has not yet been used in video compression.

In addition to the effect of saliency on CB, the graduality of the change is also effective in the change detection ability of HVS. The effect of graduality is most visible in the case of the "spot the difference" task. It has been shown that if the two images are shown in a fast-flickering manner, the observers can spot the difference quickly while if the frequency of the flickering is reduced, the observers struggle with finding the difference between the two images <cit.>. In a recent study, Alizadeh et al. <cit.> demonstrated that the human observer tends to fail following the gradual change. 

In this work, we would initially introduce our suggested novel algorithm of spatiotemporal encoding. We explore the effect of several parameters such as the introduction of CCR, the QP schedule curve, and the duration of the QP variation utilizing a subjective video quality assessment experiment. Afterward, we would focus specifically on the takeaways from our unfavorable subjective experiments.  
 


§ MATERIALS AND METHODS




 §.§ Video Dataset


The videos used in this study were taken from the JEVT dataset <cit.>. A total of four raw 10-bit videos in YUV format were selected as the source video sequence. The videos were selected in a way that they would cover both still and highly dynamic scenarios, crowded and uncrowded scenes, and the presence and absence of dominant human faces. All videos had a resolution of 1920×1080 pixels and a frame rate of 24 frames per second (fps). The duration of the video sources varied from 8.3s (199 frames) to 20.8s (499 frames). Figure  <ref> shows the four source videos employed for this study. 





 §.§ CCR Prediction Algorithm


To enable the spatially  QP-varying coding, it was necessary to predict the regions where CB was most probable to happen. These regions are referred to as Compression Candidate Regions (CCR). The prediction of CCR was carried out per frame and followed the SDSP algorithm suggested by Zhang et al. <cit.>. The algorithm relies on three types of priors, a frequency prior, a color prior, and a location prior. The frequency prior used in the SDSP algorithm is based on a research output of Achanta et al. <cit.>. The algorithm also counts for center bias which is simulated using a Gaussian map since according to the findings of Judd et al., <cit.>, human observers tend to pay more attention to the center of an image. Overall, the Saliency Detection by combining Simple Priors (SDSP) algorithm is computationally light and quite fast which makes it suitable for real-time encoding, yet as it can be seen, it is based on low-level visual features and does not take the semantic information of the scene into account. Additionally, SDSP does not account for the saliency of the moving objects and their tradeoff with center bias. Meanwhile, in the case of the chosen videos, in many instances, mobile attractive objects are present at the edges of the scene. Consequently, due to SDSP’s calculational speed, it was decided to proceed with this algorithm while excluding the center bias to partially account for the shortcomings of the saliency prediction algorithm. Figure  <ref> shows a sample of the saliency maps for each source video sequence. The complementary of these saliency maps were used to achieve the CCRs.   












 §.§ Video Encoding Tool

The videos were encoded using the Kvazaar open-source HEVC encoder <cit.>. Although the HEVC is designed to use the P-frames for inter-frame prediction, to minimize the observed effect of the frame-encoding paradigm in the suggested method, the encoding parameters were set in a way that during the entire sequence, only one I-frame (Intra-frame) would exist (Table  <ref>). As a result, all videos were encoded by starting with an I-frame followed by bidirectional frames (B-frames) throughout the video sequence and no prediction frames (P-frame) were present. The spatiotemporal varying QPs were included using the (roi) tag and a text file including the quantization parameter difference (Δ QP) values which will be described in more detail in the following section. Finally, to reduce the encoding complexity, the ultrafast preset was used.   







 §.§ Stimuli preparation
 
To introduce the CCR as well as the temporal gradual change in the HEVC encoding algorithm, a map of Δ QP for each frame was created. In the spatial domain, the Δ QP which will be referred to as the RoI map was generated by resizing the CCR map to the desired number of blocks in each dimension using cubic interpolation. Afterward, the RoI map was multiplied by the Δ QP values designated for each frame to create the spatially varying Δ QP map for each frame. Finally, the Δ QP map for all the frames was saved in a text file that would be used under the –roi tag. In our study, two scenarios were implemented. In the first scenario, the RoI map consisted of a single block, encoding the importance of the frame, and no spatial CCR was encoded. In the second scenario, each frame was divided using the block size (BS) of 10 (10 by 10 blocks). 

On the other hand, to exploit the blindness of human observers to the temporal changes, a gradual Δ QP schedule was employed that defined the designated Δ QP value for each frame throughout the video sequence. Figure  <ref> shows the two main Δ QP schedules used in our study.

The range of the frame numbers (nf) within which the gradual change happens varied between 16 frames, 32 frames, and the full video length. As for the Δ QP range, the higher limit was dictated by Kvazaar and the baseline QP. In the case of the frame numbers lower than the full video length, the schedule would repeat throughout the video sequence. As a result, to prevent the sharp descent of QP at the end of the cycle (typical in cubic schedule), the cubic schedule was only used for the full video-length schedules. 

Based on Kvazaar’s implementation, the highest possible QP while encoding using Kvazaar is 51, and since the baseline QP was set to 22 (Table 1) the upper boundary of the Δ QP range was set to 29. In order to push the video compression to an extent in which the compression artifacts are more visible, the lower boundary of the Δ QP range was set to 15. 



The combination of two frame numbers with gaussian schedule and the two schedules for the full video length frame number, with the two RoI scenarios led to an overall 8 configurations.  In the next step, the size of each encoded video was used as a target in an optimization loop through which a conventional HEVC-encoded video with constant QP was generated. Since the QP values could only be integers and the video size was not a continuous function of QP, the closeness of the two video sizes was ensured by introducing the constraint by which the constant QP videos could not be more than 5% smaller than the target video size. As a result, the final number of stimuli per image source was 16 videos and the overall number of stimuli was 64.



 §.§ Experiment design

To evaluate the general quality of the encoded videos and compare the quality of the saptio-temporally varying encoded videos with the conventional HEVC encoded videos, a psychophysical experiment was carried out. The stimuli included the 64 encoded videos with varying configurations, the four source videos along with six randomly chosen videos to account for intra-observer variability. In order to have the distance between the quality of different stimuli, the single-stimulus hidden-reference category judgment scheme was chosen <cit.>. According to the International Telecommunication Union (ITU) recommendations <cit.>, 5 absolute levels were selected ranging from “Bad” to “Excellent”.



 Figure  <ref> presents the scheme of the experiment. As it is portrayed, the experiment was carried out in two phases. Before the training phase, the observers were presented with the general instructions for the experiment. In the next step, the observers went through the training phase where they were presented with an example of an excellent and a bad video sequence. The video source used for this step was the "SunRise" video, taken from the JVET HDR dataset <cit.>. The observers also got to experience the rating interface during the training phase.







After finishing the training section, the observers moved on to the main experiment where they were presented with the randomized stimuli video sequences one at a time and were asked to rate the videos between "Very Bad" to "Excellent". The participants could observe each video sequence once and they were given unlimited time to rate the videos. After finishing the experiment, the observers were interviewed about the element they mainly focused on while observing each video source. 



 §.§ Experiment setup


In the experimental set up a 31.1-inch EIZO HDR CG3146 (4096×2160) was used which was calibrated using its built-in camera to match DCI-P3. The observers' seat was set at 1m distance from the display and a desk was put between the observer and the display to make sure that the observer did not approach the display at any time during the experiment. The experiment session was designed using MATLAB (2022b). The order of the video display was randomized and the user rating interface was written in MATLAB yet, the playback was carried out using ffplay. Using ffplay it was ensured that the video sequence is displayed in the original resolution. Since the display's resolution was higher than the resolution of the video sequence, the video sequence was shown over a neutral gray background.











 §.§ Observers

In general, 16 observers (8 male and 8 female) participated in the experiment. All the observers were students or members of the IDI department at NTNU (Gjøvik) and their range of expertise varied from familiar to highly experienced in the field of image and video processing. The observer's age varied from 22 to 36 years old and all had normal or corrected to normal vision.  



 §.§ Variability studies


In this study, the agreement percentage was employed as a measure of inter-observer variability to evaluate the variability of opinions among the observers. Agreement percentage can be expressed as the ratio of cases in which the observer's opinion score fell within the range of μ_n-1±σ_n-1 (where μ_n-1is the mean opinion score of all observers other than the studied observer and σ_n-1is the standard deviation) to all the observed cases. 
On the other hand, the consistency of observers in their opinions was calculated using the mean of the absolute opinion score distance (|Δ S|). Equation  <ref> presents the mean of absolute opinion score distance where m is the number of repeated stimuli and S_i is the opinion score for the i^th sample, Orig stands for original, and RP stands for repeated. 


    |ΔS|= ∑|S_i,Orig-S_i,RP|/m





§ RESULTS AND DISCUSSION



 §.§ Observer Variabilities

Figure  <ref>  shows the inter- (top) and intra-observer variability (bottom) among the 16 observers. As it can be seen, four observers (1,4,10 and 16) show high inter-observer variability in comparison with other observers. However, their agreement percentage is more than 50%, and looking at the intra-observer variability of these observers, none of them showed a mean opinion distance of higher than 2 which makes them reliable enough to use the data collected from them in further evaluations.  














Looking at the absolute opinion differences per each sample (Table.  <ref>), it can be seen that none of the chosen samples were highly confusing (more than 2 score distances) for all the observers. As a result, the present intra-observer variability can be related to the observer's reliability.











 §.§ Source Video Quality Ratings

Table   <ref>, portrays the MOS values for the source video sequences along with the size of the source videos. All the source videos approximately had the same bitrate of 1.2 Gbps. It can also be seen that all the source video sequences have been rated highly. Furthermore, no statistically significant quality difference was visible between the source videos. 























 §.§ Discussion

As described in section  <ref>, to answer the question of whether the proposed encoding algorithm outperforms the current HEVC encoding in terms of video quality; for each spatiotemporally encoded video, a conventionally HEVC encoded video (with constant QP) with the same size was generated. Table   <ref> lists the bitrate of the spatiotemporally encoded video along with its best size matching HEVC encoded video for each source video and each configuration scenario. As it can be seen the constant-QP size matched videos have a bit rate of less than 20% lower than their paired gradually changing-QP encoded videos. 

Figure  <ref> summarizes the quality MOS. In this figure, each column presents the cases for one source video sequence, each row represents one spatial encoding scenario, and within each diagram, different Δ QP schedules are presented. Looking at the graphs in Figure  <ref>, it can be seen that in none of the cases, the spatiotemporal encoding has led to significantly higher quality than the size-matched HEVC-encoded videos. Even in the case of the 16-frame Gaussian gradual change, the difference was reported as not significant based on the t-Student test (P =0.0926).

As it can be seen, either the qualities are rated closely or the constant-QP video is significantly better is content-dependent. However, in the case of temporal-only encoding (no CCR), no significant difference between the quality of the proposed encoding and the HEVC encoding was observed. While in the presence of CCRs, significantly lower quality was observed in the case of the spatially encoded videos with dynamic scenes (FireEaters and Runners). Based on the observed difference, it can be suggested that the saliency prediction algorithm fails to include the saliency raised from the movement of the objects. This conclusion is aligned with the points noticed in the interviews as the observers could obviously see the compression artifacts at the edges of moving objects such as the runners or the flames. Looking at the CCR maps, the presence of dark edges right next to highly salient regions would lead to strong compression artifacts at a very close distance to the salient regions. As a result, there is a high chance that the presence of the compression artifacts would capture the observers’ attention. The mentioned observations suggest that to successfully exploit change blindness using our proposed method the employed saliency prediction algorithm should predict the human observer's saliency very closely and account for the temporally-raised saliency.  

Nevertheless, it can be seen that both in the presence and absence of spatial Δ QP variation, in some video sources, the increase in the frame number has led to a significant decrease in the average perceived quality (between 16 and full video length frames). This can be due to the fact that when the gradual change occurs over a longer time frame, a bigger number of frames with low quality are consecutively shown to the observers, therefore, the observer has a longer time to recognize the artifacts and associate it with lower quality. Furthermore, a higher average perceived quality is observed in the case of the cubic schedule in comparison to the Gaussian schedule across the full video length while due to the large standard deviation, the difference is not significant. This observation can be explained by the fact that in the case of the Gaussian schedule, the highly compressed frame occurs in the middle of the video sequence leaving the observers a longer time to incorporate their observation with lower quality. Based on these observations, the frame number employed in the suggested algorithm is required to be optimized based on the Δ QP temporal schedule and in a way that the gradually appearing compression artifacts would escape observers' attention.  

By the comparison between the MOS values and the video sizes for the with and without CCR scenarios, it can be observed that the presence of CCR in encoding has led to higher bitrates. As a result, while the quality of the spatiotemporally encoded videos has not significantly changed, the size-matching HEVC files have quite larger sizes and therefore higher quality. Furthermore, the observers frequently reported paying attention to parts of the scene that the CCR map generally marks as less important. For instance, in the case of the Market, the observers paid a lot of attention to the wall and the tower, and the umbrella while the prediction model does not include any of these as highly salient regions. In another example, the observers frequently reported the runners' faces and the text on their shirts as important regions while the model notes the runner's shorts as the most salient regions. These observations suggest that not only does the SDSP algorithm fail in the case of dynamic scenes, but it also fails to properly mimic human attention even in less dynamic scenes such as the market. However, it must also be noted that all of the observers were to some extent experts in the field of image and video quality, and their attitude toward the task was affected by the bias caused by their expertise <cit.>. Based on the post-experiment interviews they were actively searching for regions where the artifacts were expected. It is known that CB is affected by attention  <cit.> and if the attention of the observer due to the task architecture or due to the expertise, is drawn to the locations where change is present, exploitation if CB becomes impractical. It is worth mentioning that although the common quality measurement methods are well-established, they are designed to measure the perceived quality in scenarios of low-level visual masking where the observer's attention does not play a major part in the perceived quality. Our observations suggest that the evaluation of perceived video quality in scenarios where high-level visual masking is present would require a new methodology where instructions, multi-sensory stimuli or distractors prevent the effect of expertise and task-allocated attention.       





§ CONCLUSION


Based on the principles of change blindness, a novel spatiotemporal video encoding approach for HDR and SDR video content within the framework of HEVC encoding was proposed. The subjective quality of the encoded videos was then compared to the quality of HEVC videos of the same size to determine the effectiveness of the proposed method. However, no significant enhancement in the subjective quality of the spatiotemporally encoded video was observed and even in some cases, the HEVC encoded video significantly outperformed the suggested encoding method. The results illustrated that the use of a proper saliency prediction algorithm for the prediction of the compression candidate region plays a key role in the proposed algorithm. It was observed that the saliency prediction should include both temporal and semantic saliency in order to be successful in exploiting change blindness. Additionally, it was observed that if the gradual change is carried out over a shorter period, the observers perceive less drop in the quality.         
Although the initial results did not show significant enhancement in the spatiotemporally encoded video, the current work can be used as a ground to explore better and more semantic CCR prediction models. Additionally, in this work, the effect of block size and its correlation with the CCR map accuracy was not studied.  Exploring these parameters can potentially lead to the enhancement of the proposed algorithm. Additionally, it was shown that the use of experts in the subjective study could considerably affect the results when exploring change blindness applications due to the alteration of their attention allocation based on their expertise-based bias. Finally, it was observed that the common quality measurement methodology is less suitable for evaluation of attention-dependent perceived quality.      



00
b1 E. Awh, E. K. Vogel, and S. H. Oh, “Interactions between attention and working memory,” Neuroscience, vol. 139, no. 1, pp. 201–208, 2006.
b2 T. J. Smith and J. M. Henderson, “Edit Blindness: The relationship between attention and global change blindness in dynamic  scenes.,” J Eye Mov Res, vol. 2, no. 2, Dec. 2008.
b3 D. J. Simons, S. L. Franconeri, and R. L. Reimer, “Change blindness in the absence of a visual disruption,” Perception, vol. 29, no. 10, pp. 1143–1154, 2000.
b23 S. le Moan, &  M. Pedersen. "A three-feature model to predict color change blindness". Vision (Switzerland), 3(4), 2019.

b4 S. Lee, S. Lee, H. Joo, and Y. Nam, “Examining factors influencing early paid over-the-top video streaming market growth: A cross-country empirical study,” Sustainability (Switzerland), vol. 13, no. 10, 2021.
b5 J. McNamee, J. Hatchett, K. Debattista, and A. Chalmers, “Live HDR video streaming on commodity hardware,” in Applications of Digital Image Processing XXXVIII, SPIE, p. 95990U, 2015.
b6 T. Wiegand, G. J. Sullivan, G. Bjøntegaard, and A. Luthra, “Overview of the H.264/AVC video coding standard,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 13, no. 7, pp. 560–576, Jul. 2003.
b7 G. J. Sullivan, J. R. Ohm, W. J. Han, and T. Wiegand, “Overview of the high efficiency video coding (HEVC) standard,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 22, no. 12, pp. 1649–1668, 2012.
b8 Z. Li, S. Qin, and L. Itti, “Visual attention guided bit allocation in video compression,” Image Vis Comput, vol. 29, no. 1, pp. 1–14, 2011.
b9 H. Hadizadeh and I. V. Bajic, “Saliency-aware video compression,” IEEE Transactions on Image Processing, vol. 23, no. 1, pp. 19–33, 2014.
b10 L. Itti, C. Koch, and E. Niebur, “Short Papers A Model of Saliency-Based Visual Attention for Rapid Scene Analysis,” 1998.
b11 S. Barua, K. Mitra, and A. Veeraraghavan, “Saliency guided Wavelet compression for low-bitrate Image and Video coding,” IEEE Global Conference on Signal and Information Processing (GlobalSIP), pp. 1185–1189, 2015.
b20 T. F., Brady, & J. B.Tenenbaum, . A "Probabilistic Model of Visual Working Memory: Incorporating Higher Order Regularities Into Working Memory Capacity Estimates". Psychological Review, 120(1), p. 85–109, 2013.
b21 M. A. Cohen, D. C. Dennett, & N. Kanwisher, "What is the Bandwidth of Perceptual Experience?", Trends in Cognitive Sciences, Vol. 20, Issue 5, pp. 324–335, 2016.
b12 S. Le Moan and I. Farup, “Exploiting Change Blindness for Image Compression,” in Proceedings - 11th International Conference on Signal-Image Technology and Internet-Based Systems, SITIS 2015, Institute of Electrical and Electronics Engineers Inc., pp. 89–95, 2016. 
b13 A. Alizadeh Mivehforoushi, “A New Descriptive Subjective Video Quality Dataset,” Master’s Thesis, NTNU, Gjovik, 2023.
b14 L. Zhang, Z. Gu, and H. Li, “SDSP: A NOVEL SALIENCY DETECTION METHOD BY COMBINING SIMPLE PRIORS,” IEEE International Conference on Image Processing, pp. 171–175, 2013.
b15 R. Achanta, S. Hemami, F. Estrada, and S. Susstrunk, “Frequency-tuned Salient Region Detection,” in IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Maiami: IEEE, pp. 1597–1604,2009.
b16 T. Judd, K. Ehinger, F. Durand, and A. Torralba, “Learning to Predict Where Humans Look,” in IEEE 12th International Conference on Computer Vision, IEEE,  pp. 2106–2113, 2009.
b17 A. Yla-Outinen, A. Lemmetti, M. Viitanen, J. Vanne, and T. D. Hamalainen, “Kvazaar: HEVC/H.265 4K30p Intra Encoder,” in Proceedings - 2017 IEEE International Symposium on Multimedia, ISM 2017, Institute of Electrical and Electronics Engineers Inc., Dec. 2017, pp. 362–363.
b18 International Telecommunication Union (ITU), “Subjective video quality assessment methods for multimedia applications,” ITU Recommendation. p. 910, 2022.
b19 A. Segall, E. Francois, W. Husak, S. Iwamura, and D. Rusanovskyy, “JVET common test conditions and evaluation procedures for HDR/WCG video,  document No, "JVET-T2011”, 2020.
b22 S. Werner&  B. Thies, "Is "Change Blindness" Attenuated by Domain-specific Expertise? An Expert-Novices Comparison of Change Detection in Football Images", Visual Cognition, 7:1-3, 163-173, 2000.



 




 Paper Title: AMAES: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation.






Augmented Masked Autoencoder Pretraining for 3D Brain Segmentation






A. Munk et al.




Department of Computer Science, University of Copenhagen, Copenhagen, Denmark Pioneer Centre for AI, Copenhagen, Denmark 

{asmu,jaam}@di.ku.dk
Equal contribution



: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation
    Asbjørn Munk1,2* Jakob Ambsdorf1,2* Sebastian Llambias1,2 Mads Nielsen1,2

    August 5, 2024
==============================================================================================



§ ABSTRACT 
This study investigates the impact of self-supervised pretraining of 3D semantic segmentation models on a large-scale, domain-specific dataset. We introduce , a dataset of 44,756 brain MRI volumes from public sources, the largest public dataset available, and revisit a number of design choices for pretraining modern segmentation architectures by simplifying and optimizing state-of-the-art methods, and combining them with a novel augmentation strategy. The resulting AMAES framework is based on masked-image-modeling and intensity-based augmentation reversal and balances memory usage, runtime, and finetuning performance. Using the popular U-Net and the recent MedNeXt architecture as backbones, we evaluate the effect of pretraining on three challenging downstream tasks, covering single-sequence, low-resource settings, and out-of-domain generalization. The results highlight that pretraining on the proposed dataset with AMAES significantly improves segmentation performance in the majority of evaluated cases, and that it is beneficial to pretrain the model with augmentations, despite pretraing on a large-scale dataset. Code and model checkpoints for reproducing results, as well as the  dataset are available at <https://github.com/asbjrnmunk/amaes>.









§ INTRODUCTION





Large-scale self-supervised pretraining promises immense potential for medical image segmentation. Works in computer vision from the natural image domain demonstrate that state-of-the-art results are achieved by scaling model capacity and pretraining on large datasets, enabling increased generalization performance, few- or zero-shot predictions, and highly semantic embeddings <cit.>. 


To enable self-supervised representation learning in vision, various methods have been proposed. Early works involved reversing image degradations, such as denoising images <cit.>, solving jigzaw puzzles <cit.>, or re-colorizing images <cit.>.
Masked-Image-Modeling (MIM) proved to be an especially successful and scalable variant, where patches of the input image are masked and learned to be reconstructed either directly or indirectly by predicting an embedding <cit.>. A Masked Autoencoder (MAE) is an encoder-decoder architecture that reconstructs the entire image from a masked version <cit.>.
Another family of approaches applied contrastive learning, where semantically similar transformed inputs are mapped closely in a latent space, while dissimilar inputs are pushed apart <cit.>. Self-distillation is based on a related idea, enforcing similar embeddings for semantically close inputs for a student and teacher network <cit.>.
Recent approaches combined these image-level contrastive or self-distillation objectives with patch-level MIM to enhance the spatial representation performance for downstream segmentation tasks <cit.>.

Evidence from medical vision indicates that the benefits of pretraining are transferable to biomedical image segmentation <cit.>. Yet, the current method of choice for solving segmentation tasks is still a fully-supervised paradigm <cit.>. Due to a heavily optimized data augmentation pipeline, the fully-supervised approach can be effective even with limited labeled data <cit.>. However, to fully leverage the potential of large-scale models in medical image segmentation as witnessed in computer vision, more training data is required, and the supervised approach is inherently limited by labeling cost. Large-scale unlabeled medical imaging datasets could resolve this need, but they are not readily available to the public, hindering progress for methodological research on how to harness unlabeled data for downstream segmentation most effectively. To address the need for pretraining datasets in the brain MRI domain, we propose , to the best of our knowledge the largest public pretraining dataset for brain MRI to date[  is an order of magnitude larger than the pretraining dataset used in SwinUNETR, the current state-of-the art methodology.] comprising of 44,756 volumes collected from public sources.

Using this large-scale dataset, we are proposing the AMAES (Augmented Masked Auto Encoder for 3D Segmentation) framework for pretraining modern semantic segmentation architectures. AMAES combines the successful Masked Image Modeling (MIM) task <cit.> with augmentation reversal, both of which can be viewed as special cases of denoising autencoding  <cit.>, with the goal of learning local, detailed representations that are invariant to intensity shifts between different domains. The resulting framework balances factors such as GPU memory usage, pretraining runtime, and downstream task performance. We compare AMAES against the state-of-the-art pretraining method proposed for SwinUNETR <cit.>, which combines MIM with a rotation prediction task and a contrastive learning objective. We find that in practice, solving the rotation problem is trivial while the contrastive loss adds little benefit provided the relatively small batch sizes possible in 3D imaging. Our simplified framework results in a significant speedup and memory reduction while improving downstream performance.

We evaluate our method using two backbone architectures: (i) The U-Net, as the most popular semantic segmentation model, and (ii) MedNeXt, as a recently proposed modernization. The  dataset and proposed AMAES framework for pretraining segmentation models contribute to building a foundation for the next generation of large-scale biomedical image segmentation models. 







§  DATASET


We propose   a dataset of 44.756 Brain MRI volumes collected from public sources, featuring a diverse set of acquisition parameters and patient populations. The dataset is a compilation of data from four large non-labelled datasets (ADNI <cit.>, OASIS3 <cit.>, OASIS4 <cit.>, PPMI <cit.>) and five challenge datasets (MSD <cit.>, BraTS21 <cit.>, ISLES22 <cit.>, WMH <cit.>, MSSEG1 <cit.>). The resulting dataset is highly heterogeneous and assembled to simulate the characteristics of clinical data. Overview of the dataset can be found in Table <ref>. The dataset includes data acquired at both 1.5T and 3T. PPMI data is openly available from PPMI. ADNI data used in preparation of this article were obtained from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database. To compile the diverse data into a suitable pretraining dataset, we transformed all volumes to the RAS coordinates system, resampled to isotropic 1mm spacing, clipped values to the 99% percentile, and z-normalized on a volume level. Afterwards, each volume is cropped to the minimum bounding box of the brain. The preprocessing pipeline is based on the  framework for medical imaging <cit.>. Code to reproduce the final version of the dataset is provided in the project repository. Usage of the dataset needs to comply with the individual legal requirements of each source.




§ METHOD



In this section we briefly describe . Our approach is highlighted in Figure <ref>. The overall goal of the framework is to provide an effective and efficient pretraining strategy for large-scale datasets, requiring a balance between memory usage, runtime and downstream performance (see Figure <ref>).



 §.§ Pretraining


  is based on a Masked Autoencoder (MAE) architecture with a light weight decoder and a specific augmentation scheme. Importantly, the architecture assumes no specific choice of encoder and can be used with any encoder that is compatible with the skip connections of the U-Net decoder.

Pre-text task. Prior work has centered around 2D natural images <cit.>, however, previous work showed that a Masked Autoencoder works well in 3D by applying it to 3D CT Brains using the SwinUNETR architecture <cit.>. The SwinUNETR pretraining objective further included both a rotation and a contrastive loss. However we find that the rotation task was near-trivial for the models to solve and thus provided no tangible utility, while the contrastive loss resulted in the memory footprint during pretraining to be doubled, without significant performance gains (see Appendix <ref>). Threfore, we reduce the pre-text task for  to masked reconstruction, and instead increase the batch size and double the number of training epochs. For each 3D input patch, the input is divided into subpatches of size 4×4×4 voxels, and each subpatch is masked with probability 0.6, following prior work <cit.>.



Decoder architecture. Since the decoder is discarded after pretraining, it is desirable to use as small a decoder as possible. This not only decreases training time and memory usage, but also forces the network to learn expressive representations in the encoder.

We find that a light weight version of the standard U-Net decoder provides a good balance between flexibility, memory efficiency and parameter count. Table <ref> shows results of pretraining with two versions of the U-Net decoder: the standard 3D U-Net decoder and one with half the number of convolutional layers (denoted light weight decoder). Interestingly, the light weight decoder not only decreases pretraining time with 20%, but also further increases finetuning performance by 1.96 dice points.

Skip connections. While skip connections are used during finetuning, our proposed encoder-decoder architecture does not include skip connections during pretraining. We found that using skip connections during pretraining tended to slightly adversely affect downstream performance, with a finetuning dice score of 56.22 with and 56.65 without, on 6-fold cross-validation on the BraTS21 dataset. While the difference is not significant, we choose to not use skip connections in the pursuit of the simplest possible setup.

Augmentations. To learn representations which are robust to changes in the input domain, we apply an extensive augmentation pipeline to the data during pretraining. The full list of augmentations include all augmentations used in nnUnet <cit.> and is given in Appendix <ref>. We want our representations to be covariant to any spatial augmentations (rotation, scaling), while we seek to learn representations that are invariant to all other augmentations, for instance additive noise. We thus only apply spatial augmentations to the reconstruction target. This is visualized in Figure <ref>. Table <ref> shows the result of applying augmentations to the pretraining data: when applying augmentations during pretraining, we see a 1.24 Dice point increase in finetuning performance, compared to no augmentations, when finetuned with only spatial augmentations.



 §.§ Finetuning


For finetuning, we transfer the backbone encoder weights from pretraining and attach a new standard U-Net decoder with skip connections between the encoder and the decoder. We only use spatial augmentations (rotation, scaling) during finetuning, since this setup performed better empirically. Upon closer inspection, the learning curves showed signs of catastrophic forgetting during the second half of training when finetuned with the full augmentation package.









§ EXPERIMENTAL DESIGN



 §.§ Datasets

We evaluate   on three different segmentation tasks covering tumor (BraTS21 <cit.>), ischemic stroke lesion (ISLES22 <cit.>) and white matter hyperintensity (WMH <cit.>). To test the models few-shot learning ability, we evaluate the models in a low resource setting: we only use n_train = 20 labeled examples for finetuning, while we evaluate the models on the remaining data. We choose n_train = 20 to realistically model what one medical expert can manually label with limited effort. Since the pretraining dataset considers each sequence individually, we evaluate the model's ability to perform downstream tasks in a single sequence setup, and thus for BraTS21 and WMH we consider only the flair modality, for ISLES22 we consider only the DWI modality.

Hyperparameters were selected on 3-fold cross-validation splits of the BraTS21 dataset. We do not tune hyperparameters individually for each fold or for each dataset. All results are averages over 6 folds. For BraTS21, we evaluate the model on a 20% holdout set. For ISLES22 we follow <cit.> and provide results from the rotating cross-validation splits. For WMH we both perform cross-validation evaluation (denoted "WMH") and test set evaluation. The test set contains data from scanners not present in the training set, which are also not present in , and thus evaluates the models ability to generalize to unseen data. We call this dataset "WMH OOD" (out-of-distribution). 



 §.§ Configurations

Our baselines, pretraining and finetuning are implemented using PyTorch, PyTorch Lightning, and the   library for medical image analysis <cit.> as a back-bone. All models are trained using 16-bit mixed-precision to limit memory footprint. Models were trained patch-based, using an input patch size of 128 × 128 × 128. Models are pretrained for 100 epochs (one epoch is one patch from each volume) and trained supervised for 2500 steps on downstream datasets. The exact hyperparameters used for each model is given in Appendix <ref>. All models are pretrained on two H100 GPUs, and fine-tuned on a single H100 GPU.



 §.§ Baselines

We evaluate  along two dimensions: (i) How does   compare to the pretrained SwinUNETR?, and (ii) What is the impact of pretraining on downstream performance?

To evaluate (i), we pretrain SwinUNETR with the exact configuration given in <cit.> on , which includes both a contrastive loss, a rotation loss and a reconstruction loss, as well as a different choice of masking ratio and mask size. To ensure a fair comparison, SwinUNETR is pretrained for 100 epochs, similar to , and is pretrained with patch size 128^3, instead of patch size 96^3 and 90 epochs. To address (ii), we apply   to a set of convolutional backbones and evaluate the difference between training from scratch and finetuning the pretrained model. The backbone models include a U-Net in two sizes: XL (90M parameters) and B. The U-Net B is very similar in size to the one used in nnUNet (22M parameters, when trained with the default setting  of 12 GB). Further, we explore using the modernized U-Net architecture MedNeXt <cit.>, which uses depth-wise seperable convolution to introduce compound scaleable medical segmentation models. All MedNeXt models are trained with kernel size 3 and do not use UpKern. The networks are finetuned and trained from scratch using the same hyperparameters. Models trained from scratch use the full augmentation pipeline.







§ RESULTS AND DISCUSSION


The results of our experiments are given in Table <ref>. Both U-Net B, U-Net XL and MedNeXt models benefit from pretraining in all scenarios. MedNeXt M does not benefit from pretraining on BraTS21 and ISLES22, however exhibits very good performance on WMH both in and out of domain. We believe this is mainly due to strong baseline performance, before pretraining is added. Note that models do not use the full augementation pipeline during finetuning, while the baselines from scratch use all augmentations. All models trained with   outperform the pretrained SwinUNETR, except for the standard U-Net on WMH OOD. Generally, model size seems to be correlated with performance, with two notable exceptions: The SwinUNETR model seems to be worse than both MedNeXt M model and the standard U-Net, both of which have less than a third of the parameters of SwinUNETR. The results on WMH OOD demonstrate that pretraining also increases performance when applied to out-of-domain data – even with respect to the pretraining domain.

Limitations. This study only considers single sequence and low resource settings. We have not applied   to the multi-sequence versions of the dataset and neither have we applied it to datasets with more datapoints than n_train = 20. How to leverage pretraining for multi-sequence tasks is a relevant avenue for future work. This study neither tests the SwinUNETR framework using different backbones and nor tests   using a SwinUNETR backbone, since it is not readily compatible with a U-Net decoder.[Since it uses non-pretrained UNETR <cit.> encoder blocks in the interface between the Swin-Transformer and the convolutional decoder.]



§ CONCLUSION

We present , a framework based on Masked Image Modeling to pretrain 3D-native segmentation models which balances memory usage, runtime, and finetuning performance. We further introduce , the largest dataset of brain MRI for pretraining to date. We pretrain   on   and evaluate the models on three challenging datasets. We compare the framework to models trained from scratch as well as the pretrained SwinUNETR model. Our results demonstrate that when  is used to pretrain U-Net models or MedNeXt models, the models outperform similar models trained from scratch, as well as SwinUNETR models, both pretrained and not in the challenging tasks evaluated.
























    





Supplementary Material



§ HYPERPARAMETERS






§ AUGMENTATIONS







§ ADNI ACKNOWLEDGEMENT


Data used in preparation of this article were obtained from the Alzheimer’s Disease Neuroimag- ing Initiative (ADNI) database (adni.loni.usc.edu). As such, the investigators within the ADNI contributed to the design and implementation of ADNI and/or provided data but did not par- ticipate in analysis or writing of this report. A complete listing of ADNI investigators can be found 
http://adni.loni.usc.edu/wp-content/uploads/how_to_apply/ ADNI_Acknowledgement_List.pdfhere.

The ADNI was launched in 2003 as a public-private partnership, led by Principal Investigator Michael W. Weiner,MD. The primary goal of ADNI has been to test whether serial magnetic reso- nance imaging (MRI), positron emission tomography (PET), other biological markers, and clinical and neuropsychological assessment can be combined to measure the progression of mild cognitive impairment (MCI) and early Alzheimer’s disease (AD). For up-to-date information, see <www.adni-info.org>.

Data collection and sharing for this project was funded by the Alzheimer’s Disease Neuroimaging Initiative (ADNI) (National Institutes of Health Grant U01 AG024904) and DOD ADNI (Depart- ment of Defense award number W81XWH-12-2-0012). ADNI is funded by the National Institute on Aging, the National Institute of Biomedical Imaging and Bioengineering, and through generous contributions from the following: AbbVie, Alzheimer’s Association; Alzheimer’s Drug Discovery Foundation; Araclon Biotech; BioClinica, Inc.; Biogen; Bristol-Myers Squibb Company; CereSpir, Inc.; Cogstate; Eisai Inc.; Elan Pharmaceuticals, Inc.; Eli Lilly and Company; EuroImmun; F. Hoffmann-La Roche Ltd and its affiliated company Genentech, Inc.; Fujirebio; GE Healthcare; IXICO Ltd.;Janssen Alzheimer Immunotherapy Research & Development, LLC.; Johnson & John- son Pharmaceutical Research & Development LLC.; Lumosity; Lundbeck; Merck & Co., Inc.;Meso Scale Diagnostics, LLC.; NeuroRx Research; Neurotrack Technologies; Novartis Pharmaceuticals Corporation; Pfizer Inc.; Piramal Imaging; Servier; Takeda Pharmaceutical Company; and Transi- tion Therapeutics. The Canadian Institutes of Health Research is providing funds to support ADNI clinical sites in Canada. Private sector contributions are facilitated by the Foundation for the National Institutes of Health (www.fnih.org). The grantee organization is the Northern California Institute for Research and Education, and the study is coordinated by the Alzheimer’s Therapeutic Research In- stitute at the University of Southern California. ADNI data are disseminated by the Laboratory for Neuro Imaging at the University of Southern California.




§ CONTRASTIVE AND ROTATION LOSS CURVES






§ RESULTS WITH STANDARD DEVIATION








 Paper Title: FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication.


	
	
	
	 
	
FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication
    
		
			Yuna Yan1, 
			Xin Zhang1, 
			Lixin Li1, 
			Wensheng Lin1, 
            Rui Li2,
            Wenchi Cheng3,
			and Zhu Han4 
		1School of Electronics and Information, Northwestern Polytechnical University, Xi’an, China, 710129
        2Samsung AI Center, Cambridge, UK
        3State Key Laboratory of Integrated Services Networks, Xidian University, Xi'an, China, 710071
		4Department of Electrical and Computer Engineering, University of Houston, Houston, TX, 77004
	
    August 5, 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
In this paper, we address the problem of image semantic communication in a multi-user deployment scenario and propose a federated learning (FL) strategy for a Swin Transformer-based semantic communication system (FSSC). Firstly, we demonstrate that the adoption of a Swin Transformer for joint source-channel coding (JSCC) effectively extracts semantic information in the communication system. Next, the FL framework is introduced to collaboratively learn a global model by aggregating local model parameters, rather than directly sharing clients' data. This approach enhances user privacy protection and reduces the workload on the server or mobile edge.
Simulation evaluations indicate that our method outperforms the typical JSCC algorithm and traditional separate-based communication algorithms. Particularly after integrating local semantics, the global aggregation model has further increased the Peak Signal-to-Noise Ratio (PSNR) by more than 2dB, thoroughly proving the effectiveness of our algorithm. 




Semantic communication, federated learning, swin Transformer, privacy protection.




§ INTRODUCTION

Wireless communication is embracing an explosive growth of data traffic semantic nature, generated by various contemporary multi-media applications such as AR/VR, video streaming, gaming and telehealth. Traditional communication systems take the source messages (e.g. images and texts) and first encode these messages with source coding algorithms on the application layer, in order to remove redundant information. Subsequently, when the compressed information reaches the physical layer, it is then further channel-encoded by injection of redundant bits, to be resilient to the noise in the communication channel. There are three major factors that make this current approach inefficient in the 5th Generation Mobile Communication Technology(5G) and future generations of mobile networks. Firstly, although modern channel coding algorithm on its own has theoretically reached the Shannon limit on the physical layer <cit.>, when jointly considering the source coding and channel coding, the Shannon's optimality theory only holds when the message blocks are asymptotically infinitely long, which is impractical in the real-world <cit.>. Secondly, the separation-based approach<cit.> compresses and then upscales messages, thus inherently introducing extra latency, which lead to inferior performances for major time-sensitive 5G/6G and upcoming applications, e.g. tactile internet and autonomous driving. Third, channel coding strategies treat all bits of information equally, missing the opportunity to utilize the abundant semantic information available in the messages, especially in image and video contents.

To address these issues in the traditional communication regime, semantic communication<cit.> a.k.a. joint source and channel coding, has been proposed and studied, which focuses on the transmission of meaning rather than bits, helping to escape the “Shannon trap”<cit.>. 
Compared to conventional communication methods, semantic communication is centered around the relevant semantic features essential for transmitting information, thereby eliminating unnecessary redundancy. This approach leads to a significant reduction in data transmission latency and bandwidth consumption. To enhance the versatility of semantic communication systems, we propose the employment of a Swin Transformer-based semantic communication (STSC) system that primarily focuses on the reconstruction of global image signals. Its unique architecture enables the image-oriented semantic communication system to effectively capture deeper semantic features from potential semantic information and represent complex semantic relationships within the input data. In addition, the Swin Transformer is capable of reducing computational complexity and speeding up both model training and inference compared to traditional Transformer models. As a result, it is well-suited for handling large-scale semantic communication, especially in multi-user deployments.



However, image-oriented semantic communication targets applications that often involves users' private data, such as their faces, home interior, license plates, etc., which cannot always be uploaded to the cloud server. Moreover, most of the prior art on image-based semantic communication are point-to-point systems, which do not fit in many applications where multiple users are involved. In this paper, we propose a multi-user federated learning (FL) framework for semantic communication, showcasing the potential of efficient learning of a global semantic communication model that combines the learned knowledge from multiple users, without direct access to users' private data. We demonstrate that the proposed FL semantic system, when powered with a hierachical transformer architecture, i.e. the Swin Transformer, achieves a better image transmission quality at a faster convergence rate compared to the single-user.

In summary, the main contributions of this paper are summarized as follows:


	
  * We propose a Swin Transformer based semantic communication architecture, which aims at minimizing the difference between the transmitted and received images, so as to effectively communicate image semantic information.
	
  * In contrast to the centralised approach of sharing the user data directly to a cloud server in order to train a transformer neural network for semantic communication, we consider a cross-silo FL setting that employs multiple users to train locally their own model using local data, and share the local updates with a global server, thereby preserving users' data-privacy.
	
  * Through the use of well-established dataset for computer vision, simulation results show that our proposed framework surpasses the representative JSCC methods, and further boosts peak signal-to-noise ratio (PSNR) by over 2dB on top of aggregating local semantics.


The rest of this paper is organized as follows. In Section <ref>, we introduce the framework of the image semantic communication system based on FL and formulate the corresponding research problem. We then present the federated learning strategy for a Swin Transformer-based semantic communication (FSSC) in Section <ref>. In Section <ref>, we present numerical simulation results and showcase the semantic communication performance of FSSC as well as the convergence of the system. Finally, the paper is concluded in Section <ref>.




§ SYSTEM SETUP


In this paper, we consider a multi-user FL semantic communication system, as shown in Fig. <ref>, which consists of a cloud server and multiple clients. An FSSC client can be an edge server or a terminal device. When the client is an edge server, it is able to perform computational tasks, hence the training and inference of STSC happens on the client locally. When the client is a light-weight terminal device equipped with only limited or no computational power, the client should be served by a corresponding edge server which is trusted by the client to share data with in order to perform training. The federated training process is detailed as follows:

	
  * The cloud server initialize the STSC model for semantic communication, and distribute it to all participating clients. Each client will then perform n epochs of local training using its own data. Once the local training is completed, all clients transmit the updated model parameters to the cloud server, and the server will then aggregate the updated model parameters following Federated Averaging <cit.>, which completes one round of the federated learning. This process is repeated until the model converges;
	
  * During deployment of a trained STSC model, a user will initiate a communication request and sends a semantic message to the client. The client will then process this message by performing an inference using the trained STSC model with the message sample as the joint source-channel encoding process, and then transmit the output of the STSC model through the uplink channel to the receiver;
	
  * The transmitted message will then go through a noisy channel which typically add noisy corruption to the message. Once the receiver receives the message, it will then use the decoding part of the STSC model to recover the message. Note that the receiver in this setting could be either a cloud server, or a client.




Obviously, semantic information extraction and reconstruction of images are considered to be the key to the success of this system. Therefore, in this paper, the average mean square error (MSE) between the original input image and the reconstructed image is used as the loss function, which is defined as,


    MSE=1/L∑_i=1^n(x_i-x̂_i)^2,

where x_i is the original image vector, x̂_i is the reconstructed image vector, and L  represents the length of the image vector. A smaller MSE indicates that the reconstructed image is closer to the original image and the quality of image reconstruction is better.




§ IMAGE SEMANTIC ENCODER AND DECODER NETWORK


In this section, we first present an edge architecture for Swin Transformer-based semantic communication (STSC), which aims to capture long-term hierarchical semantic information from images. Additionally, to enhance the accuracy of semantic information extraction and ensure user privacy, we introduce the federated learning strategy for a Swin Transformer-based semantic communication (FSSC), which enables the training of edge models using FL across multiple devices hosted on cloud servers. As a result, our proposed scheme can effectively learn semantic representations from images transmitted by different users, and ensure reliable communication on various channels.



 §.§ STSC Model

Fig. <ref> shows the semantic communication network structure based on Swin Transformer. In this paper, joint source channel coding is adopted for semantic communication, in which the Swin Transformer module  <cit.> is adopted for encoding and decoding. Meanwhile, a non-trainable fully connected layer is used to simulate the physical channel.

At the transmitter, a set of N images 𝐗={x_i}_i=1^N is given, where x_i ∈ℝ^3 × H × W denotes the i-th image, and H and W denote height and width of the image, respectively. Firstly, the images are divided into non-overlapping patches through patch partion module. The size of patch is 4 × 4  and the vector dimension is converted to (H / 4,  W / 4,48). 
Then, the feature block sequence is linearly embedded. And through the learnable embedding matrix E, the feature blocks can be projected into an  embedding representation of arbitrary dimension C. After passing through the Swin Transformer block, the vector dimension is transformed to  (H / 4,  W / 4,C), 
and the whole process is called Stage 1. After the embedded representation, the feature blocks are fed into the Swin Transformer module, as shown in Fig. <ref>. As can be seen from the figure, the first patch merging layer and Swin Transformer block are combined into Stage 2. The patch merging layer concatenates each group of adjacent patches of size 2 × 2, so that the number of patch tokens becomes 1 / 4 of the original, i.e., H / 8×  W / 8. Meanwhile, the dimension of patch token is expanded by four times, i.e., 4C. In order to reduce the output dimension and realize the downsampling of the feature map, the patch merging layer performs a fully connected operation (implemented by a 1 × 1  convolutional layer). After this process, the dimension of the concatenated feature patch is reduced from 4C to 2C. Then, the patch goes through the swin Transformer block for feature transformation. Finally, the output dimension becomes  (H / 8,  W / 8,2C). The output x̃ passes through a fully connected layer to facilitate transmission in the channel. Specifically, it can be expressed as,

    s=W_1 x̃+b_1,

where, W_1 is the weight parameter matrix of the full connection, and  b_1 is the bias of the full connection.

In order to realize the joint training of encoder and decoder, a layer of fully connected neural network is used to simulate the channel. The neural network is actually a pair of input-output mappings, where the mapping relationship is determined by the neuron weight W_n and the bias b_1. Here, the weight W_n represents the channel gain, and the bias b_n is a random variable added to simulate the noise in the channel. The variance of the variable corresponds to the power of the channel noise, and its value mainly depends on the SNR and transmission power. After transmission in the physical channel, the signal s receives interference from channel noise and becomes the signal y when it reaches the receiver. 

At the receiver, the decoder consists of the Swin Transformer blocks. The function is to decode the signal y to restore the signal x̂.




 §.§ FL Training

In this paper, we develop the FSSC algorithm, by applying STSC to a FL framework, which can extract semantic information accurately and carry out semantic communication while protecting user privacy.While Fig. <ref> depics the overal architecture of FSSC, we detail below the specific training process of FSSC algorithm.

We use Fedavg<cit.> as the aggregation algorithm for our federated training. Assume that N clients participate in FL training. D_k  denotes the local Dataset of client E_k, and |D_k| is the size of the corresponding dataset. Thus, the training loss of the client based on the model in the local sample is

    Loss^k(ω)=1/|D_k|∑_i=1^|D_k|Loss_i^k(ω),

where  Loss^k(·) is the training Loss of the local client E_k, and Loss_i^k(·) is the loss of the sample x_i corresponding to the client. In the FSSC algorithm, Loss(·) is represented in (<ref>).

Therefore, the global federated training loss can be obtained by weighted average of the loss of each client according to the dataset size. The specific formula is as follows.


    Loss^∘(ω)=∑_k=1^N |D_k|/|D_o|Loss^k(ω),

where Loss^∘(ω) is the global loss after cloud aggregation, and |D_o| is the global dataset size. The purpose of the FL training is to find the global optimal model that minimizes the sum of training losses over all data. The specific federated training process is as follows.

Initialization. The dataset is non-uniformly split and distributed on each client. The semantic communication network based on Swin Transformer is initialized on the cloud server, and the initialized global model parameters W^o are broadcast to each client to build a local STSC network.

Model training. In the communication round t, the client participating in the training will conduct local training on STSC by using the local data set. And the stochastic gradient descent method is used to update the local model parameters. For each batch i, the process can be expressed as,


    ω_i+1^k ←ω_i^k-η∂Loss(ω_i^k ; b)/∂ω_i^k,

where  η is the learning rate. Each client saves its own trained local parameters and uploads them to the cloud server.

Parameter aggregation. The cloud server aggregates model parameters from different clients, and updates W_t^o to W_t+1^o . The aggregation mechanism shown in the following equation. Through this process, the update of the global shared model can be completed.


    w_t+1←∑_k=1^K |D_k|/|D_o| w_t+1^k,

where K is the total number of clients participating in training. This aggregation mechanism provides a certain guarantee for the stability of the system from the perspective of model convergence. Specifically, a certain degree of client-side fluctuations can be accepted without affecting the final convergence of the model.

Model convergence. The parameters of the global model W_t+1^o  are downloaded to each client. Then, the next round of training begins. Repeat this step until a specified global training cycle is reached or the model converges.



§ EXPERIMENT AND NUMERICAL RESULTS





 §.§ Simulation Settings

In this paper, we use the CIFAR-10 dataset<cit.> to evaluate the performance of the proposed FSSC algorithm. This dataset contains 50,000 training images and 10,000 test images, each with a size of 32 × 32. In addition, the number of clients is set to 3, and each client is assigned a different amount of non-overlapping training data. All clients use the same set of validation and test images. The training and testing environment are Windows 11+ CUDA 12.1, and the deep learning framework is Pytorch2.2.1.



In the experiments, the process of downloading and uploading model parameters between the client and the cloud server is regarded as an ideal situation. The focus is on the effects of global models obtained through local training and federated aggregation. The following introduces the initialization parameter settings for the semantic communication network based on Swin Transformer. Firstly, the network hyper-parameter C is set to 32. The network training all uses MSE as the loss function, the number of communication round is set to 60, the batch size of the data is 64, the learning rate is 0.001, and the Adam optimizer<cit.> is used. In all the tests, the compression ratio (CR) of the images is set to 0.33. 

To demonstrate the superiority of our method, we choose the typical JSCC <cit.> and the traditional communication algorithm as the comparison algorithm. The traditional communication algorithm adopts JPEG <cit.> as the source coding method, the channel coding adopts LDPC <cit.>, and the modulation mode and the demodulation mode are Quadrature Amplitude Modulation(QAM). And The modulation order is set to 4. 



 §.§ Simulation Results

In order to evaluate the semantic representation of the FSSC model, the experiments in this section compared the performance results of the local model algorithm STSC, JSCC and the traditional communication algorithms on AWGN channels, and verified the changes of the PSNR of the three methods under different SNR conditions, as shown in Fig. <ref>. It can be seen from Fig. <ref> that the PSNR of STSC and JSCC gradually increase with the increase of SNR. The STSC algorithm outperformed JSCC in terms of PSNR, with an average improvement of over 5 decibels. This significant advantage demonstrates the superior efficiency of the STSC algorithm in semantic feature extraction and reconstructing image data, resulting in a more accurate representation of the original image. At the same time, it can be noted that the PSNR of JPEG+LDPC+QAM remains unchanged when the channel conditions are very bad or greatly improved. This is because when the SNR is low, the traditional communication algorithm basically cannot transmit any semantic information. However, when the SNR is high, the PSNR reaches the performance saturation of the traditional communication algorithm. Thus, the accuracy of the proposed algorithm in this paper is significantly better than that of the traditional communication algorithms, especially when the SNR is improved, the performance improvement is very obvious. This is because while LDPC codes and QAM modulation can improve the robustness of data transmission, the JPEG and BPG compression algorithm are lossy and can lead to irreversible information loss. Combining the JPEG, LDPC and QAM techniques may reduce the error resilience of the overall system, leading to degraded image quality or data integrity in the presence of transmission errors. In comparison, the STSC algorithm retains the semantic information to the greatest extent, which ensures the effective transmission of images.



In order to evaluate the convergence of the proposed STSC model under the FL framework, the convergence analysis experiment of the FSSC model is also carried out in this section. Taking SNR=12 as an example, Fig. <ref> shows how the target loss varies as the training rounds/time increases in the case of FL training and local training, where the target loss is related to the degree of image reconstruction. As can be seen from the figure, on AWGN channels, the target loss of the communication model decreases rapidly in a short period of time, then fluctuates for a period of time, and finally remains unchanged. In other words, the model eventually reaches convergence. Moreover, the convergence efficiency of the FL training model is higher than that of the local training of the client. At the same time, the final training loss of global training model is relatively lower. This is because the local model uses a local single data set for local training, which makes it impossible to fully learn the global semantic distribution characteristics of the data. However, FL training has the advantages of data integration and improving model performance, and thus solves this problem well. 


Fig. <ref> shows how the PSNR change with the channel SNR for the global model and the local model without federated aggregation after model convergence. From the figure, it can be observed that with the increase of SNR, the PSNR of the global model and the local model without joint aggregation also significantly improved. In contrast to local models devoid of collaborative aggregation, global models can further enhance the PSNR by 2-3dB. This is attributed to the global model's ability to learn richer noise models and finer image semantic features through effective information aggregation and synchronization. Consequently, under high signal-to-noise ratio conditions, its PSNR can continue to improve, demonstrating superior generalization and noise resistance capabilities. This is because the image reconstruction is better as the channel condition improves. Moreover,  FSSC avoids the sharing of sensitive data, thus protecting the privacy of users and motivating users to participate in semantic communication under the FL framework to a certain extent.



§ CONCLUSION


In this paper, we proposed a semantic communication architecture for image transmission over wireless channels in multi-user cases, called FSSC. In this architecture, the local STSC semantic communication system model accurately extracts semantic information on the respective dataset and carries out the semantic communication model training. Then, joint training is performed through FL parameter aggregation to minimize the target loss function of the reconstructed image. Simulation results show that the algorithm can converge effectively. This approach is able to combine the semantic information exists in the diverse data from clients, and distributes the computational workload from edge to users, without compromising the user's privacy. Under the premise of privacy protection, the proposed algorithm can greatly reduce the amount of data required by each client while ensuring the fidelity in recovered images. Therefore, the FSSC proposed in this paper is a promising candidate scheme for multi-user image semantic communication system. 

00

ref1
K. Lu, Q. Zhou, R. Li, Z. Zhao, X. Chen, J. Wu, and H. Zhang, “Rethinking modern communication from semantic coding to semantic communication,” IEEE Wireless Communications, vol. 30, no. 1, pp. 158-164, February 2023.

ref2
E. Bourtsoulatze, D. Burth Kurka and D. Gündüz, “Deep Joint Source-Channel Coding for Wireless Image Transmission," IEEE Transactions on Cognitive Communications and Networking, vol. 5, no. 3, pp. 567-579, September 2019.

ref3
N. Farsad, M. Rao, and A. Goldsmith, “Deep learning for joint source-channel coding of text,” in 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), Calgary, AB, Canada, April 2018. 

ref4
H. Xie, Z. Qin, G. Y. Li, and B. H. Juang, “Deep learning enabled semantic communication systems” IEEE Transactions on Signal Processing, vol. 69, pp. 2663–2675, April 2021. 

ref5
K. Choi, K. Tatwawadi, A. Grover, T. Weissman, and S. Ermon, “Neural joint source-channel coding,” in Proceedings of the 36th International Conference on Machine Learning (ICML), Long Beach, CA, United states, May 2019. 

ref6
J. Xu, B. Ai, W. Chen, A. Yang, P. Sun, and M. Rodrigues, “Wireless image transmission using deep source channel coding with attention modules,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 32, no. 4, pp. 2315-2328, April 2022. 

ref7
Claude E Shannon, “A mathematical theory of communication,” The Bell system technical journal, vol. 27, no. 3, pp. 379–423, July 1948.

ref8
D. Gündüz, Z. Qin, I. E. Aguerri, H. S. Dhillon, Z. Yang, A. Yener, K. K. Wong, and C.-B. Chae, “Beyond transmitting bits: Context, semantics, and task-oriented communications,” IEEE Journal on Selected Areas in Communications, vol. 41, no. 1, pp. 5–41, November 2022.

ref9
M. Kountouris and N. Pappas, “Semantics-empowered communication for networked intelligent systems,” IEEE Communications Magazine, vol. 59, no. 6, pp. 96–102, June 2021.

ref10
W. Yang et al., “Semantic Communications for Future Internet: Fundamentals, Applications, and Challenges,” IEEE Communications Surveys & Tutorials, vol. 25, no. 1, pp. 213-250, November 2023.

ref11
Z. Liu et al., “Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,” in IEEE/CVF International Conference on Computer Vision (ICCV), Montreal, QC, Canada, October 2021.

ref12
K. Yang, S. Wang, J. Dai, K. Tan, K. Niu and P. Zhang, “WITT: A Wireless Image Transmission Transformer for Semantic Communications,” in IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Rhodes Island, Greece, June 2023.

ref13
H.B. Mcmahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” in International Conference on Artificial Intelligence and Statistics (AISTATS), Fort Lauderdale, Florida, USA, February 2017.

ref14
A. Krizhevsky, “Learning multiple layers of features from tiny images,” Master's thesis, University of Tront, April 2009.

ref15
D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” arXiv preprint arXiv:1412.6980, December 2014.


ref16
G. K. Wallace, “The JPEG still picture compression standard,” IEEE Transactions on Consumer Electronics, vol. 38, no. 1, pp.xviii-xxxiv, February 1992.

ref17
R. Gallager, “Low-density parity-check codes,” IRE Transactions on Information Theory, vol. 8, no. 1, pp. 21–28, January 1962.








 Paper Title: Regional quality estimation for echocardiography using deep learning.



Generalised BBGKY hierarchy for near-integrable dynamics
    Jacopo De Nardis
    Submitted August 5, 2024
========================================================

§ ABSTRACT 
Automatic estimation of cardiac ultrasound image quality can be beneficial for guiding operators and ensuring the accuracy of clinical measurements. Previous work often fails to distinguish the view correctness of the echocardiogram from the image quality. Additionally, previous studies only provide a global image quality value, which limits their practical utility. In this work, we developed and compared three methods to estimate image quality: 1) classic pixel-based metrics like the generalized contrast-to-noise ratio (gCNR) on myocardial segments as region of interest and left ventricle lumen as background, obtained using a U-Net segmentation 2) local image coherence derived from a U-Net model that predicts coherence from B-Mode images 3) a deep convolutional network that predicts the quality of each region directly in an end-to-end fashion. We evaluate each method against manual regional image quality annotations by three experienced cardiologists. The results indicate poor performance of the gCNR metric, with Spearman correlation to the annotations of ρ=0.24. The end-to-end learning model obtains the best result, ρ=0.69, comparable to the inter-observer correlation, ρ=0.63. Finally, the coherence-based method, with ρ=0.58, outperformed the classical metrics and is more generic than the end-to-end approach.


    Keywords:


Cardiac segmentation, Ultrasound, image quality, Coherence, Signal-To-Noise Ratio






§ INTRODUCTION


Image quality is one of the main challenges in ultrasound imaging and can differ significantly between patients and imaging equipment. In echocardiography, many factors influence image quality such as the ultrasound scanner, the patient, and the probe. Several quantitative measurements using the images are performed. However, this requires image quality good enough for the given measurement. Different measurements have different image quality requirements, for instance, left ventricular (LV) volume, ejection fraction (EF), and strain measurements require good image quality in the entire myocardium. On the other hand, mitral annular plane systolic excursion (MAPSE) only requires good image quality in the annulus. Good image quality should generally provide measurement values with low uncertainty. Estimating image quality can be useful in the following ways:

    
  * To guide operators to achieve as good image quality as possible while scanning.
    
  * To automatically select the best images, recordings, and the best cardiac cycles to use for a given measurement.
    
  * As quality assurance, e.g. to warn the user when an image is not good enough for a measurement, and to automatically approve/disapprove individual myocardial segments based on image quality.
    
  * In data mining projects, to exclude cases with insufficient quality for reliable measurements.


We distinguish between two types of quality of ultrasound images: view quality/correctness and image quality. 
In this work, we will focus on image quality specifically. For view correctness, previous work has demonstrated that 3D ultrasound can serve as training data to automatically identify the transducer rotation and tilt in relation to the desired standard view and can guide the user to the correct position <cit.>. 
For image quality, the classic ultrasound signal-processing metrics are the contrast ratio (CR) <cit.>, contrast-to-noise ratio (CNR) <cit.>, and generalized CNR (gCNR) <cit.>. These three metrics need a region of interest (ROI) and a background region to compare against. More recently, global image coherence (GIC) <cit.> has been proposed as a general quality metric that does not require the selection of these two regions. The image coherence measures how well the signals of the transducer elements align after delay compensation, with more alignment corresponding to clearer and sharper images. From the above mentioned methods, only the GIC can be used directly and automatically for measuring image quality separately as it does not require selecting an ROI and noise region. However, this approach requires channel data, which is not readily available in practice and does not give regional metrics.  

Several automatic methods for measuring ultrasound image quality have been published applicable to cardiac imaging. Abdi et al. <cit.> used a recurrent neural network to predict the global quality of cardiac cine loops. The criteria for quality assessment take both image quality and view correctness into account. In subsequent studies <cit.>, they used an architecture that performs both view classification and global quality prediction simultaneously. The image quality metric is a global criterion based on the manual judgment of the clarity of the blood-tissue interface. Labs et al. <cit.> used a multi-stream neural network architecture where each stream takes in a sequence of frames and predicts a specific quality criterion. The criteria are global and take both view correctness and image quality into account. Karamalis et al. <cit.> detect attenuated shadow regions with random walks resulting in a pixel-level confidence map. Unlike the other methods above, this method is not based on deep learning. It provides a local, pixel-level metric, but it only measures the visibility of regions and not the quality of their content. 



All of the automatic methods mentioned above have the limitation that they only provide a global image quality evaluation and/or do not assess the image quality separate from the view correctness.  




§ METHODS

In this work, we developed and compared three fully automatic methods to asses regional image quality in cardiac ultrasound separate from the view correctness:

    
  * Classical ultrasound image quality metrics, such as contrast ratio and contrast to noise ratio, applied in cardiac regions automatically extracted using deep-learning segmentation.
    
  * Deep-learning predicted ultrasound coherence, which is a measure of how coherent a signal is received by the transducer elements, together with deep-learning segmentation.
    
  * End-to-end prediction of regional image quality.


The rest of this section first presents the datasets used to develop these methods, and then presents each of the three methods.



 §.§ Datasets




  §.§.§ VLCD


The Very Large Cardiac Channel Data Database (VLCD) consists of channel data from 33280 frames from 538 recordings of 106 study participants <cit.>. It contains parasternal short axis (PSAX), parasternal long axis (PLAX), apical long axis (ALAX), apical two-chamber (A2C), and apical four-chamber (A4C) views. We split the VLCD dataset on the study participant level into train, validation, and test sets, 70%, 15%, and 15% respectively. 



  §.§.§ HUNT4


The Nord-Trøndelag Health Study dataset (HUNT4Echo) is a clinical dataset including among others PSAX, PLAX, ALAX, A2C, and A4C views. Each recording contains 3 cardiac cycles. We use two subsets of the HUNT4Echo dataset.

    
  * Segmentation annotation dataset A fraction of 311 study participant exams, the segmentation annotation set <cit.>, contains single frame segmentation annotations in both ED and ES as pixel-wise labels of the left ventricle (LV), left atrium (LA), and myocardium (MYO) in ALAX, A2C, and A4C views.
    
  * Regional image quality dataset For this work, we created an additional dataset of image quality labels. The local image quality labels are manual annotations that asses the image quality of the cardiac regions of interest on a subset of the HUNT4 dataset in ALAX, A2C, and A4C views. Section <ref> describes the annotation process in more detail.












 §.§ Regional image quality annotation on HUNT4
  
An annotation tool was developed specifically for this project using the open-source Annotation Web software[<https://github.com/smistad/annotationweb>] <cit.>.
The tool was made to enable clinicians to annotate regional image quality as efficiently and accurately as possible. The tool is freely available and can be adapted to other image quality projects. The image quality annotations were performed by three experienced cardiologists using the following protocol:


  * Annotate the end-diastole (ED) and end-systole (ES) frame of each recording, and optionally other frames if the image quality changes significantly during the recording.

  * If the majority of the cardiac regions of interest is out-of-sector, label it as out-of-sector. Otherwise, label the part of the region that is inside the sector according to the definitions in Table <ref>. We ignore the out-of-sector regions in the remainder of this work.

For the first round of annotations, each of the three clinicians annotated the same 10 frames from 5 recordings of 2 study participants. We used this dataset to calculate the inter-observer variability. For the second round of annotations, the three clinicians collectively annotated 458 frames from 158 recordings of 65 study participants. The annotations from the second round form the regional image quality dataset. This dataset was split randomly at the study participant level into train, validation, and test sets, allocating 70%, 15%, and 15% of the data to each set respectively.








 §.§ Regional image quality estimation









  §.§.§ Classical ultrasound image quality metrics

For the classical image quality metrics, deep-learning segmentation is used to extract the annulus regions and each of the myocardial segments as regions of interest and the LV as the background region. Appendix <ref> gives more details about the procedure for dividing the segmentation into regions. The four classical ultrasound image quality metrics below were tested.
We apply histogram matching <cit.> to a Gaussian distribution (μ=127 , σ=32) for the B-Mode grayscale images before calculating pixel-based quality metrics.




  
  • Pixel intensity is the average pixel intensity value in each region.


  
  • Contrast Ratio (CR) <cit.> is defined as 
  
    CR=μ_segment/μ_LV

  where μ_segment is the average intensity in each region and μ_LV is the average intensity inside the LV lumen.
  
  
  • Contrast to Noise Ratio (CNR) <cit.> is defined as 
    
    CNR=μ_segment-μ_LV/√(σ_segment^2+σ_LV^2)

    where σ_segment is the standard deviation in each region and σ_LV is the standard deviation inside the LV lumen.

  
  • Generalized CNR (gCNR) <cit.> is defined as the maximum performance that can be expected from a hypothetical pixel classifier based on intensity using a set of optimal thresholds. It is calculated as 

    gCNR= 1-1/2∑_i=0^MAX_imin{p_segment(i), p_LV(i)}

where p_segment(x) is the probability density function of the pixel intensities inside the region the gCNR is calculated for, p_LV(x) the probability density function of the pixel intensities inside the LV lumen, and MAX_i the maximum possible pixel intensity. 
 Fig. <ref> shows an example of the probability density functions for one of the regions as ROI and the LV lumen as background.








  §.§.§ Local, deep-learning predicted image coherence as quality metric
 

We use the VLCD dataset to calculate the coherence factor <cit.> for each pixel in the ultrasound image. This factor is the ratio between the amplitude of the sum of the received signals to the sum of the amplitudes of those signals, 

    CF = ∑_n=1^NS_i/∑_n=1^N|S_i|

where S_i is the delayed signal for the i-th transducer element. This is equivalent to taking the coherent sum of the signal and dividing it by the incoherent sum of each signal. Thus, the coherence factor measures of how well the complex signals of all transducer elements align. The remainder of the signal-processing chain is the native processing of the GE HealthCare Vivid E95 system[ Gundersen et al. <cit.> describe this signal-processing chain in more detail.] but without the log compression. The result is a coherence image with the same dimension as the B-Mode image.  The final preprocessing step applies gamma normalization with γ=0.5 on the coherence images,

    t_i,normalized=t_i^γ

where t_i are the pixels of the target coherence image. The corresponding B-mode images are generated from the channel data using the same, native signal-processing pipeline. 



The HUNT4 dataset, like most ultrasound datasets, does not include channel data. Therefore, VLCD was used to train an image-to-image network that takes as input the grayscale B-mode image and predicts the coherence image, which is then used to calculate local image coherence. We use a lightweight U-Net architecture inspired by the U-Net 1 architecture in <cit.>, with characteristics listed in Table <ref>. As coherence is related to image quality, we only apply augmentations that do not influence the quality of the image. Furthermore, the coherence should be invariant to different gain settings, so we additionally augment with brightness adjustments on the B-mode image while keeping the target coherence image unchanged. During training and validation of the coherence prediction model, we sample a random frame from each recording in the train and validation set respectively during each epoch. During testing, we use all frames in the test set. The local image coherence quality metric of a region is the average pixel value of all pixels corresponding to the region in the coherence image. This is the same as the pixel intensity metric above but applied to the coherence image instead of the B-Mode image.









  §.§.§ End-to-end deep-learning quality prediction
 

The end-to-end learning approach trains a convolutional neural network on the regional image quality dataset to predict the quality of each region directly. We treat the problem as a regression task where the model predicts a score for each segment. Table <ref> shows the correspondence between quality scores and annotation labels. As architecture, we use a MobileNetV2 <cit.> that predicts the image quality labels of all regions simultaneously. Appendix <ref> describes the ablation study conducted to justify this setup. 








§ EXPERIMENTAL SETUP





 §.§ Evaluation of coherence prediction from B-mode


We use the structural similarity index (SSIM) <cit.>, peak signal-to-noise ratio (PSNR)[The maximum pixel value for coherence images is 1.], and relative pixel error (RPE) to evaluate the coherence image prediction network. We define the RPE as 

    RPE = |t_i-p_i|/max(t_i,ϵ)

where t_i are the pixel values of the target coherence image, p_i the pixel values of the predicted coherence image, and ϵ=1e-4. Section <ref> contains the results of this experiment.



 §.§ Evaluation of quality metrics



The correlation and accuracy of each quality metric were measured by comparing them to the expert annotations on the test set of the regional image quality dataset. For the classic image quality methods and regional coherence method, linear regression models were used to map the quality metric values to image quality labels. The train and validation set were used together to fit the linear regression model and evaluate it on the test set. Section <ref> contains the results of this experiment.



 §.§ Comparison to inter-observer variability




We compare the end-to-end, local image coherence, and gCNR-based model to the inter-observer variability on the data obtained in the first round of annotation, explained in subsection <ref>. For the local image coherence and gCNR-based models, the same linear model as subsection <ref> was used to map the metrics to quality labels. The inter-observer variability is calculated from the aggregate of the three unique pairwise score errors between each of the three annotators:

    e_inter-observer = e_12∪ e_23∪ e_13

where e_ij is the score difference between operator i and j. The error metrics of the automatic methods are calculated from the aggregate of the pairwise score errors between the output of the method and each of the three annotators:

    e_M = e_1M∪ e_2M∪ e_3M

where e_iM is the score difference between operator i and method M. 
Section <ref> contains the results of this experiment.






 §.§ Relation to variability in clinical measurements


This experiment evaluates whether there is a relation between the predicted quality and the agreement between different methods for clinical measurements. The hypothesis is that with lower image quality the variability, and thus the uncertainty, of the measurements between methods and between experts increases. More specifically,
this analysis compares peak global longitudinal strain (GLS) and ejection fraction (EF) measurements obtained either fully automatically with AI tools or manually by using GE HealthCare EchoPAC software on HUNT4 <cit.>. For AI estimation of GLS and EF, the deep-learning methods proposed by Østvik et al. <cit.> and
Smistad et al. <cit.> were used respectively. The study participants in HUNT4 used for model development were excluded from the analysis. For GLS, the predicted quality is the average quality of all segments over the full recording.
For EF, the predicted quality is the average quality of all segments in the end-diastole (ED) and end-systole (ES) frames of all cycles in the recording. Section <ref> contains the results of this experiment.












§ RESULTS






 §.§ Results of coherence prediction from B-mode
 

Table <ref> summarizes the average metric values on the test set. Fig. <ref> shows an example of the best, median, and worst-case predictions according to the relative pixel error. Fig. <ref> illustrates how the predicted coherence images are almost independent of the brightness/gain and contrast/dynamic range of the input B-Mode images. The main finding is that the difference between the estimated and ground truth coherence images is small and thus the predicted coherence images can be used to obtain coherence-based quality metrics for B-mode for which the channel data is not available. 























 §.§ Results of quality metrics
 

Table <ref> summarizes the results of the evaluation of the quality metrics. Fig. <ref> shows box plots of the quality metrics per image quality label for the end-to-end, coherence, gCNR, and intensity models. Fig. <ref> shows examples of B-mode images with varying quality together with labels from the annotators and automatic quality metrics. The main finding is that the end-to-end model performs the best, followed by the local image coherence metric. The classical ultrasound image quality metrics perform poorly.
















 §.§ Results of comparison to inter-observer variability
 

Fig. <ref> shows the bar plot comparing the automatic methods to the inter-observer variability and Table <ref> lists the corresponding average metric values. Using the Wilcoxon signed-rank test <cit.> and a significance level of p=0.05, we find that the difference in Mean absolute error (MAE) between each of the methods is statistically significant. The difference between the inter-observer MAE and the MAE of each of the methods is also statistically significant, i.e. the inter-observer MAE is higher than the MAE of the end-to-end model and lower than the MAE of the other two models.












 §.§ Results of relation to variability in clinical measurements
 

Fig. <ref> shows box plots visualizing the agreement between the measurements obtained automatically and with EchoPAC for each predicted quality category. The standard deviations in these plots represent how well the AI estimates agree with the manual references. The main finding is that the limits of agreement are narrower for higher qualities.












§ DISCUSSION




 §.§ Challenges and considerations



Assessing image quality based on human perception is inherently a subjective task, even when supported by clear definitions of image quality categories. This creates challenges for training and evaluation as there is no ground truth as the reference labels are a subjective estimation themselves. Therefore, it is not realistic to expect the automatic models to agree with reference labels as well as on tasks with well-defined ground truth labels. This, together with a rather fine scale of image quality categories, explains the low accuracies in Tables <ref> and <ref>. Fig. <ref> shows how the end-to-end model has on average less error than the annotators between each other. This indicates the model has learned to produce quality labels that strike a middle ground between the subjective assessments of the annotators. 

The end-to-end learning model overestimates low-quality regions and underestimates high-quality regions, as can be seen in Fig. <ref>. This means that the model can only explain a limited amount of variability in the image quality labels and is a result of minimizing the mean squared error (MSE) while dealing with subjective, and thus noisy reference values with fixed boundaries. We can eliminate this effect by fitting a linear model on the validation set that maps predicted image quality to image quality labels and applying it when doing inference on the test set. This increases MSE but gives more uniform performance over the image quality labels.

One reason for the weak correlation between the annotations and the pixel-based methods is the rough selection of ROI and background region. Fig. <ref> shows how the average metrics of the classic pixel-based and coherence metrics increase for each quality label until the good label, and then drop again for excellent.  This is because on the one hand in these high-quality images, the blood speckles can be visible inside the LV lumen, which is used as background region, and on the other hand the myocardium tissue, which is used as ROI, is less blurred resulting in a smaller spread of pixels with high intensity. This can be seen for the anterolateral wall and apex in the rightmost column of Fig. <ref>. One possibility to only select regions belonging to the tissue is to perform automatic pixel selection methods like Otsu thresholding <cit.> or percentile filters, but in our experiments this reduced the performance even further. 





 §.§ Design choices


The different methods in this study have a trade-off between accuracy and versatility. The default end-to-end network gives the best results but requires specific image quality labels for the task. Next, the coherence-based method is more generic and can potentially be applied more generally without the need for view-specific image quality annotations. Rindal et al. <cit.> showed that the GIC is not significantly different between apical views, but is higher for apical views than parasternal views. Thus, while a single image-to-image model can learn to predict coherence for different views, the mapping from coherence to image quality should be done for each group of views separately. Another advantage is that coherence can be used to give a global image quality metric without the need for a segmentation model. Finally, the pixel-based methods can be applied automatically in the most general way given a segmentation model to select ROI and background regions but also give the lowest accuracy. 

The ablation study of the end-to-end learning model showed that increasing the complexity of the model did not improve the performance. This is a result of the relatively small dataset size and the specific task of the model. For a more general model of image quality prediction with more varied input, e.g. one model for all cardiac views, a larger dataset and more complex model may be required.



 §.§ Clinical use


Image quality estimation can be the first step towards a method for giving reliability estimates to clinical measurements and quality control of fully automatic methods. Fig. <ref> shows that the variability in clinical measurements goes down with higher predicted quality. However, image quality is only one source of variability, so a reliability model would also need to include view correctness and other factors that determine whether a given input is difficult to assess. 

More direct use cases of the quality prediction model include the automatic selection of the best frame to perform a clinical measurement when multiple options are available, data cleansing in data mining, and automatic disapproval of segments for regional strain analysis. All the methods explored in this work are computationally efficient and can be run in real-time while scanning, and can thus be used as a guidance tool to enable clinicians to acquire images with better image quality.




§ REAL-TIME DEMO APPLICATION


To showcase the functionality of the end-to-end, real-time quality network, a real-time application was developed using the FAST framework <cit.>. The demo is a split-screen application that shows the B-Mode input to the left and the segmentation regionally color-coded by the quality as predicted by the end-to-end network to the right. Fig. <ref> provides a screenshot of the application in use. We provide a demo video <cit.> illustrating the application in action while a clinician operates a GE Vivid E95 scanner. The video can be accessed at <https://doi.org/10.6084/m9.figshare.26413984.v2>.






§ CONCLUSION


In this work, we developed and compared different deep-learning methods for regional image quality estimation in cardiac ultrasound. We show that classic pixel-based methods, such as (g)CNR, together with automatic image segmentation, give low agreement with the quality assessment of cardiologists. We developed a U-Net model to predict the coherence factor for each pixel in the ultrasound image and showed that the resulting coherence image can be used to assess the image quality in a pixel-based way with better performance than the classic measures. The best results, below inter-observer variability, are obtained by using an end-to-end deep-learning model. Finally, we show higher predicted quality is associated with lower limits of agreement between fully automatic and manual methods for the clinical measurements EF and GLS. 





IEEEtran







§ EXTRACTION OF CARDIAC REGIONS OF INTEREST
 


We use the nnU-Net <cit.> architecture to segment the cardiac images. The nnU-Net is used out of the box using the default configuration but without the final ensemble step. Instead, we train and validate on a single predefined 80% train, 10% validation, and 10% test split from the HUNT4 segmentation annotation dataset. Table <ref> summarizes the characteristics of the nnU-Net architecture. This model is described in more detail and compared to other segmentation models in our previous work <cit.>. We use two nnU-Nets, one for apical two- (A2C) and four-chamber (A4C) views, and one for apical long axis (ALAX) views. The nnU-Net for A2C and A4C views segments the left ventricle (LV), left atrium (LA) and myocardium (MYO). The nnU-Net for ALAX views additionally segments the aorta (AO).






    


The segmentation of the MYO is divided into eight regions using the following algorithm:

  
  * Extract the annulus points. For A2C/A4C views, these are the points where the MYO meets the LA. For ALAX views, these are the points where the MYO meets the LA and AO. Points A and B are the annulus points in Fig. <ref>.
  
  * Extract the apex of the LV, defined as the furthest points from the base points within the lv lumen. This is point C in Fig. <ref>.
  
  * Divide both the left and right part of the endocardium border, defined as the border between the LV and MYO regions, into three parts with equal length. This gives points D, E, F and G in Fig. <ref>.
  
  * Find the closest points on the outer MYO border for points C, D, E, F and G. These are points H, I, J, K and L in Fig. <ref>.
  
  * Fill in the regions by connecting the points via the contour, resulting in the MYO divided into six regions. 
  
  * Draw circles[Due to the unequal pixel spacing in depth and width, the annulus regions become ovals in the 256x256 segmentation maps. When plotting the images with equal spacing in width and depth, these regions become circles again.] with a radius of 2  millimeters around the annulus points, points A and B in Fig. <ref>. 
  We use these additional two regions to asses the local image quality of the annulus points in the image. The result are the eight regions, as in Fig. <ref>.
  
  * Remove any parts of regions that fall outside of the sector. The apical top regions in Fig.<ref>, i.e., the yellow and white masks, are examples of this. If more than 50% of all pixels inside the region fall outside the sector, we exclude the region from analysis.

The goal is to automatically quantify the image quality in each of these eight regions. The LV lumen is used as background region.








§ ABLATION STUDY OF THE END-TO-END LEARNING MODEL



In the ablation study, we evaluated the impact of modifying the architecture of the end-to-end learning model. Three different network architectures were tested: Cardiac View Classification (CVC) network <cit.>, MobileNetV2 <cit.>, and EfficientNet <cit.>. 
We tested approaching the problem both as a classification and regression task, with only the final dense layer and loss function being changed accordingly. Additionally, three basic network attention variations were tested using the automatic segmentation output. The default model did not use any attention and predicts each label directly, as in Fig.<ref>. For the other two versions, the region masks extracted from the segmentation were dilated with a square dilation filter of size 50x50 pixels and used as an additional input to the networks which then predict the label of one region at a time. The dilation filter reveals the direct vicinity around each region so the boundary between tissue and background becomes visible. The first variant used this dilated mask as hard attention by blacking out the other parts of the image, as shown in Fig. <ref>. For the second variant, the masks were used as soft attention as input to a side branch of the network, as proposed by Eppel <cit.>. For this version, we created an attention map and added it element-wise to the output of the first layer, corresponding to version 'c' in Eppel <cit.>. Fig. <ref> shows this configuration. 






The ablation study consists of two parts. Both parts used the training configuration listed in Table <ref> and data from the regional image quality dataset. In the first part, we examined the effect of changing the convolutional backbone and the effect of framing the problem as a classification or regression task. Table <ref> compares the predictions with the annotations on the test set for the different configurations using the default end-to-end model. In the second part of the ablation study, the backbone was fixed to MobileNetV2 <cit.> and the problem was framed as a regression task. Next, the different variants shown in Fig. <ref> were compared. Table <ref> summarizes the results on the test set. The default end-to-end model with no attention, the MobileNetV2 <cit.> architecture, and the problem framed as a regression task gave the best results and were thus used for this paper.









 Paper Title: Image Super-Resolution with Taylor Expansion Approximation and Large Field Reception.


























Shell et al.: Bare Demo of IEEEtran.cls for IEEE Journals
Image Super-Resolution with Taylor Expansion Approximation and Large Field Reception
    Jiancong Feng, Yuan-Gen Wang, Mingjie Li, and Fengchuang Xing




Feng, J., Wang, Y.-G., Li, M., and Xing F. are with the School of Computer Science and Cyber Engineering, Guangzhou University, Guangzhou 510006, China (e-mail: fengjiancong@e.gzhu.edu.cn; wangyg@gzhu.edu.cn; limingjie@gzhu.edu.cn; xfchuang@e.gzhu.edu.cn). 




    August 5, 2024
===========================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Self-similarity techniques are booming in blind super-resolution (SR) due to accurate estimation of the degradation types involved in low-resolution images. However, high-dimensional matrix multiplication within self-similarity computation prohibitively consumes massive computational costs. We find that the high-dimensional attention map is derived from the matrix multiplication between Query and Key, followed by a softmax function. This softmax makes the matrix multiplication between Query and Key inseparable, posing a great challenge in simplifying computational complexity. To address this issue, we first propose a second-order Taylor expansion approximation (STEA)  to separate the matrix multiplication of Query and Key, resulting in the complexity reduction from 𝒪(N^2) to 𝒪(N). Then, we design a multi-scale large field reception (MLFR) to compensate for the performance degradation caused by STEA. Finally, we apply these two core designs to laboratory and real-world scenarios by constructing LabNet and RealNet, respectively. Extensive experimental results tested on five synthetic datasets demonstrate that our LabNet sets a new benchmark in qualitative and quantitative evaluations. Tested on the RealWorld38 dataset, our RealNet achieves superior visual quality over existing methods. Ablation studies further verify the contributions of STEA and MLFR towards both LabNet and RealNet frameworks.





Blind Super-Resolution, Self-similarity, Taylor Expansion, High-dimensional matrix multiplication 










§ INTRODUCTION



Single-Image Super-Resolution (SISR) <cit.> has obtained widespread application in various fields, such as remote sensing imagery, security surveillance, and medical imaging. SISR methods are generally divided into two categories: non-blind SR and blind SR. The non-blind SR <cit.> assumes that the degradation type is known, which degradation function is typically represented as:

    𝐱_lr = 𝐱_hr↓_s,

where 𝐱_hr and 𝐱_lr denote the HR and LR images respectively, and ↓_s denotes the bicubic downsampling operation with a scale factor of s. However, the non-blind SR represented by Eq. (<ref>) fails to accurately approximate the real-world degradation due to the inaccessibility of the reference image.



Subsequently, researchers have developed blind SR methods, which assume that the LR image is obtained from the HR image through a degradation function with a complex and unknown degradation type, and this degradation function is typically assumed to follow the equation:

    𝐱_lr = (𝐱_hr⊗𝐊𝐞𝐫𝐧𝐞𝐥) ↓_s + N,

where ⊗ denotes the convolutional operation, 𝐊𝐞𝐫𝐧𝐞𝐥 represents the unknown blur kernel, ↓_s signifies the downsampling operation with a scale factor of s, and N accounts for sophisticated noise. According to Eq. (<ref>), we can clearly see that a single HR image corresponds to infinitely many possible LR images, which is typically an ill-posed problem. Nonetheless, a multitude of blind SR methods has been proposed to overcome this problem. They can be broadly divided into two main types: methods without blur kernel estimation <cit.> and methods with blur kernel estimation <cit.>.

As for methods without blur kernel estimation, SRMD <cit.> is the first to combine the preset blur kernel, noise, and LR image as the input to the SR network. After that, <cit.> incorporated Maximum A Posteriori Probability into iterative optimization. For these methods, the selection of parameters for the preset blur kernel is crucial for achieving effective results in practical applications. However, it is challenging to manually choose appropriate preset blur kernel parameters in the real world. To overcome this challenge, researchers have developed some methods to estimate the blur kernel. IKC <cit.> proposed to iteratively estimate the blur kernel, which divides the SR reconstruction process into two steps: 1) estimating the blur kernel from the LR image and 2) combining the estimated blur kernel and the LR image to reconstruct the SR image. However, the independence between the training of the blur kernel estimation model and the SR model leads to incompatibility issues. To address this issue, DANv1 <cit.> proposed an alternating optimization strategy between the blur kernel estimation and SR models, which makes the two models form a closed loop that can be trained in an end-to-end manner. Although DANv1 has achieved impressive performance in blind SR tasks, NLCUNet <cit.> has made significant advancements by incorporating the Non-Local Sparse Attention (NLSA) <cit.> module. It screens global features in the feature map by calculating self-similarities, enabling implicit estimation of the blur kernel. Therefore, the self-similarity technique has received the widest application in blind SR tasks. However, blind SR models that incorporate the calculation of self-similarity suffer from high computational complexity, making them unsuitable for resource-constrained devices. It is worth noting that the commonly used module for calculating self-similarity is Non-Local Attention (NLA) <cit.>.

In this paper, we go deep into the execution process of NLA. We find that during the self-similarity calculation, a matrix multiplication operation between the high-dimensional attention map and the Value is required, leading to high computational complexity. The attention map is obtained by multiplying the Query matrix with the Key matrix. Then the softmax function is further performed, which hinders the separation of Query and Key. Based on this finding, we propose a second-order Taylor expansion approximation (STEA) to significantly reduce the computational complexity. Compared to self-similarity calculations, STEA loses the ability to filter global features. Hence, we design a multi-scale large field reception (MLFR) to compensate for the performance loss caused by STEA. Finally, we apply our two core designs to laboratory and real-world scenarios by constructing LabNet and RealNet, respectively. Extensive experiments validate the effectiveness of the proposed method. The major contributions of this paper are summarized as follows.

    
  * We propose the second-order Taylor expansion approximation (STEA) strategy, which significantly reduces the computational complexity in self-similarity-based blind SR tasks.
    
  * We design a multi-scale large field reception (MLFR) to compensate for the performance loss caused by STEA, showcasing competitive performance over the state-of-the-art.
    
  * We construct the LabNet and the RealNet by integrating our two core designs (STEA and MLFR) to respectively verify the effectiveness of our designs in the laboratory and real-world scenarios. Extensive experimental results demonstrate that our LabNet achieves state-of-the-art performance while maintaining linear computational complexity, and our RealNet effectively handles various real-world degradations and achieves superior performance in diverse blind SR scenarios.






The structure of this paper is organized as follows: Section II provides a comprehensive overview of the proposed methodology. Specifically, Section II-A presents a detailed explanation of the intricate aspects of STEA. Section II-B elucidates the specific details of MLFR. Section II-C outlines the implementation specifics of LabNet, while Section II-D delves into the implementation details of RealNet. Subsequently, Section III encompasses the dataset and evaluation metrics (Section III-A), implementation details of LabNet (Section III-B), implementation details of RealNet (Section III-C), comparison with state-of-the-art techniques (Section III-D), ablation studies and analysis of LabNet (Section III-E), ablation study and analysis of RealNet (Section III-F) and adjustable parameter details in RealNet (Section III-G). Section IV concludes the paper.




§ METHODOLOGY

In this section, we provide a detailed explanation of the STEA and MLFR algorithms, as well as their deployment to the LabNet (testing for laboratory settings) and RealNet (testing for real-world scenarios).











 §.§ Second-Order Taylor Expansion Approximation (STEA)

The self-similarity has been widely applied to the blind SR task, which is commonly implemented by a Non-Local Attention (NLA) <cit.>. However, we found that the NLA brings extremely high computational complexity. This mainly arises from the matrix multiplication between the high-dimensional Attention Map and the V matrix, as highlighted by the red dashed box of Fig. <ref>. Without loss of generality, Q, K, V ∈ℛ^N× d denote query, key, and value matrix, respectively. Then, we know the computational complexity of NLA is 𝒪(N^2× d), where N and d represent the dimension of the feature map and the number of channels, respectively. In this part, we propose an STEA algorithm to significantly reduce the computational complexity of NLA. To this end, we first analyze the execution process of the NLA module. As shown in Fig. <ref>, the Attention Map is derived from the matrix multiplication of Q and K^T. Due to the limitation imposed by the softmax operation, the matrix multiplication between Q and K^T cannot be separated, resulting in a 𝒪(N^2) complexity. To address this challenge, we propose a Second-Order Taylor Expansion to approximate the NLA module. Here, we mathematically rewrite the execution process of the NLA as follows:

    Q = XW_q, K = XW_k, V = XW_v,

where W_q, W_k and W_v stand for the parameters of three 1× 1 convolutions, respectively. Q, K and V represent the Query, Key and Value, respectively. X denotes the input feature map. Next, we perform the matrix multiplication between Q and K^T, followed by the softmax operation, to obtain the Attention Map:

    AttMap = SoftMax(QK^T),

where AttMap represents the Attention Map, and SoftMax(·) denotes the softmax function. By substituting Eq. (<ref>) to Eq. (<ref>), we have

    AttMap = SoftMax(XW_qW_k^TX^T).

We denote W_qk = W_qW_k^T, resulting in the below expression:

    AttMap = SoftMax(XW_qkX^T).

In the original NLA module, the multiplication between Q and K^T followed by the softmax operation gives rise to softmax attention. Although softmax attention can effectively enhance the performance of the SR network by making the matrix full-rank, it introduces high computational complexity since such a high-dimensional Attention Map (AttMap) must be multiplied with V in the subsequent computation. For that, we propose to simplify the softmax function as an exponential form:

    AttMap = SoftMax(QK^T) = 1/k× e^QK^T = 1/k× e^XW_qkX^T,

where k is a learnable parameter used to approximate the denominator of softmax. By now, we can see that the operation between QK^T and V is still nonlinear due to the presence of an exponential function, which disables us from separating Q and K from QK^T. In subsequent multiplication by V, the basic operation of the matrix cannot be flexibly used, failing to reduce the high computational complexity. For this purpose, we exploit a second-order Taylor expansion as a substitute for the exponential function, achieving a good trade-off between



efficiency and accuracy. 

We know that in principle, the more remainder terms kept in the Taylor expansion of the exponential function, the more the expansion approximates the raw function. According to our empirical results and considering a good trade-off between accuracy and complexity, we adopt the second-order Taylor expansion on Eq. (<ref>) as follows:

    AttMap = 1/k× e^XW_qkX^T
     = 1/k× (I_n + XW_qkX^T + (XW_qkX^T)^2/2! + O(XW_qkX^T)),

where I_n stands for identity matrix and O(·) stands for the high-order remainder of the Taylor expansion. By ignoring the constant and remainder terms of Eq.(<ref>), we achieve:

    AttMap= 1/k× (I_n + XW_qkX^T + 
     XW_qkX^TXW_qkX^T).

Then, we multiply AttMap by V and refer to Eq.(<ref>) to obtain:

    Output = AttMap× V = 1/k× (XW_v + 
     XW_qkX^TXW_v + XW_qkX^TXW_qkX^TXW_v).

Eq. (<ref>) clearly indicates that the constraint of the softmax function has been removed, and therefore we can apply the associative property of the matrix operation in Eq. (<ref>), which allows us to calculate X^TX first. Considering that X^TX is a symmetric matrix, we can transform X^TX into a diagonalization expression as follows:  

    X^TX = ZBZ^T,

where B is a diagonal matrix with its diagonal elements being the eigenvalues of X^TX, and Z is an orthogonal matrix whose columns are the eigenvectors of X^TX corresponding to the eigenvalues. By substituting Eq. (<ref>) into Eq. (<ref>), we obtain:

    Output = 1/k× (XW_v + XW_qkZBZ^TW_v + 
     XW_qkZBZ^TW_qkZBZ^TW_v).

According to Eq. (<ref>), we can see that obtaining B through the matrix eigen-decomposition operation is computationally expensive. If we randomly initialize a learnable diagonal matrix B to avoid the high computational complexity from the eigen-decomposition operation of X^TX. Unfortunately, this strategy makes the matrix B independent of the input feature maps (X). Besides, in the random initialization, the matrix B is required to be a fixed size, making X have to be the fixed size. This disables the input LR image to be an arbitrary size. To reduce the complexity of the matrix eigen-decomposition and implement the SR task on the LR images with an arbitrary size, we design two separative eigenvalue extraction modules to respectively obtain two diagonal matrices B_1 and B_2, as shown in Fig. <ref>. Our design avoids randomly initializing a single B, and meanwhile enhances the representation ability of the model. Therefore, we can obtain a learning-based approximation expression as follows: 

    Output = 1/k× (XW_v + XW_qkZB_1Z^TW_v + 
     XW_qkZB_1Z^TW_qkZB_2Z^TW_v).

It is observed from Fig. <ref> that our design modules can dynamically extract global information from the matrices, which can well match the matrix eigen-decomposition where global information is exploited dynamically. 

Similarly, we adopt learnable parameters to substitute Z, solving the high computational complexity of the eigen-decomposition of X^TX. Next, we let W_1 = W_qkZ, W_2 = Z^TW_qkZ, and W_3 = Z^TW_v. Hence, we obtain:

    Output = 1/k× (XW_v + XW_1B_1W_3 + XW_1B_1W_2B_2W_3).

By now, we equate the softmax function to an exponential function and perform a Taylor expansion, making the operations of Q, K, and V more flexible. It is worth noting that the softmax function can transform the resulting matrix derived from Q and K^T into a full-rank matrix, while the matrix obtained by direct multiplication of Q and K^T is not a full-rank matrix, which significantly reduces the feature diversity, in turn, leads to a degradation of the model's performance <cit.>. Therefore, as done in <cit.>, we perform a depthwise convolution (DWC) operation on XW_v of Eq. (<ref>) to compensate for the performance degradation. Finally, we obtain:

    FinalOutput = 1/k× [DWC(XW_v) + 
     XW_1B_1W_3 + XW_1B_1W_2B_2W_3].

As shown in Fig. <ref>, we design the same network block to learn W_1, W_2, and W_3, which first reduces the number of channels to a quarter and then restores it to the original channel number. Meanwhile, an efficient Conv_1×1 operation is adopted to approximate W_v. 

As shown in Fig. <ref>, N and d are a multiplication of the width and height of a feature map and the number of channels. In general, N is far larger than d. According to Eq. (<ref>), our STEA strategy reduces the complexity of the original NLA from 𝒪(N^2× d) to 𝒪(N× d^2). The subsequent experimental results validate that our STEA is more efficient compared to the NLA module, however, exhibits a slight performance degradation due to the ignoring of remaining terms of Taylor extension and the use of the learning-based matrix. To compensate for such performance degradation, we design a Multi-Scale Large Field Reception (MLFR), which will be detailedly described in the following.





 §.§ Multi-Scale Large Field Reception (MLFR)

As mentioned in Section II-A, compared to the NLA module, our STEA achieves higher efficiency but sacrifices



the ability to extract global features from the feature maps. To this end, we design a Multi-Scale Large Field Reception (MLFR) strategy to obtain a large receptive field, as shown in Fig. <ref>. For the unfilled parts of the dilated convolution, we use normal convolution to fill in the dilation. Specifically, we process the input feature map X through three separate branches. In the first branch, the input is processed by a dilated convolution with an expansion rate of 6 and a depth-wise convolution with a kernel size of 7×7, respectively. The output results are concatenated along the channel dimension. Simultaneously, in the second branch, the input is processed by a dilated convolution with an expansion rate of 4 and a depth-wise convolution with a kernel size of 5×5, respectively. Similarly, the output results are concatenated along the channel dimension. In the last branch, the input is processed by a dilated convolution with an expansion rate of 2 and a depth-wise convolution with a kernel size of 3×3, respectively. Again, the output results are concatenated along the channel dimension. We represent the above execution process using the following equation: 

    Output_(i,j)=Concat(DDWC(X)_dr=i, 
    DWC(X)_ks=j)_i=2,4,6
    j=3,5,7,

where Concat(·,·) stands for concatenate operation, DDWC(·) stands for dilation convolution operation on the input, dr is the dilation rate, and DWC(·) stands for depth-wise convolution operation on the input, ks is the convolution kernel size. Furthermore, each value of i corresponds uniquely to one value of j. Subsequently, we employ 1×1 convolutions to perform channel-wise feature fusion on the outputs of the three branches, effectively reducing the dimension of the channels:

    Fusion_(i,j) = Conv_1×1(Output_(i,j))_i=2,4,6
    j=3,5,7,

where Fusion_(i,j) represents the output obtained after channel-wise feature fusion of the input, and Conv_1×1(·) denotes the 1×1 operation performed on the input. Similarly, each value of i corresponds uniquely to one value of j. In addition, the input X undergoes a 1×1 bypass convolution in addition to the aforementioned three branches:

    Bypass = Conv_1×1(X).

The purpose of the 1×1 convolution in Eq. (<ref>) is to perform subtle adjustments to the channels of the input X. Finally, we concatenate the outputs of the three branches with the output of the bypass 1×1 convolution along the channel dimension and perform channel-wise feature fusion using a 1×1 convolution to reduce the channel dimension:

    FinalOutput=Conv_1×1(Concat(Fusion_(i,j), 
     Bypass)_i=2,4,6
    j=3,5,7).

In summary, we utilize multiple dilated convolutions with different scales to extract features with various receptive fields. For the unpadded regions in dilated convolutions, we employ depth-wise convolutions to fill in the dilation. Additionally, we utilize 1×1 convolutions to perform subtle adjustments on the input as bypass convolutions. In the final step of feature fusion, we concatenate all the branches along the channel dimension and then employ a 1×1 convolution for channel-wise feature fusion, resulting in the final output as shown in Eq. (<ref>).







 §.§ Application to Laboratory Scenario (LabNet)

To show the effectiveness of our STEA and MLFR designs in a laboratory scenario, we construct a novel network termed LabNet for testing. As shown in Fig. <ref>, our LabNet is an end-to-end blind SR framework composed of three parts: shallow feature extraction, deep feature extraction, and HR image reconstruction. For the shallow feature extraction part, we employ a 3×3 convolution to extract shallow features from the input LR image. Subsequently, the extracted shallow features are fed into the UNet for deep feature extraction. Each LSTEA module in the UNet consists of multiple LSTEA blocks. When the input features pass through an LSTEA block, they undergo a channel-wise LayerNorm operation, followed by two branches. One branch utilizes multiple depth-wise convolutions for local feature extraction, while the other branch employs the STEA and MLFR modules for large-scale feature extraction. The channel features are then fused using a 1×1 convolution and filtered using the channel attention module. Subsequently, the features pass through the GDFN <cit.> module, which performs a nonlinear mapping operation on the input to enhance network performance. Finally, for the HR image reconstruction part, the channel dimension is adjusted using a 3×3 convolution. The H and W dimensions are then expanded to the target size using a pixel shuffle operation, and a final refinement is performed using a 3×3 convolution to obtain the ultimate output results.



 §.§ Application to Real-world Scenario (RealNet)

To show the effectiveness of our STEA and MLFR designs in real-world scenarios, we construct a novel network termed RealNet for testing. As illustrated in Fig. <ref>, our RealNet comprises three key components: shallow feature extraction, deep feature extraction, and HR image reconstruction. Firstly, we perform shallow feature extraction on the input LR image using two separate 3×3 convolutions. Subsequently, the extracted features are separately subjected to deep feature extraction in the denoising branch and the deblurring branch. Both the Denoising and Deblurring modules are composed of multiple LSTEA blocks. The features extracted from the Denoising module and the Deblurring module are then fused in proportion using Adapter modules. Finally, we employ a Fusion module for HR image reconstruction.

Adapter Module. For images contaminated by noise, there are numerous detailed features within the image that can be learned by the deblurring branch. Similarly, for blurry images, there are many low-frequency features within the image that can be learned by the denoising branch. Therefore, as depicted in Fig. <ref>, we propose the Adapter module. The parameters α and β control the proportions of denoising and deblurring, respectively. Analyzing the denoising branch, we take the features extracted by the deblurring branch in the previous step, apply a 3×3 convolution, and perform a channel-wise concatenation with the features extracted by the denoising branch. Subsequently, a 1×1 convolution is employed for channel-wise feature fusion, followed by a channel attention module for channel-wise feature selection. Similarly, a similar processing approach is utilized for the deblurring branch.

Fusion Module. To adjust the proportions of the outputs from the final denoising module and deblurring module, we also utilize two adjustable parameters (α and β). After adjusting the two sets of output feature maps, we perform a channel-wise concatenation and then apply a 1×1 convolution for channel information fusion and adjustment of channel dimensions. Finally, the output feature maps are resized to match the target size using a pixel shuffle operation and fine-tuned with a 3×3 convolution.

RealNet-GAN. To achieve better visual results, we employ a GAN training strategy for the proposed RealNet. Specifically, we use RealNet as the generator and utilize a U-Net as the discriminator <cit.>. The training process consists of two parts. Firstly, we train a PSNR-oriented model using the ℓ_1 loss, which is then used as the generator in the GAN training strategy. We then train a GAN-oriented model using a combination of ℓ_1 loss <cit.>, perceptual loss <cit.>, and adversarial loss <cit.>. The overall loss function of our network is

    ℒ_entire = ℒ_1 + ℒ_perc + 0.1 ×ℒ_adv,

where ℒ_1, ℒ_perc, and ℒ_adv represent the ℓ_1 loss, perceptual loss, and adversarial loss, respectively. According to our empirical study, their weights are set to be 1, 1, and 0.1, respectively. 
We utilize the {conv1,..., conv5} feature maps before activation in the pre-trained VGG19 <cit.> model (with weights {0.1, 0.1, 1, 1, 1}) as the perceptual loss.







§ EXPERIMENTS





 §.§ Datasets and Evaluation Metrics

We train all our proposed SR models using the DF2K dataset (DIV2K <cit.> + Flickr2K <cit.>). There are 3450 2K-resolution images in the dataset. For the quantitative evaluation, we conducted tests on five standard super-resolution image benchmark datasets: Set5 <cit.>, Set14 <cit.>, BSD100 <cit.>, Urban100 <cit.>, and Manga109 <cit.>. To conduct a qualitative evaluation, we compared the visual results of our approach with previous methods on the RealWorld38 <cit.> dataset. The quantitative performance metrics we use for comparison are the Peak Signal-to-Noise Ratio (PSNR) <cit.> and Structural Similarity (SSIM) <cit.>.





































 §.§ Implementation Details of LabNet

Our LabNet is built on the foundation of the LSTEA module, forming a U-Net. Each LSTEA module consists of multiple LSTEA blocks, as shown in Fig. <ref>, the number of LSTEA blocks in each LSTEA module from left to right is: {6, 8, 10, 6, 8, 8}. During the training process, we set the size of the input image to 64 × 64, and we use the Mean Absolute Error (MAE) as the loss function. The total iterations are 1.2×10^6. The initial learning rate is set to 4×10^-4 and is halved every 3×10^5 iterations. We use Adam <cit.> as the optimizer with β_1 = 0.9 and β_2 = 0.99. All models are trained on Tesla T4 GPUs.

Data Preprocessing of LabNet. For the training dataset, a blur operation with a blur kernel size of 21 × 21 is applied. For different scaling factors (×2, ×3, and ×4), the variance of the blur kernel is uniformly sampled during the training process from the ranges [0.2, 2.0], [0.2, 3.0], and [0.2, 4.0], respectively. For the test dataset, we uniformly select 8 blur kernels from the ranges [0.80, 1.60], [1.35, 2.40], and [1.8, 3.2] for different scaling factors (×2, ×3, and ×4), respectively.



 §.§ Implementation Details of RealNet

Our RealNet is constructed by two branches: denoising and deblurring. As illustrated in Fig. <ref>, both the denoising branch and the deblurring branch consist of 4 modules, each composed of 5 LSTEA modules. The entire network is trained for 30 epochs. The initial learning rate is set to 1× 10^-3 and is halved every 15 epochs. We use Adam as the optimizer with β_1 = 0.9 and β_2 = 0.99. All models are trained on Tesla T4 GPUs.

Data Preprocessing of RealNet. As shown in Fig. <ref>, we first independently and in parallel apply second-order noise and second-order blur processing to the HR image, resulting in two corresponding LR images. Then, these two LR images are passed through the denoising and deblurring branches of RealNet, respectively. Finally, a fusion operation is performed to obtain the restored SR image.






 §.§ Comparisons with State-of-the-art

Quantitative Evaluations in Lab Scenarios. From Table <ref>, it can be observed that our LabNet significantly outperforms other state-of-the-art deep SR models on nearly all benchmarks and scaling factors. For instance, compared to NLCUNet <cit.> with a scale factor of ×2, our LabNet achieves performance improvements of 0.03 dB, 0.05 dB, 0.11 dB, and 0.02 dB on the Set14, BSD100, Urban100, and Manga109 datasets, respectively. Similarly, compared to NLCUNet with a scale factor of ×4, our LabNet achieves performance improvements of 0.04 dB, 0.02 dB, and 0.11 dB on the Set14, BSD100, and Manga109 datasets, respectively. Our LabNet consistently outperforms DCLS <cit.> across all scale factors. Additionally, as shown in Table <ref>, when the input image size is (1,3,64,64) and the scaling factor is ×4, our LabNet achieves an inference time of 1.18 seconds on a GPU, which is 1.05 seconds faster than NLCUNet but slightly slower than DCLS. Furthermore, our LabNet has significantly fewer GFLOPs compared to NLCUNet and is comparable to DCLS. In summary, our LabNet not only achieves comparable performance to NLCUNet but also exhibits higher efficiency. This can be attributed to the utilization of LSTEA in our LabNet, which combines STEA and MLFR to capture multi-scale features with a large receptive field while maintaining linear complexity.


Qualitative Evaluations in Lab Scenarios. As shown in Fig. <ref>, by comparing the reconstruction result of the 'zebra' in Fig. <ref>, we can observe that the black and white stripes on the legs of the zebra generated by our LabNet are close to the GT image. Furthermore, by comparing the reconstruction result of 'YamatoNoHane', we can observe that our LabNet performs better in recovering the textual content in the image compared to other methods.


Quantitative Evaluation in Real-World Scenarios. As commonly observed, GAN-based SR methods often generate lower values in terms of PSNR and SSIM metrics while producing visually captivating results. As shown in Table <ref>, Real-ESRGAN and BSRGAN perform significantly lower than non-GAN-based SR methods on the five benchmark test datasets. However, Fig. <ref> reveals the impressive visual outcomes achieved by Real-ESRGAN and BSRGAN. Additionally, due to the absence of corresponding GT images for the degraded images in the RealWorld38 dataset, we only conducted visual comparisons between our RealNet and other blind SR methods.


Qualitative Evaluation in Real-World Scenarios. The experimental results are presented in Fig. <ref>. We evaluated the generalization capability of our model using the RealWorld38 dataset, which consists of real-world images. Our RealNet outperforms previous methods for reconstructing LR images in real-world scenarios. Among all the methods, the non-deep learning bicubic interpolation method performs the worst, followed by DANv2, DCLS, and NLCUNet. These methods only employ Gaussian blur and downsampling during training, resulting in visually unsatisfactory outcomes. BSRGAN <cit.> and Real-ESRGAN achieve favorable visual results due to their training on datasets with complex degradations. However, upon comparing the 'OST_009' image in Fig. <ref>, we observe that the result obtained by BSRGAN exhibits a significant color deviation from the LR input, while the result produced by Real-ESRGAN evokes a strong sense of discordance. In contrast, our RealNet reconstruction result closely resembles the color of the input LR image and demonstrates improved overall coherence.





 §.§ Ablation Study and Analysis of LabNet



  §.§.§ Impact of STEA

Table <ref> showcases the effect of employing first-order Taylor expansion approximation (FTEA) and second-order Taylor expansion approximation (STEA) on the performance of blind SR tasks.
By comparing the PSNR and SSIM values of the first and second rows in Table <ref>, we observed an improvement in performance across all test datasets when DWC(·) was applied to the Value. Furthermore, comparing the PSNR and SSIM values of the second and third rows in Table <ref>, we found that utilizing STEA for the exponential function outperformed FTEA across all test datasets. Particularly on the Urban100 dataset, the PSNR was higher by 0.11 dB when using a STEA compared to a FTEA. Furthermore, we found that when the order of Taylor expansion is greater than 2, the performance does not show a significant improvement. 

Efficiency Analysis. As shown in Table <ref>, when applying DWC(·) to Value, there is a slight increase in the number of parameters and a slight rise in MFLOPs. By comparing rows 2 to 4 in Table <ref>, as the order of Taylor expansion increases, the third-order Taylor expansion approximation (TTEA) parameter count and MFLOPs are greater than STEA, which are greater than FTEA. Ultimately, we used STEA to build our LabNet and RealNet.



  §.§.§ Impact of MLFR

Table <ref> highlights the influence of incorporating multi-scale dilated convolutions in MLFR on blind SR performance. By comparing the PSNR and SSIM values of the first to third rows in Table <ref>, we found that in MLFR, when utilizing dilated convolutions with dilations of 2, 4, and 6 simultaneously, the performance was the best across all test datasets. Particularly on the Urban100 dataset, the PSNR was 0.12 dB higher when using dilated convolutions with dilations of 2, 4, and 6 compared to dilated convolutions with dilations of 4 and 6. 

Efficiency Analysis. As shown in Table <ref>, in MLFR, the maximum number of parameters and MFLOPs are achieved when utilizing dilated convolutions with Dilated = 6, 4, and 2 simultaneously.




  §.§.§ Impact of STEA and MLFR

As shown in Table <ref>, LabNet achieves the best performance on all five benchmarks when both STEA and MLFR are used simultaneously.



 §.§ Ablation Study and Analysis of RealNet

The experimental results from Table <ref> highlight the following insights: Our RealNet without NLA has the fastest inference speed, processing a 64 × 64 input image in just 0.40 seconds. In contrast, Our RealNet with NLA and NLSA have much higher computational costs, with inference times of 4.21 seconds and 5.15 seconds, respectively. Our RealNet with the STEA, achieves high efficiency by processing an input image in only 0.48 seconds. Our RealNet with the STEA and the MLFR strikes a favorable balance between efficiency and visual result, processing an input image in 0.53 seconds.



 §.§ Details of adjustable parameters in RealNet

Adjusting a single parameter: Users can adjust the values of the denoising and deblurring parameters in RealNet to achieve the desired effects. In the visualization shown in Fig. <ref>, Increasing α_1 slightly enhances image details, but excessive increase may lead to distorted edges. Gradually increasing α_3 improves the distinction between objects and the background, but setting it too high can negatively impact the image. α_2 progressively enhances image visualization as it increases, but setting it to extremely large values can damage the image. Careful adjustment of α_4 is necessary to prevent color distortion, as it is related to the deep layer. It is advisable to avoid setting β_1 to β_4 to excessively large values to prevent color distortion in the resulting image. 

Joint modification of multiple parameters: Users can still fine-tune the denoising and deblurring ratios (i.e., α_i and β_i) according to their preferences, even after the training process. In the visualization shown in Fig. <ref>, the second and third rows display the restored results using different ratios for denoising and deblurring. This demonstrates that RealNet is capable of restoring rich texture information from LR images, catering to different user preferences.



§ CONCLUSION

In this paper, we presented a second-order Taylor expansion approximation (STEA) to handle the high computational complexity of self-similarity techniques. Then, we designed a multi-scale large-field reception (MLFR) to compensate for the super-resolution (SR) performance loss caused by STEA. To confirm the superiority of our designed STEA and MLFR, we respectively constructed the LabNet for testing in laboratory scenarios and the RealNet for testing in real-world scenarios. Extensive experimental results demonstrated the effectiveness of our method in blind SR tasks. The ablation study further validates their contribution to the overall framework.
IEEEbib



























 Paper Title: Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM.






Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM
    Xiaofeng Liu, Jonghye Woo, Chao Ma, Jinsong Ouyang, and Georges El Fakhri
X. Liu, C. Ma, J. Ouyang and G. El Fakhri are with the Department of Radiology and Biomedical Imaging, Yale University, New Haven, CT 06519.
J. Woo is with the Massachusetts General Hospital and Harvard Medical School, Boston, MA 02114.

    August 5, 2024
===========================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Delineating lesions and anatomical structure is important for image-guided interventions. Point-supervised medical image segmentation (PSS) has great potential to alleviate costly expert delineation labeling. However, due to the lack of precise size and boundary guidance, the effectiveness of PSS often falls short of expectations. Although recent vision foundational models, such as the medical segment anything model (MedSAM), have made significant advancements in bounding-box-prompted segmentation, it is not straightforward to utilize point annotation, and is prone to semantic ambiguity. In this preliminary study, we introduce an iterative framework to facilitate semantic-aware point-supervised MedSAM. Specifically, the semantic box-prompt generator (SBPG) module has the capacity to convert the point input into potential pseudo bounding box suggestions, which are explicitly refined by the prototype-based semantic similarity. This is then succeeded by a prompt-guided spatial refinement (PGSR) module that harnesses the exceptional generalizability of MedSAM to infer the segmentation mask, which also updates the box proposal seed in SBPG. Performance can be progressively improved with adequate iterations. We conducted an evaluation on BraTS2018 for the segmentation of whole brain tumors and demonstrated its superior performance compared to traditional PSS methods and on par with box-supervised methods.


 




§ INTRODUCTION
 



The costly labeling effort in medical image delineation significantly hinder the development of data-driven AI models. There are increasing interests on weakly supervised segmentation to utilize bounding box or even a single point as label for training supervision <cit.>. However, due to the absence of accurate size and boundary guidance, there is a large performance gap between point-supervised medical image segmentation (PSS) and mask/box-supervised counterparts <cit.>. 

The recent progress of vision foundational models has achieved breakthroughs in several weakly supervised tasks, benefited by their strong zero-shot generalizability. For instance, the point-prompt natural image segment anything model (SAM) has been applied to enhance pseudo labels for point-supervised segmentation <cit.>. Recently, the medical image version of SAM (MedSAM) <cit.> has been trained with 1.5 million segmented medical images to achieve generalizable box-prompt segmentation. MedSAM takes both an image slice and a bounding box, i.e., prompt, to predict the possible segmentation mask within the box. Notably, the natural image SAM itself follows point-prompt design. Therefore, it is not straight forward to integrate box-prompt MedSAM to PSS task as <cit.>. In addition, a significant limitation of SAM/MedSAM is the lack of classification ability, resulting in class / structure-agnostic segmentation results <cit.>. They are designed for general use, while fails to accurately delineate the specific lesion or structure as desired <cit.>.  



 

To our knowledge, this is the first attempt to integrate MedSAM to facilitate semantic-aware PSS. We adopt an iterative framework to achieve coarse-to-fine progression of the bounding box. The point prompt is first converted to a box-proposal seed. Then, we configure a semantic box-prompt generator (SBPG) to propose and pick the reasonable box according to semantic similarity as <cit.>. It is followed by a box prompt guided spatial refinement (PGSR) to utilize the generalizable MedSAM to predict the segmentation mask. In addition, the smallest box that covers the mask is further used as the box seed in the next round of SBPG. We do not need to fine-tune the MedSAM to be semantic aware, which is prone to catastrophic forgetting. Notably, only the prompt refiner θ with the encoder part of the segmentor model will be trained, which involves about half the parameters of mask-supervised UNet training. 

We demonstrated its effectiveness in BraTS2018 for PSS of the whole brain tumor segmentation from T2-weighted MRI slices. We show that in testing, with 3 to 5 rounds of iteration, the point-prompt can achieve superior performance to approximate the box-supervised MedSAM <cit.>.








§ METHODOLOGY

In PSS, we are given tuplet {x^i,ρ_c^i,s^i}, in which x^i can be an MR slice, while the point prompt ρ^i_c={ρ^i_x,ρ^i_y,c} indicates the 2D spatial coordinates and the interested class c, e.g., tumor or normal tissue. We would expect the segmentation mask prediction ŝ^i_c to approximate the expert annotation s^i_c. 

To catering the box-prompt MedSAM <cit.>, an initial box proposal seed of class c, i.e., b_c:0^i, is generated from ρ^i_c, which is centered on point location (ρ^i_x,ρ^i_y) with the size of X× Y. We do not expect b_c:0^i fit the region of interests very well as there is no prior information about the size in the initial step. Then, the proposal bag ℬ^i_c={b_c:n^i}_n=1^N is created by scaling b_c:0^i with N different scales. Therefore, the boxes in group ℬ^i_c has strong spatial correlation with ρ_c^i, which avoid the inefficient random proposal in whole image <cit.>. 

To enable the model be aware of the specific semantic class, for example tumor in our task, a parameterized prompt refiner θ with a segmentor encoder and two fully connected layers is applied. Specifically, for the training sample with segmentation mask label s^i_c, we convert s^i_c to the corresponding smallest rectangle bounding box label b^i_c. Then, we store a set of feature f^i_c:n=θ(x^i,b_c:n^i) in the most recent M batches to a memory buffer, and calculate the mean of {f^i_c:nm}_n,m=1^N,M as the prototype 𝒱_c. As conventional multiple instance learning (MIL) based methods <cit.>,
the instance level probability indicating the likelihood of b^i_c:n matches 𝒱_c can be approximated with 
    p(b^i_c:n,𝒱_c)=e^cos(f^i_c:n,𝒱_c)/∑_c=1^C e^cos(f^i_c:n,𝒱_c),
where cos(·) indicates the cosine similarity, and C is the number of segmentation category. With the ground truth b^i_c, θ is trained with the binary cross entropy loss of 

    ℒ=-∑_c=1^C {c^i_1log∑_n^N p(b^i_c:n,𝒱_c) + (1-c^i_1)log(1-∑_n^N p(b^i_c:n,𝒱_c))},
where c^i_1 is the one-hot class label. Therefore, the prompt refiner is optimized to be semantic aware for the specific anatomical structure. The prototype in buffer is also updated with the new model θ. With p(b^i_c:n,𝒱_c) for each b^i_c:n, we pick the highest probability one  in ℬ^i_c as our optimal proposal b^i_c:* in current stage, which is also our bounding box inference.


Then, in the PGSR module, the off-the-shelf MedSAM <cit.> takes both x^i and b^i_c:* to predict the possible segmentation mask ŝ^i_c. Notably, no information about the interested class can be directly informed in MedSAM. Although the MedSAM has strong ability of zero-shot segmentation to delineate the relatively accurate mask within the current box prompt b^i_c:*, the alignment of ŝ_c^i is highly dependent on the precise b^i_c:*. This motivated us to refine b^i_c:* with a more suitable proposal seed b_c:0^i in SBPG. 

In practice, we can simply apply the smallest bounding box b̂^i_c:0 for ŝ^i_c:0, and use b̂^i_c:0 as our box proposal seed in the next round of SBPG. With T rounds of iteration, we expect that the box-prompt can be progressively refined to inform the MedSAM to predict an accurate segmentation mask. 

In particular, in the last round of testing, we do not need to generate b̂^i_c:0. Instead, the T-th round ŝ^i_c:0 is used as our final prediction.









§ EXPERIMENTS AND RESULTS

We evaluated on BraTS dataset for whole tumor segmentation as <cit.>, and use T2-weighted slices as our input. We used the 80%/20% split for training or testing. We generate the smallest bounding box for the mask and use its center as ρ^i_c. Our prompt refiner adopted the encoder part of the segmentor in <cit.>, which is followed by two fully connected layers of 1024 and 256 dimensions, respectively. We empirically set X× Y=21×21 and search proper T∈{1,2,3,5,10}. We followed <cit.> to compare with the conventional point-supervised method of WISE-Net without the use of the vision fundation model. Also, we directly input the ground truth of the bounding box into MedSAM, which can be “upper bound" of the performance of pointwise method. In Tab.<ref>, we can see that our model with T≥5 outperforms WISE-Net by a large margin. The Dice score is close to the box-supervised MedSAM.   





§ CONCLUSION

In this work, we proposed an efficient iterative framework to enable box-prompt MedSAM to point-supervised medical image segmentation. A lightweight prompt refiner with only encoder is specifically trained for the interested structural class. We show that the fixed off-the-shelf large model can flexibly target the specific downstream task without the need of large scale fine-tuning. The medical PSS performance can be largely improved by the advance of vision foundation model, e.g., MedSAM, which show great promise for facilitating the image-guided interventions. 
 


 


§ ACKNOWLEDGMENT

This work is supported by R21EB034911, R01CA165221, R01DC018511, P41EB022544, and NAIRR240016. 

ieee





