Ultrasound Imaging Based on the Variance of a Diffusion Restoration Model

PaperLink

The structure of the repository:

├── <DRUSvar>  
│   ├── <configs> 
│   │   ├──imagenet_256_1c.yml  # The configure file for the 1-channel case
│   │   ├──imagenet_256_3c.yml  # The configure file for the 3-channel case 
│   ├── <exp>  
│   │   ├──<logs>             # The fined-tuned diffusion models, downloaded from: https://uncloud.univ-nantes.fr/index.php/s/SWamKLe3W5JTbSo
│   │   │   ├──<image1c>      # -- 1-channel model trained only on ImageNet   
│   │   │   ├──<vitro1c>      # -- 1-channel model trained on in-vitro CIRS datasets
│   │   │   ├──<both1c>       # -- 1-channel model trained on in-vitro CIRS and in-vivo carotid datasets
│   │   │   ├──<vtro>         # -- 3-channel model trained on in-vitro CIRS datasets
│   │   │   ├──<vivo>         # -- 3-channel model trained on in-vivo carotid datasets
│   │   ├──<slurms>                 # Slurm files if use HPC
│   │   │   ├──jobTemplate.sl       # -- a template file
│   │   │   ├──numeCyst-m1.sl       # -- for the numerical-cysts experiment, m1 means with the first simulated multiplicative map (in paper we simulated 9 independent multiplicative maps for a statistical analysis)
│   │   │   ├──numeScatterers-m1.sl # -- for the numerical-scatterers experiment
│   │   │   ├──picmusDRUS.sl        # -- our approach, use the 'both1c' diffusion model to restore all picmus images
│   │   │   ├──picmusDRUSvitro.sl   # -- our approach, use the 'vitro' diffusion model to restore picmus vitro images
│   │   │   ├──picmusDRUSvivo.sl    # -- our approach, use the 'vivo' diffusion model to restore picmus vivo images
│   │   │   ├──picmusDENO.sl        # -- denoising approach (BH \approx I), use the 'both1c' diffusion model to restore all picmus images
│   │   │   ├──picmusDENOvitro.sl   # -- denoising approach (BH \approx I), use the 'vitro' diffusion model to restore picmus vitro images
│   │   │   ├──picmusDENOvivo.sl    # -- denoising approach (BH \approx I), use the 'vivo' diffusion model to restore picmus vivo images
│   ├── <functions>           
│   │   ├──ckpt_util.py       
│   │   ├──denoising.py       # DDRM algorithm
│   │   ├──svd_replacement.py # includes operations with the SVD components    
│   ├── <guided_diffusion>    
│   ├── <runners>  
│   │   ├──diffusion.py       # IMPORTANT!!!
│   ├── main.py               # Define arguments & load configurations
│   ├── *.err, *.out          # log files generated by running the slurms
│   │ 
├── <MATLABfiles> 
│   ├── <numerical>             
│   │   ├──<SimulatedData>      # simulated data
│   │   │   ├──<additiveNoises> # -- the simulated additive noise maps
│   │   │   ├──<cysts>          # -- the simulated reflectivity maps for the cyst phantom
│   │   │   ├──<scatterers>     # -- the simulated reflectivity maps for the scatterer phantom
│   │   │   ├──a_Simu_n.m       # -- to simulate additive noise 
│   │   │   ├──b_Simu_p.m       # -- to simulate echogenicity map 
│   │   │   ├──c_Simu.mp.m      # -- to simulate reflectivity map
│   │   │   ├──anisoKernel.m    # -- to illustrate the anisotropic degradation kernel.
│   │   ├──<Test_cysts>         # RESULTS AND IMAGES USED IN THE PAPER
│   │   ├──<Test_scatterers>    # RESULTS AND IMAGES USED IN THE PAPER
│   ├── <picmus>              
│   │   ├──<DAS>         # the baseline results of using DAS
│   │   ├──<Observation> # By (measurements), the inputs to diffusion models
│   │   │   ├──<DENO>         
│   │   │   ├──<DRUS>         
│   │   ├──<SVD>         # SVD components, required if use DRUS, can be downloaded from: https://drive.google.com/drive/folders/10KwoH5G-s8Gk_aCj7WxTZ_L3596u44dI?usp=sharing
│   │   │   ├──Sigma.mat      
│   │   │   ├──Ud.mat         
│   │   │   ├──Vd.mat         
│   │   │   ├──compute_mainDRUS_svdBH.m   
│   │   ├──<Test_picmus> # RESULTS AND IMAGES USED IN THE PAPER
│   │ 
│   ├── <src>        # help resources
├── environment.yml  

Fine-tuning

The ultrasound datasets and the code for fine-tuning can be found in this repository.

