Deep Learning for Photoacoustic Imaging Reconstruction

Implementation of models for the work presented in the following articles.

Jensen, Cristian Perez, Navchetan Awasthi, and Kalloor Joseph Francis. "Transformer-based optical attenuation compensation and denoising in photoacoustic imaging." Journal of Biomedical Optics 30.11 (2025): 116004-116004.https://www.spiedigitallibrary.org/journals/journal-of-biomedical-optics/volume-30/issue-11/116004/Transformer-based-optical-attenuation-compensation-and-denoising-in-photoacoustic-imaging/10.1117/1.JBO.30.11.116004.full

Cristian P. Jensen, Kalloor Joseph Francis, and Navchetan Awasthi "Image depth improvement in photoacoustic imaging using transformer based generative adversarial networks", Proc. SPIE PC12842, Photons Plus Ultrasound: Imaging and Sensing 2024, PC128421V (13 March 2024); https://doi.org/10.1117/12.3001537

Please cite this if you find this useful in your work.

Models

The following image-to-image translation models are implemented using PyTorch and PyTorch Lightning:

• Pix2Pix (Isola et al. 2018);
• Attention U-net (Oktay et al. 2018);
• Residual U-net with the following basic blocks:
  • Res18 (He et al. 2015);
  • Res50 (He et al. 2015);
  • ResV2 (He et al. 2016);
  • ResNeXt (Xie et al. 2017).
• Trans U-net (Chen et al. 2021);
• Palette (Saharia et al. 2022).

More models can easily be added by using the UnetWrapper class.

Loss functions

The following loss functions are implemented:

• GAN loss, using Pix2Pix' discriminator (to change the used adversarial network, you must change the Discriminator class in models/wrapper.py);
• MSE loss;
• SSIM loss;
• PSNR loss.
• Combination of SSIM and PSNR loss.

Data organisation

The organisation of your data does not matter. The only important thing is the data file, a YAML file containing a list of input-ground truth entries. The input and ground truth files must be relative to the directory of the data file. For example:

- input: input/00001.png
  ground_truth: ground_truth/00001.png
- input: input/00002.png
  ground_truth: ground_truth/00002.png
- input: input/00003.png
  ground_truth: ground_truth/00003.png

Training a model

To train a model, run the following:

python main.py <run name> <options>

When training, the model with the highest SSIM on the validation dataset will be selected as the "best" checkpoint.

Testing a model

To test a trained model, run the following:

python report.py <report name> <options>

It essentially takes a model checkpoint and test data file as input and outputs metrics and information about the model. The following metrics are reported:

• SSIM per image;
• PSNR per image;
• Mean SSIM;
• Mean PSNR;
• Mean RMSE;
• FLOPs;
• Parameter count;
• SSIM over depth (vertically) of the image (this is only relevant for PAI reconstruction).
• Outputs of the model.