Semester Project Changes

dataset.py

@@ -59,7 +59,7 @@ def __init__(self,
             target_folder: str,
             input_size: int,
             depth_size: int,
-            input_channel: int = 3,
+            input_channel: int = 4,
             transform=None,
             target_transform=None,
             full_channel_mask=False,
@@ -82,14 +82,24 @@ def pair_file(self):
         target_files = sorted(glob(os.path.join(self.target_folder, '*')))
         pairs = []
         for input_file, target_file in zip(input_files, target_files):
-            assert int("".join(re.findall("\d", input_file))) == int("".join(re.findall("\d", target_file)))
+           # assert int("".join(re.findall("\d", input_file))) == int("".join(re.findall("\d", target_file))) 
             pairs.append((input_file, target_file))
         return pairs
 
     def label2masks(self, masked_img):
         result_img = np.zeros(masked_img.shape + ( self.input_channel - 1,))
-        result_img[masked_img==LabelEnum.BRAINAREA.value, 0] = 1
-        result_img[masked_img==LabelEnum.TUMORAREA.value, 1] = 1
+       # result_img[masked_img==LabelEnum.BACKGROUND.value, 0] = 1
+       # result_img[masked_img==LabelEnum.TUMORAREA.value, 1] = 1
+       # result_img[masked_img==LabelEnum.BS.value, 0] = 1
+       # result_img[masked_img==LabelEnum.CBM.value, 1] = 1
+       # result_img[masked_img==LabelEnum.CSF.value, 2] = 1
+       # result_img[masked_img==LabelEnum.GM.value, 3] = 1
+       # result_img[masked_img==LabelEnum.LV.value, 4] = 1
+       # result_img[masked_img==LabelEnum.SGM.value, 5] = 1
+       # result_img[masked_img==LabelEnum.WM.value, 6] = 1
+        result_img[masked_img==LabelEnum.FLUID.value, 0] = 1
+        result_img[masked_img==LabelEnum.CORTEX.value, 1] = 1
+        result_img[masked_img==LabelEnum.REST.value, 2] = 1
         return result_img
 
     def read_image(self, file_path, pass_scaler=False):
@@ -118,7 +128,7 @@ def resize_img(self, img):
 
     def resize_img_4d(self, input_img):
         h, w, d, c = input_img.shape
-        result_img = np.zeros((self.input_size, self.input_size, self.depth_size, 2))
+        result_img = np.zeros((self.input_size, self.input_size, self.depth_size, 3)) # 2  
         if h != self.input_size or w != self.input_size or d != self.depth_size:
             for ch in range(c):
                 buff = input_img.copy()[..., ch]
@@ -151,7 +161,7 @@ def __getitem__(self, index):
         input_img = self.read_image(input_file, pass_scaler=self.full_channel_mask)
         input_img = self.label2masks(input_img) if self.full_channel_mask else input_img
         input_img = self.resize_img(input_img) if not self.full_channel_mask else self.resize_img_4d(input_img)
-
+   
         target_img = self.read_image(target_file)
         target_img = self.resize_img(target_img)
 

diffusion_model/modules.py

@@ -136,7 +136,7 @@ def forward(self, x):
 
 class GroupNorm32(nn.GroupNorm):
     def forward(self, x):
-        return super().forward(x.float()).type(x.dtype)
+        return super().forward(x.float()).type(x.dtype) # super().forward(x.float()).type(x.dtype) 
 
 class AttentionPool2d(nn.Module):
     """

diffusion_model/trainer.py

@@ -271,7 +271,7 @@ def p_losses(self, x_start, t, condition_tensors = None, noise = None):
         else:
             x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
             x_recon = self.denoise_fn(x_noisy, t)
-
+        
         if self.loss_type == 'l1':
             loss = (noise - x_recon).abs().mean()
         elif self.loss_type == 'l2':
@@ -381,7 +381,7 @@ def load(self, milestone):
     def train(self):
         backwards = partial(loss_backwards, self.fp16)
         start_time = time.time()
-
+        
         while self.step < self.train_num_steps:
             accumulated_loss = []
             for i in range(self.gradient_accumulate_every):
@@ -434,6 +434,7 @@ def train(self):
         print('training completed')
         end_time = time.time()
         execution_time = (end_time - start_time)/3600
+
         self.writer.add_hparams(
             {
                 "lr": self.train_lr,

diffusion_model/unet.py

@@ -262,7 +262,7 @@ def forward(self, x, timesteps, y=None):
 
         hs = []
         emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
-
+        
         if self.num_classes is not None:
             assert y.shape == (x.shape[0],)
             emb = emb + self.label_emb(y)
@@ -286,14 +286,14 @@ def create_model(
     channel_mult="",
     learn_sigma=False,
     class_cond=False,
-    use_checkpoint=False,
-    attention_resolutions="16",
-    num_heads=1,
+    use_checkpoint=False, # False 
+    attention_resolutions="4,8,16", # "16"
+    num_heads=4,
     num_head_channels=-1,
     num_heads_upsample=-1,
     use_scale_shift_norm=False,
     dropout=0,
-    resblock_updown=False,
+    resblock_updown=False, 
     use_fp16=False,
     use_new_attention_order=False,
     in_channels=1,
@@ -304,8 +304,10 @@ def create_model(
             channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
         elif image_size == 256:
             channel_mult = (1, 1, 2, 2, 4, 4)
+        elif image_size == 160:
+            channel_mult = (1, 1, 2, 2, 3, 4) 
         elif image_size == 128:
-            channel_mult = (1, 1, 2, 3, 4)
+            channel_mult = (1, 1, 2, 3, 4) 
         elif image_size == 64:
             channel_mult = (1, 2, 3, 4)
         else:

sample.py

@@ -62,11 +62,20 @@ def resize_img_4d(input_img):
 
 def label2masks(masked_img):
     result_img = np.zeros(masked_img.shape + (in_channels-1,))
-    result_img[masked_img==LabelEnum.BRAINAREA.value, 0] = 1
-    result_img[masked_img==LabelEnum.TUMORAREA.value, 1] = 1
+   # result_img[masked_img==LabelEnum.BRAINAREA.value, 0] = 1
+   # result_img[masked_img==LabelEnum.TUMORAREA.value, 1] = 1
+   # result_img[masked_img==LabelEnum.BS.value, 0] = 1
+   # result_img[masked_img==LabelEnum.CBM.value, 1] = 1
+   # result_img[masked_img==LabelEnum.CSF.value, 2] = 1
+   # result_img[masked_img==LabelEnum.GM.value, 3] = 1
+   # result_img[masked_img==LabelEnum.LV.value, 4] = 1
+   # result_img[masked_img==LabelEnum.SGM.value, 5] = 1
+   # result_img[masked_img==LabelEnum.WM.value, 6] = 1
+    result_img[masked_img==LabelEnum.FLUID.value, 0] = 1
+    result_img[masked_img==LabelEnum.CORTEX.value, 1] = 1
+    result_img[masked_img==LabelEnum.REST.value, 2] = 1
     return result_img
 
-
 input_transform = Compose([
     Lambda(lambda t: torch.tensor(t).float()),
     Lambda(lambda t: (t * 2) - 1),
@@ -83,7 +92,7 @@ def label2masks(masked_img):
     image_size = input_size,
     depth_size = depth_size,
     timesteps = args.timesteps,   # number of steps
-    loss_type = 'L1', 
+    loss_type = 'l1', 
     with_condition=True,
 ).cuda()
 diffusion.load_state_dict(torch.load(weightfile)['ema'])
@@ -130,7 +139,7 @@ def label2masks(masked_img):
             counter = counter + 1
             sampleImage = sampleImages[b][0]
             sampleImage = sampleImage.numpy()
-            sampleImage=sampleImage.reshape(refImg.shape)
+            sampleImage=sampleImage.reshape((input_size, input_size, depth_size))
             nifti_img = nib.Nifti1Image(sampleImage, affine=ref.affine)
             nib.save(nifti_img, os.path.join(img_dir, f'{counter}_{msk_name}'))
             nib.save(ref, os.path.join(msk_dir, f'{counter}_{msk_name}'))

scripts/sample.sh

@@ -1 +1 @@
-python3 ../sample.py --inputfolder ../dataset/whole_head/mask --exportfolder ../exports/ --input_size 128 --depth_size 128 --num_channels 64 --num_res_blocks 1 --batchsize 1 --num_samples 1 --weightfile ../model/model_128.pt
+python3 ../sample.py --inputfolder ../../dataset/testing/mask_preprocessed/ --exportfolder ../../synthesized_samples/Med-DDPM/testing_samples/ --input_size 160 --depth_size 160 --num_channels 64 --num_res_blocks 1 --batchsize 1 --num_samples 1 --num_class_labels 4 --timesteps 500 --weightfile ../checkpoints/feta_ddpm_l1/model-99.pt

scripts/train.sh

@@ -1 +1 @@
-python3 ../train.py --with_condition --inputfolder ../dataset/whole_head/mask/ --targetfolder ../dataset/whole_head/image/ --batchsize 1 --epochs 100000 --input_size 128 --depth_size 128 --num_channels 64 --num_res_blocks 1 --timesteps 250 --save_and_sample_every 1000 --resume_weight ../model/model_128.pt
+python3 ../train.py --with_condition --inputfolder ../../dataset/training/mask_augmented/ --targetfolder ../../dataset/training/image_augmented/ --batchsize 1 --epochs 100000 --input_size 160 --depth_size 160 --num_channels 64 --num_res_blocks 1 --num_class_labels 4 --timesteps 500 --save_and_sample_every 1000

train.py

@@ -1,6 +1,6 @@
 #-*- coding:utf-8 -*-
 # +
-from torchvision.transforms import RandomCrop, Compose, ToPILImage, Resize, ToTensor, Lambda
+from torchvision.transforms import RandomCrop, CenterCrop, RandomRotation, Compose, ToPILImage, Resize, ToTensor, Lambda
 from diffusion_model.trainer import GaussianDiffusion, Trainer
 from diffusion_model.unet import create_model
 from dataset import NiftiImageGenerator, NiftiPairImageGenerator
@@ -26,7 +26,7 @@
 parser.add_argument('--timesteps', type=int, default=250)
 parser.add_argument('--save_and_sample_every', type=int, default=1000)
 parser.add_argument('--with_condition', action='store_true')
-parser.add_argument('-r', '--resume_weight', type=str, default="model/model_128.pt")
+parser.add_argument('-r', '--resume_weight', type=str, default="")
 args = parser.parse_args()
 
 inputfolder = args.inputfolder
@@ -51,7 +51,7 @@
 
 input_transform = Compose([
     Lambda(lambda t: torch.tensor(t).float()),
-    Lambda(lambda t: (t * 2) - 1),
+    Lambda(lambda t: (t * 2) - 1), 
     Lambda(lambda t: t.permute(3, 0, 1, 2)),
     Lambda(lambda t: t.transpose(3, 1)),
 ])
@@ -79,7 +79,6 @@
 in_channels = num_class_labels if with_condition else 1
 out_channels = 1
 
-
 model = create_model(input_size, num_channels, num_res_blocks, in_channels=in_channels, out_channels=out_channels).cuda()
 
 diffusion = GaussianDiffusion(
@@ -105,9 +104,9 @@
     train_batch_size = args.batchsize,
     train_lr = train_lr,
     train_num_steps = args.epochs,         # total training steps
-    gradient_accumulate_every = 2,    # gradient accumulation steps
+    gradient_accumulate_every = 2,    # gradient accumulation steps, 2
     ema_decay = 0.995,                # exponential moving average decay
-    fp16 = False,#True,                       # turn on mixed precision training with apex
+    fp16 = True,#True,                       # turn on mixed precision training with apex
     with_condition=with_condition,
     save_and_sample_every = save_and_sample_every,
 )

utils/dtypes.py

@@ -3,8 +3,18 @@
 
 class LabelEnum(IntEnum):
     BACKGROUND = 0
-    TUMORAREA = 2
-    BRAINAREA = 1
+   # TUMORAREA = 2
+   # BRAINAREA = 1
+   # BS = 1
+   # CBM = 2
+   # CSF = 3
+   # GM = 4
+   # LV = 5
+   # SGM = 6
+   # WM = 7
+    FLUID = 1
+    CORTEX = 2
+    REST = 3
 
 class FilterMethods(Enum):
     CUBIC = "CUBIC"