[WIP] 519 Add Radiogenomic GAN

.gitignore

@@ -124,3 +124,18 @@ temp/
 # temporary testing data MedNIST
 tests/testing_data/MedNIST*
 tests/testing_data/*Hippocampus*
+
+model_out
+MedNIST/
+MedNIST.tar.gz
+hand-gan*/
+ModelOut/
+research/radiogenomic-gan-nodule-synthesis/ModelOut/
+
+examples/workflows/runs/
+
+output/
+
+research/radiogenomic-gan-nodule-synthesis/datasets.py
+
+.vscode/

examples/synthesis/gan_training.py

@@ -82,11 +82,11 @@ def main():
     real_dataloader = DataLoader(real_dataset, batch_size=batch_size, shuffle=True, num_workers=10)
 
     # define function to process batchdata for input into discriminator
-    def prepare_batch(batchdata):
+    def prepare_batch(batchdata, device):
         """
         Process Dataloader batchdata dict object and return image tensors for D Inferer
         """
-        return batchdata["hand"]
+        return batchdata["hand"].to(device)
 
     # define networks
     disc_net = Discriminator(
@@ -179,7 +179,7 @@ def generator_loss(gen_images):
         discriminator_loss,
         d_prepare_batch=prepare_batch,
         d_train_steps=disc_train_steps,
-        latent_shape=latent_size,
+        latent_size=latent_size,
         key_train_metric=key_train_metric,
         train_handlers=handlers,
     )
@@ -190,7 +190,7 @@ def generator_loss(gen_images):
     # Training completed, save a few random generated images.
     print("Saving trained generator sample output.")
     test_img_count = 10
-    test_latents = make_latent(test_img_count, latent_size).to(device)
+    test_latents = make_latent(test_img_count, latent_size, device)
     fakes = gen_net(test_latents)
     for i, image in enumerate(fakes):
         filename = "gen-fake-final-%d.png" % (i)

monai/engines/trainer.py

@@ -16,7 +16,7 @@
 from torch.utils.data import DataLoader
 
 from monai.engines.utils import CommonKeys as Keys
-from monai.engines.utils import GanKeys, default_make_latent, default_prepare_batch
+from monai.engines.utils import GanKeys, default_gan_prepare_batch, default_make_latent, default_prepare_batch
 from monai.engines.workflow import Workflow
 from monai.inferers import Inferer, SimpleInferer
 from monai.transforms import Transform
@@ -176,9 +176,8 @@ class GanTrainer(Trainer):
         d_optimizer: D optimizer function.
         d_loss_function: D loss function for optimizer.
         g_inferer: inference method to execute G model forward. Defaults to ``SimpleInferer()``.
-        d_inferer: inference method to execute D model forward. Defaults to ``SimpleInferer()``.
         d_train_steps: number of times to update D with real data minibatch. Defaults to ``1``.
-        latent_shape: size of G input latent code. Defaults to ``64``.
+        latent_size: size of G input latent code. Defaults to ``64``.
         d_prepare_batch: callback function to prepare batchdata for D inferer.
             Defaults to return ``GanKeys.REALS`` in batchdata dict.
         g_prepare_batch: callback function to create batch of latent input for G inferer.
