PA OmniNet: A retraining-free, generalizable deep learning framework for robust photoacoustic image reconstruction

For clinical translation of photoacoustic imaging cost-effective systems development is necessary. One approach is the use of fewer transducer elements and acquisition channels combined with sparse sampling. However, this approach introduces reconstruction artifacts that degrade image quality. While deep learning models such as U-net have shown promise in reconstructing images from limited data, they typically require retraining for each new system configuration, a process that demands more data and increased computational resources. In this work, we introduce PA OmniNet, a modified U-net model designed to generalize across different system configurations without the need for retraining. Instead of retraining, PA OmniNet adapts to a new system using only a small set of example images (between 4 and 32), known as a context set. This context set conditions the model to effectively remove artifacts from new input images in various sparse sampling photoacoustic imaging applications. We evaluated PA OmniNet against a standard U-net using multiple datasets, including in vivo data from mouse and human subjects, synthetic data, and images captured at different wavelengths. PA OmniNet consistently outperformed the traditional U-net in generalization tasks, achieving average improvements of 8.3% in the Structural Similarity Index, a 11.6% reduction in Root Mean Square Error, and a 1.55 dB increase in Peak Signal-to-Noise Ratio. In 66% of our test cases, the generalized PA OmniNet even outperformed U-net models trained specifically on the new dataset. Code is available at https://github.com/olivierstam4/PA_OmniNet.