Merge pull request #3 from psweens/1.2.0

1.2.0

README.md

@@ -1,17 +1,29 @@
-# VAN-GAN: Vessel Segmentation Generative Adversarial Network
+# VAN-GAN: An Unsupervised Deep Learning Tool for Vascular Network Segmentation
+
+***VAN-GAN offers an accessible and efficient solution for the segmentation of vascular networks from 3D images without the need for human-annotated labels.***
 
 ## Introduction
-Innovations in imaging hardware have led to a revolution in our ability to visualise vascular networks in 3D at high resolution. The segmentation of microvascular networks from these 3D image volumes and interpretation of their meaning in the context of physiological and pathological processes unfortunately remains a time consuming and error-prone task. Deep learning has the potential to solve this problem, but current supervised analysis frameworks require human-annotated ground truth labels. To overcome these limitations, we present an unsupervised image-to-image translation deep learning model called the ***Vessel Segmentation Generative Adversarial Network (VAN-GAN)***. VAN-GAN integrates synthetic blood vessel networks that closely resemble real-life anatomy into its training process and learns to replicate the underlying physics of an imaging system in order to learn how to segment vasculature from 3D biomedical images. By leveraging synthetic data to reduce the reliance on manual labelling, VAN-GAN lower the barriers to entry for high-quality blood vessel segmentation to benefit imaging studies of vascular structure and function.
+
+Recent advancements in imaging hardware have opened up new possibilities for visualising 3D vascular networks with exceptional resolution. However, the task of segmenting microvascular networks from these 3D image volumes and interpreting their physiological and pathological significance remains a time-consuming and error-prone process. Deep learning holds great potential in addressing this challenge, but existing supervised frameworks rely on labor-intensive human-annotated ground truth labels. To overcome these limitations, we introduce ***Vessel Segmentation Generative Adversarial Network (VAN-GAN)***, an unsupervised image-to-image translation deep learning model.
+
+![alt text](VANGAN_Overview.jpg)
+*(A) VAN-GAN performs unsupervised training to learn the mappings between the imaging and segmentation domains to train a generator which can segment 3D vasculatures from biomedical image volumes. (B) The two domains are define as: 1) Imaging domain, the dataset of real biomedical image volumes; and 2) segmentation domain, the dataset of computer-generated 3D vascular labels. Example of mouse ear imaged using mesoscopic photoacoustic imaging shown.*
 
 ## Methodology