Notes: The original diffusion model was trained on ImageNet color images. Based on this model, we fine-tuned several checkpoints, some of which still utilize RGB channels. These models are referred to as 3-channel (3c) models. For a 3c model, two-dimensional images are repeated three times before being input into the network, and the output is processed by averaging the dimensions, reducing [3, 256, 256] to [1, 256, 256].

Although the 3-channel models demonstrated good restoration quality, we further modified the fine-tuning process to support 1-channel (1c) data. This adjustment eliminates the need for the repeat and average operations, streamlining the process for single-channel ultrasound images.

Reproduction

Fig. 1 (left)      -- run the script: `../MATLABfiles/numerical/Test_cysts/a_phantom_display.m`.
Fig. 1 (right)     -- run the script: `../MATLABfiles/numerical/Test_scatterers/a_phantom_display.m`
Fig. 2 (lineCharts)-- run the script: `../MATLABfiles/numerical/Test_cysts/b_lineCharts.m`
Fig. 2 (images)    -- run the script: `../MATLABfiles/numerical/Test_cysts/c_images.m`
Fig. 3 (lineCharts)-- run the script: `../MATLABfiles/numerical/Test_scatterers/b_lineCharts.m`
Fig. 3 (images)    -- run the script: `../MATLABfiles/numerical/Test_scatterers/c_images.m`
Fig. 4             -- run the script: `../MATLABfiles/picmus/Test_picmus/picmusImages2.m`

Run Diffusion Sampling

on HPC using slurm, or as follows:

python main.py --ni --config {CONFIG.yml} --doc {MODELFOLDER} --ckpt {CKPT} --deg {DEG} --matlab_path {MATLABPATH}

where

• CONFIG is the name of the config file (see DRUSvar/configs/), including hyperparameters such as batch size and network architectures.
• MODELFOLDER is the name of the folder saving the diffusion model checkpoints
• CKPT is the name of the selected diffusion model checkpoint
• DEG is the degradation type: DRUS | DRUSvitro | DRUSvivo | DENO | DENOvitro | DENOvivo | NumericalCysts| NumericalScatterers
• MATLABPATH is the path of the folder <MATLABfiles/>.

For example (Apply DRUS and DENO)

(DENO is faster but performs less effectively than DRUS.)

• on the picmus 6 datasets (4 vitro + 2 vivo) with the fine-tuned both1c model. Results will be saved in ../MATLABfiles/picmus/Test_picmus/results/DRUS/ and ../MATLABfiles/picmus/Test_picmus/results/DENO/.

  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_1c.yml --doc both1c  --ckpt model004000.pt  --deg DRUS --matlab_path /home/.../MATLABfiles/
  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_1c.yml --doc both1c  --ckpt model004000.pt  --deg DENO --matlab_path /home/.../MATLABfiles/
  

  After finishing the diffusion sampling process for DRUS and DENO, by running picmusImages.m, you can see a image similar to the paper's Fig. 4.
  We observed that processing in-vitro and in-vivo datasets separately leads to a better image quality, so if you wanna see better images, you can try the code as follows:

• process 4 vitro datasets with the fine-tuned vitro model. Results will be saved in ../MATLABfiles/picmus/Test_picmus/results/DRUSvitro/ and ../MATLABfiles/picmus/Test_picmus/results/DENOvitro/.

  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_3c.yml --doc vitro  --ckpt model006000.pt  --deg DRUSvitro --matlab_path /home/.../MATLABfiles/
  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_3c.yml --doc vitro  --ckpt model006000.pt  --deg DENOvitro --matlab_path /home/.../MATLABfiles/
  

  process 2 vivo datasets with the fine-tuned vivo model. Results will be saved in ../MATLABfiles/picmus/Test_picmus/results/DRUSvivo/ ../MATLABfiles/picmus/Test_picmus/results/DENOvivo/.

  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_3c.yml --doc vivo  --ckpt model004000.pt  --deg DRUSvivo --matlab_path /home/.../MATLABfiles/
  python /home/.../DRUSvar/main.py --ni --config /home/.../DRUSvar/configs/imagenet_256_3c.yml --doc vivo  --ckpt model004000.pt  --deg DENOvivo --matlab_path /home/.../MATLABfiles/
  

  After finishing, run picmusImages2.m

References and Acknowledgements

@inproceedings{DRUSvar,
    title={Ultrasound Imaging based on the Variance of a Diffusion Restoration Model},
    author={Zhang, Yuxin and Huneau, Cl{\'e}ment and Idier, J{\'e}r{\^o}me and Mateus, Diana},
    booktitle={EUSIPCO},
    year={2024}
}

This implementation is based on / inspired by:

• https://ddrm-ml.github.io/ (the DDRM repo)