@@ -209,10 +208,9 @@ def __init__(
         d_optimizer: Optimizer,
         d_loss_function: Callable,
         g_inferer: Inferer = SimpleInferer(),
-        d_inferer: Inferer = SimpleInferer(),
         d_train_steps: int = 1,
-        latent_shape: int = 64,
-        d_prepare_batch: Callable = default_prepare_batch,
+        latent_size: int = 64,
+        d_prepare_batch: Callable = default_gan_prepare_batch,
         g_prepare_batch: Callable = default_make_latent,
         g_update_latents: bool = True,
         iteration_update: Optional[Callable] = None,
@@ -240,9 +238,8 @@ def __init__(
         self.d_network = d_network
         self.d_optimizer = d_optimizer
         self.d_loss_function = d_loss_function
-        self.d_inferer = d_inferer
         self.d_train_steps = d_train_steps
-        self.latent_shape = latent_shape
+        self.latent_size = latent_size
         self.g_prepare_batch = g_prepare_batch
         self.g_update_latents = g_update_latents
 
@@ -263,9 +260,9 @@ def _iteration(
         if batchdata is None:
             raise ValueError("must provide batch data for current iteration.")
 
-        d_input = self.prepare_batch(batchdata).to(engine.state.device)
+        d_input = self.prepare_batch(batchdata, engine.state.device)
         batch_size = self.data_loader.batch_size
-        g_input = self.g_prepare_batch(batch_size, self.latent_shape, batchdata).to(engine.state.device)
+        g_input = self.g_prepare_batch(batch_size, self.latent_size, engine.state.device, batchdata)
         g_output = self.g_inferer(g_input, self.g_network)
 
         # Train Discriminator
@@ -281,8 +278,8 @@ def _iteration(
 
         # Train Generator
         if self.g_update_latents:
-            g_input = self.g_prepare_batch(batch_size, self.latent_shape, batchdata).to(engine.state.device)
-        g_output = self.g_inferer(g_input, self.g_network)
+            g_input = self.g_prepare_batch(batch_size, self.latent_size, engine.state.device, batchdata)
+            g_output = self.g_inferer(g_input, self.g_network)
         self.g_optimizer.zero_grad()
         g_loss = self.g_loss_function(g_output)
         g_loss.backward()

monai/engines/utils.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Optional, Sequence, Tuple, Union
+from typing import Dict, List, Optional, Sequence, Tuple
 
 import torch
 
@@ -74,17 +74,53 @@ def get_devices_spec(devices: Optional[Sequence[torch.device]] = None) -> List[t
     return devices
 
 
-def default_prepare_batch(
-    batchdata: Dict[str, torch.Tensor]
-) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+def default_prepare_batch(batchdata: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     assert isinstance(batchdata, dict), "default prepare_batch expects dictionary input data."
     if CommonKeys.LABEL in batchdata:
         return (batchdata[CommonKeys.IMAGE], batchdata[CommonKeys.LABEL])
-    elif GanKeys.REALS in batchdata:
-        return batchdata[GanKeys.REALS]
     else:
         return (batchdata[CommonKeys.IMAGE], None)
 
 
-def default_make_latent(num_latents: int, latent_size: int, real_data: Optional[torch.Tensor] = None) -> torch.Tensor:
-    return torch.randn(num_latents, latent_size)
+def default_gan_prepare_batch(batchdata: Dict[str, torch.Tensor], device: torch.device) -> torch.Tensor:
+    """
+    Prepares batchdata for GAN Discriminator training. Sends Tensor to executing device.
+
+    Args:
+        batchdata: Dictionary data returned by DataLoader
+        devices: torch.device to store tensor for execution
+
+    Raises:
+        AssertionError: If input data is not a dictionary.
+        RuntimeError: If dictionary does not have CommonKeys.IMAGE or GanKeys.REALS key.
+
+    Returns:
+        Batch size tensor of images on device.
+    """
+    assert isinstance(batchdata, dict), "default prepare_batch expects dictionary input data."
+    if GanKeys.REALS in batchdata:
+        data = batchdata[GanKeys.REALS]
+    elif CommonKeys.IMAGE in batchdata:
+        data = batchdata[CommonKeys.IMAGE]
+    else:
+        raise RuntimeError("default gan_prepare_batch expects '%s' or '%s' key." % (CommonKeys.IMAGE, GanKeys.REALS))
+    return data.to(device)
+
+
+def default_make_latent(
+    num_latents: int, latent_size: int, device: torch.device, batchdata: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """
+    Prepares a latent code for GAN Generator training. Sends Tensor to executing device.
+    If Generator needs additional input from Dataloader, override this func and process batchdata.
+
+    Args:
+        num_latents: number of latent codes to generate (typically batchsize)
+        latent_size: size of latent code for Generator input
+        device: torch.device to store tensor for execution
+        batchdata: Minibatch from dataloader, ignored by default.
+
+    Returns:
+        Randomly generated latent codes.
+    """
+    return torch.randn(num_latents, latent_size).to(device)

monai/networks/layers/factories.py

@@ -218,6 +218,7 @@ def batch_factory(dim: int) -> Type[Union[nn.BatchNorm1d, nn.BatchNorm2d, nn.Bat
 
 Norm.add_factory_callable("group", lambda: nn.modules.GroupNorm)
 Act.add_factory_callable("elu", lambda: nn.modules.ELU)
+Act.add_factory_callable("glu", lambda: nn.modules.GLU)
 Act.add_factory_callable("relu", lambda: nn.modules.ReLU)
 Act.add_factory_callable("leakyrelu", lambda: nn.modules.LeakyReLU)
 Act.add_factory_callable("prelu", lambda: nn.modules.PReLU)

research/radiogenomic-gan-nodule-synthesis/ConvertDicomMask.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+
+import ast
+import os
+
+import nibabel as nib
+import numpy as np
+import pydicom
+from pydicom.data import get_testdata_files
+
+
+def findTag(dataset, tag, range_z):
+    for data_element in dataset:
+        if data_element.VR == "SQ":
+            for sequence_item in data_element.value:
+                findTag(sequence_item, tag, range_z)
+        else:
+            if data_element.name in tag:
+                repr_value = repr(data_element.value)
+                range_z.append(repr_value)
+                # print("{0:s} = {1:s}".format(data_element.name, repr_value))
+                # print(range_z)
+    return range_z
+
+
+def convert_one_pair(imgFile, mskFile, outFile):
+    if os.path.isfile(mskFile):
+            i = 121
+            data = pydicom.dcmread(mskFile)
+            seg = data.pixel_array
+            seg = np.swapaxes(seg, 0, 2)
+
+            image = nib.load(imgFile)
+            affine = image.affine
+            img = image.dataobj
+
+            if np.shape(seg) == np.shape(img):
+                nib.save(nib.Nifti1Image(seg.astype(np.uint8), affine), outFile)
+            elif i == 83:
+                segOut = np.zeros(np.shape(img))
+                segOut[: seg.shape[0], : seg.shape[1], : seg.shape[2]] = seg
+                nib.save(nib.Nifti1Image(segOut.astype(np.uint8), affine), outFile)
+            else:
+                tag = ["Image Position (Patient)"]
+                range_z = findTag(data, tag, [])
+                print("Case", i, "different in dimension")
+                print("Seg Size \n", np.shape(seg))
+                print("Img Size \n", np.shape(img))
+                print("Matrix \n", affine)
+
+                spaZ = affine[2][2]
+                oriZ = affine[2][3]
+                if spaZ < 0:
+                    spaZ = -spaZ
+                    oriZ = -oriZ
+
+                startIdx = range_z[0]
+                startIdx = float(ast.literal_eval(startIdx)[2])
+                startIdx = int((startIdx - oriZ) / spaZ)
+
+                endIdx = range_z[len(range_z) - 1]
+                endIdx = float(ast.literal_eval(endIdx)[2])
+                endIdx = int((endIdx - oriZ) / spaZ)
+
+                print(startIdx, endIdx)
+
+                segOut = np.zeros(np.shape(img))
+                segOut[: seg.shape[0], : seg.shape[1], startIdx : (startIdx + np.shape(seg)[2])] = seg
+                nib.save(nib.Nifti1Image(segOut.astype(np.uint8), affine), outFile)
+
+
+def processFolder():
+    ct = 0
+    for i in range(98, 163):
+        mskFile = "/media/ziyue/Research/Data/NSCLC_Radiogenomics/ImagesTrain/R01-%03d/Mask/000000.dcm" % (i)
+        imgFile = "/media/ziyue/Research/Data/NSCLC_Radiogenomics/ImagesTrain/R01-%03d/image.nii.gz" % (i)
+        outFile = "/media/ziyue/Research/Data/NSCLC_Radiogenomics/ImagesTrain/R01-%03d/mask.nii.gz" % (i)
+        if os.path.isfile(mskFile):
+            ct = ct + 1
+            data = pydicom.dcmread(mskFile)
+            seg = data.pixel_array
+            seg = np.swapaxes(seg, 0, 2)
+
+            image = nib.load(imgFile)
+            affine = image.affine
+            img = image.dataobj
+
+            if np.shape(seg) == np.shape(img):
+                nib.save(nib.Nifti1Image(seg.astype(np.uint8), affine), outFile)
+            elif i == 83:
+                segOut = np.zeros(np.shape(img))
+                segOut[: seg.shape[0], : seg.shape[1], : seg.shape[2]] = seg
+                nib.save(nib.Nifti1Image(segOut.astype(np.uint8), affine), outFile)
+            else:
+                tag = ["Image Position (Patient)"]
+                range_z = findTag(data, tag, [])
+                print("Case", i, "different in dimension")
+                print("Seg Size \n", np.shape(seg))
+                print("Img Size \n", np.shape(img))
+                print("Matrix \n", affine)
+
+                spaZ = affine[2][2]
+                oriZ = affine[2][3]
+                if spaZ < 0:
+                    spaZ = -spaZ
+                    oriZ = -oriZ
+
+                startIdx = range_z[0]
+                startIdx = float(ast.literal_eval(startIdx)[2])
+                startIdx = int((startIdx - oriZ) / spaZ)
+
+                endIdx = range_z[len(range_z) - 1]
+                endIdx = float(ast.literal_eval(endIdx)[2])
+                endIdx = int((endIdx - oriZ) / spaZ)
+
+                print(startIdx, endIdx)
+
+                segOut = np.zeros(np.shape(img))
+                segOut[: seg.shape[0], : seg.shape[1], startIdx : (startIdx + np.shape(seg)[2])] = seg
+                nib.save(nib.Nifti1Image(segOut.astype(np.uint8), affine), outFile)
+
+            print("Total image processed: ", ct)
+
+
+if __name__ == "__main__":
+    # processFolder()
+    imgfile = '/nvdata/NSCLC-Radiogenomics-Fresh/test/120_chest.nii.gz'
+    segfile = '/nvdata/NSCLC-Radiogenomics-Fresh/test/120_seg.dcm'
+    outfile = '/nvdata/NSCLC-Radiogenomics-Fresh/test/120_mask_b.nii.gz'
+    convert_one_pair(imgfile, segfile, outfile)