-This Python package utilises image-to-image translation to segment 3D biomedical image volumes of vascular networks. Our unsupervised deep learning framework builds upon [CycleGAN](https://arxiv.org/abs/1703.10593) in several ways:
-* Extend the design to 3D for image volume generation using 3D convolutions.
-* Utilise a deep residual U-Net architecture for generators.
-* Apply random Gaussian noise to discriminator inputs and convolution layers for improved training stability and regularisation.
-* Use a modified objective loss function:
-  * Introduce a structure similarity reconstruction loss between real and cycled biomedical image volumes.
-  * Introduce a [spatial and topological constraint](https://arxiv.org/abs/2003.07311) between real and cycled segmentation labels.
-  * Exclude identity loss.
+
+VAN-GAN is a Python package that employs image-to-image translation techniques to segment 3D biomedical image volumes of vascular networks. Our framework builds upon the principles of CycleGAN with several extensions tailored for vascular network segmentation:
+
+* **3D Extension**: We extend the design to support 3D image volume generation using 3D convolutions, enabling accurate segmentation of vascular structures in three dimensions.
+* **Deep Residual U-Net Architecture**: We leverage a deep residual U-Net architecture for generators, providing a powerful tool for capturing fine details and improving segmentation performance.
+* **Training Stability and Regularisation**: We introduce random Gaussian noise to discriminator inputs and convolution layers to enhance training stability and regularisation, leading to more robust and reliable segmentation results.
+* **Modified Objective Loss Function**: The loss function is enhanced to include the following components:
+  * **Structure Similarity Reconstruction Loss**: This loss fosters similarity between real and cycled biomedical image volumes, aiding in preserving essential features during image-to-image translation. 
+  * **Spatial and Topological Constraint**: A constraint inspired by spatial and topological considerations is introduced to align real and cycled segmentation labels, further refining the segmentation accuracy. 
+  * **Exclusion of Identity Loss**: The identity loss component is excluded from the objective, streamlining the segmentation process and avoiding unnecessary complexities. 
+* **Synthetic 3D Images**: VAN-GAN leverages synthetic blood vessel networks closely resembling real-life anatomy during training, effectively removing the reliance on manual labelling.
+
+![alt text](Generator_Labelling_Architecture.jpg)
+*The VAN-GAN model trains a 3D deep residual U-Net architecture to segment 3D images of real blood vessel networks. An example of a tumour vasculature imaged using mesoscopic photoacoustic imaging is shown where the image is provided as input and the 3D segmentation mask predicted as output.*
 
 ## Installation
 To install the package from source, download the latest release on the VAN-GAN repository or run the following in a terminal window:
@@ -23,7 +35,7 @@ Please ensure that cudatoolkit, cudnn and Tensorflow are correctly installed as
 
 Tensorflow and the remaining Python packages below can be installed in a [_conda_](https://www.anaconda.com/download/) environment.
 
-The remaining required Python packages can be install using _pip_ in a terminal window:
+The remaining required Python packages can be installed using _pip_ in a terminal window:
 ```bash
 pip install opencv-python scikit-image tqdm tensorflow_addons tensorflow-mri joblib matplotlib
 ```
@@ -36,7 +48,6 @@ VAN-GAN has been tested on Ubuntu 22.04.2 LTS with Python 3.9.16 and the followi
 * Scikit-image 0.20.0
 * [tqdm 4.65.0](https://github.com/tqdm/tqdm)
 * [Tensorflow-addons 0.20.0](https://www.tensorflow.org/addons)
-* [Tensorflow-mri 0.22.0](https://github.com/mrphys/tensorflow-mri)
 * Joblib 1.2.0
 * Matplotlib 3.7.1
 

custom_callback.py

@@ -15,15 +15,15 @@ class GanMonitor:
     def __init__(self,
                  args,
                  dataset=None,
-                 Alist=None,
-                 Blist=None,
+                 imaging_val_data=None,
+                 segmentation_val_data=None,
                  process_imaging_domain=None):
 
         self.imgSize = args.INPUT_IMG_SIZE
-        self.test_AB = dataset.valFullDatasetA
-        self.test_BA = dataset.valFullDatasetB
-        self.Alist = Alist
-        self.Blist = Blist
+        self.imaging_val_full_vol_data = dataset.imaging_val_full_vol_data
+        self.segmentation_val_full_vol_data = dataset.segmentation_val_full_vol_data
+        self.imaging_val_data = imaging_val_data
+        self.segmentation_val_data = segmentation_val_data
         self.process_imaging_domain = process_imaging_domain
         self.period = args.PERIOD_2D_CALLBACK,
         self.period3D = args.PERIOD_3D_CALLBACK,
@@ -172,7 +172,7 @@ def stitch_subvolumes(self, gen, img, subvol_size,
                           start_col:(start_col + kW),
                           start_dep:(start_dep + kD)]
 
-                    if process_img == True and self.process_imaging_domain is not None:
+                    if process_img and self.process_imaging_domain is not None:
                         arr = self.process_imaging_domain(arr)
 
                     arr = gen(np.expand_dims(arr,
@@ -263,7 +263,7 @@ def imagePlotter(self, epoch, filename, setlist, dataset, genX, genY, nfig=6, ou
                 tf.image.random_crop(sample, size=(self.imgSize[1], self.imgSize[2], self.imgSize[3], self.imgSize[4])),
                 axis=0)
 
-        if process_img == True and self.process_imaging_domain is not None:
+        if process_img and self.process_imaging_domain is not None:
             sample = self.process_imaging_domain(sample)
 
         prediction = genX(sample, training=False)
@@ -323,7 +323,7 @@ def imagePlotter(self, epoch, filename, setlist, dataset, genX, genY, nfig=6, ou
         plt.close()
 
         # Generate 3D predictions, stitch and save
-        if epoch % self.period3D == 1 and outputFull:  # and epoch > 160:
+        if epoch % self.period3D == 1 and outputFull and epoch > 160:
             self.stitch_subvolumes(genX, storeSample.numpy(),
                                    self.imgSize, epoch=epoch, name=sampleName, process_img=process_img)
 
@@ -463,8 +463,10 @@ def on_epoch_end(self, model, epoch, logs=None):
         """
 
         # Generate 2D plots
-        self.imagePlotter(epoch, "genAB", self.Alist, self.test_AB, model.gen_IS, model.gen_SI, process_img=True)
-        self.imagePlotter(epoch, "genBA", self.Blist, self.test_BA, model.gen_SI, model.gen_IS, outputFull=True)
+        self.imagePlotter(epoch, "genIS", self.imaging_val_data, self.imaging_val_full_vol_data, model.gen_IS,
+                          model.gen_SI, process_img=True)
+        self.imagePlotter(epoch, "geSI", self.segmentation_val_data, self.segmentation_val_full_vol_data, model.gen_SI,
+                          model.gen_IS, outputFull=True)
 
     def run_mapping(self, model, test_set, sub_img_size=(64, 64, 512, 1), segmentation=True, stride=(25, 25, 1),
                     padFactor=0.25, filetext=None, filepath=''):