deepliif 1.1.9__py3-none-any.whl → 1.1.11__py3-none-any.whl

cli.py

@@ -14,7 +14,7 @@ from deepliif.models import init_nets, infer_modalities, infer_results_for_wsi,
 from deepliif.util import allowed_file, Visualizer, get_information, test_diff_original_serialized, disable_batchnorm_tracking_stats
 from deepliif.util.util import mkdirs, check_multi_scale
 # from deepliif.util import infer_results_for_wsi
-from deepliif.options import Options
+from deepliif.options import Options, print_options
 
 import torch.distributed as dist
 
@@ -59,29 +59,6 @@ def set_seed(seed=0,rank=None):
 def ensure_exists(d):
     if not os.path.exists(d):
         os.makedirs(d)
-
-def print_options(opt):
-    """Print and save options
-
-    It will print both current options and default values(if different).
-    It will save options into a text file / [checkpoints_dir] / opt.txt
-    """
-    message = ''
-    message += '----------------- Options ---------------\n'
-    for k, v in sorted(vars(opt).items()):
-        comment = ''
-        message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
-    message += '----------------- End -------------------'
-    print(message)
-
-    # save to the disk
-    if opt.phase == 'train':
-        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
-        mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write(message)
-            opt_file.write('\n')
         
         
 @click.group()
@@ -212,6 +189,18 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
     plot, and save models.The script supports continue/resume training.
     Use '--continue_train' to resume your previous training.
     """
+    assert model in ['DeepLIIF','DeepLIIFExt','SDG'], f'model class {model} is not implemented'
+    if model == 'DeepLIIF':
+        seg_no = 1
+    elif model == 'DeepLIIFExt':
+        if seg_gen:
+            seg_no = modalities_no
+        else:
+            seg_no = 0
+    else: # SDG
+        seg_no = 0
+        seg_gen = False
+    
     d_params = locals()
 
     if gpu_ids and gpu_ids[0] == -1:
@@ -241,11 +230,26 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
         d_params['padding'] = 'zero'
         print('padding type is forced to zero padding, because neither refection pad2d or replication pad2d has a deterministic implementation')
 
+    # infer number of input images
+    dir_data_train = dataroot + '/train'
+    fns = os.listdir(dir_data_train)
+    fns = [x for x in fns if x.endswith('.png')]
+    img = Image.open(f"{dir_data_train}/{fns[0]}")
+    
+    num_img = img.size[0] / img.size[1]
+    assert int(num_img) == num_img, f'img size {img.size[0]} / {img.size[1]} = {num_img} is not an integer'
+    num_img = int(num_img)
+    
+    input_no = num_img - modalities_no - seg_no
+    assert input_no > 0, f'inferred number of input images is {input_no}; should be greater than 0'
+    d_params['input_no'] = input_no
+    d_params['scale_size'] = img.size[1]
+
     # create a dataset given dataset_mode and other options
     # dataset = AlignedDataset(opt)
 
     opt = Options(d_params=d_params)
-    print_options(opt)
+    print_options(opt, save=True)
     
     dataset = create_dataset(opt)
     # get the number of images in the dataset.
@@ -468,28 +472,30 @@ def trainlaunch(**kwargs):
 
 
 @cli.command()
-@click.option('--models-dir', default='./model-server/DeepLIIF_Latest_Model', help='reads models from here')
+@click.option('--model-dir', default='./model-server/DeepLIIF_Latest_Model', help='reads models from here')
 @click.option('--output-dir', help='saves results here.')
-@click.option('--tile-size', type=int, default=None, help='tile size')
-@click.option('--device', default='cpu', type=str, help='device to load model, either cpu or gpu')
+#@click.option('--tile-size', type=int, default=None, help='tile size')
+@click.option('--device', default='cpu', type=str, help='device to load model for the similarity test, either cpu or gpu')
 @click.option('--verbose', default=0, type=int,help='saves results here.')
-def serialize(models_dir, output_dir, tile_size, device, verbose):
+def serialize(model_dir, output_dir, device, verbose):
     """Serialize DeepLIIF models using Torchscript
     """
-    if tile_size is None:
-        tile_size = 512
-    output_dir = output_dir or models_dir
+    #if tile_size is None:
+    #    tile_size = 512
+    output_dir = output_dir or model_dir
     ensure_exists(output_dir)
     
     # copy train_opt.txt to the target location
     import shutil
-    if models_dir != output_dir:
-        shutil.copy(f'{models_dir}/train_opt.txt',f'{output_dir}/train_opt.txt')
+    if model_dir != output_dir:
+        shutil.copy(f'{model_dir}/train_opt.txt',f'{output_dir}/train_opt.txt')
     
-    sample = transform(Image.new('RGB', (tile_size, tile_size)))
+    opt = Options(path_file=os.path.join(model_dir,'train_opt.txt'), mode='test')
+    sample = transform(Image.new('RGB', (opt.scale_size, opt.scale_size)))
+    sample = torch.cat([sample]*opt.input_no, 1)
     
     with click.progressbar(
-            init_nets(models_dir, eager_mode=True, phase='test').items(),
+            init_nets(model_dir, eager_mode=True, phase='test').items(),
             label='Tracing nets',
             item_show_func=lambda n: n[0] if n else n
     ) as bar:
@@ -508,7 +514,7 @@ def serialize(models_dir, output_dir, tile_size, device, verbose):
     
     # test: whether the original and the serialized model produces highly similar predictions
     print('testing similarity between prediction from original vs serialized models...')
-    models_original = init_nets(models_dir,eager_mode=True,phase='test')
+    models_original = init_nets(model_dir,eager_mode=True,phase='test')
     models_serialized = init_nets(output_dir,eager_mode=False,phase='test')
     if device == 'gpu':
         sample = sample.cuda()
@@ -590,11 +596,12 @@ def test(input_dir, output_dir, tile_size, model_dir, gpu_ids, region_size, eage
                         filename.replace('.' + filename.split('.')[-1], f'_{name}.png')
                     ))
 
-                with open(os.path.join(
-                        output_dir,
-                        filename.replace('.' + filename.split('.')[-1], f'.json')
-                ), 'w') as f:
-                    json.dump(scoring, f, indent=2)
+                if scoring is not None:
+                    with open(os.path.join(
+                            output_dir,
+                            filename.replace('.' + filename.split('.')[-1], f'.json')
+                    ), 'w') as f:
+                        json.dump(scoring, f, indent=2)
 
 @cli.command()
 @click.option('--input-dir', type=str, required=True, help='Path to input images')

deepliif/data/aligned_dataset.py

@@ -26,6 +26,8 @@ class AlignedDataset(BaseDataset):
         self.output_nc = opt.input_nc if opt.direction == 'BtoA' else opt.output_nc
         self.no_flip = opt.no_flip
         self.modalities_no = opt.modalities_no
+        self.seg_no = opt.seg_no
+        self.input_no = opt.input_no
         self.seg_gen = opt.seg_gen
         self.load_size = opt.load_size
         self.crop_size = opt.crop_size
@@ -52,6 +54,8 @@ class AlignedDataset(BaseDataset):
             num_img = self.modalities_no + 1 + 1 # +1 for segmentation channel, +1 for input image
         elif self.model == 'DeepLIIFExt':
             num_img = self.modalities_no * 2 + 1 if self.seg_gen else self.modalities_no + 1 # +1 for segmentation channel   
+        elif self.model == 'SDG':
+            num_img = self.modalities_no + self.seg_no + self.input_no
         else:
             raise Exception(f'model class {self.model} does not have corresponding implementation in deepliif/data/aligned_dataset.py')
         w2 = int(w / num_img)
@@ -85,6 +89,19 @@ class AlignedDataset(BaseDataset):
                     BS_Array.append(BS)
 
             return {'A': A, 'B': B_Array, 'BS': BS_Array,'A_paths': AB_path, 'B_paths': AB_path}
+        elif self.model == 'SDG':
+            A_Array = []
+            for i in range(self.input_no):
+                A = AB.crop((w2 * i, 0, w2 * (i+1), h))
+                A = A_transform(A)
+                A_Array.append(A)
+            
+            for i in range(self.input_no, self.input_no + self.modalities_no + 1):
+                B = AB.crop((w2 * i, 0, w2 * (i + 1), h))
+                B = B_transform(B)
+                B_Array.append(B)
+            
+            return {'A': A_Array, 'B': B_Array, 'A_paths': AB_path, 'B_paths': AB_path}
         else:
             raise Exception(f'model class {self.model} does not have corresponding implementation in deepliif/data/aligned_dataset.py')
         

deepliif/models/SDG_model.py

@@ -0,0 +1,189 @@
+import torch
+from .base_model import BaseModel
+from . import networks
+
+
+class SDGModel(BaseModel):
+    """ This class implements the Synthetic Data Generation model (based on DeepLIIFExt), for learning a mapping from input images to modalities given paired data."""
+
+    def __init__(self, opt):
+        """Initialize the DeepLIIF class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseModel.__init__(self, opt)
+
+        self.mod_gen_no = self.opt.modalities_no
+
+        # weights of the modalities in generating segmentation mask
+        self.seg_weights = [0, 0, 0]
+        if opt.seg_gen:
+            self.seg_weights = [0.3] * self.mod_gen_no
+            self.seg_weights[1] = 0.4
+
+        # self.seg_weights = opt.seg_weights
+        # assert len(self.seg_weights) == self.seg_gen_no, 'The number of the segmentation weights (seg_weights) is not equal to the number of target images (modalities_no)!'
+        # print(self.seg_weights)
+        # loss weights in calculating the final loss
+        self.loss_G_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+        self.loss_GS_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+
+        self.loss_D_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+        self.loss_DS_weights = [1 / self.mod_gen_no] * self.mod_gen_no
+
+        self.loss_names = []
+        self.visual_names = ['real_A']
+        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
+        for i in range(1, self.mod_gen_no + 1):
+            self.loss_names.extend(['G_GAN_' + str(i), 'G_L1_' + str(i), 'D_real_' + str(i), 'D_fake_' + str(i)])
+            self.visual_names.extend(['fake_B_' + str(i), 'real_B_' + str(i)])
+
+        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
+        if self.is_train:
+            self.model_names = []
+            for i in range(1, self.mod_gen_no + 1):
+                self.model_names.extend(['G_' + str(i), 'D_' + str(i)])
+
+        else:  # during test time, only load G
+            self.model_names = []
+            for i in range(1, self.mod_gen_no + 1):
+                self.model_names.extend(['G_' + str(i)])
+
+        # define networks (both generator and discriminator)
+        self.netG = [None for _ in range(self.mod_gen_no)]
+        for i in range(self.mod_gen_no):
+            self.netG[i] = networks.define_G(self.opt.input_nc * self.opt.input_no, self.opt.output_nc, self.opt.ngf, self.opt.net_g, self.opt.norm,
+                                             not self.opt.no_dropout, self.opt.init_type, self.opt.init_gain, self.opt.gpu_ids, self.opt.padding)
+            print('***************************************')
+            print(self.opt.input_nc, self.opt.output_nc, self.opt.ngf, self.opt.net_g, self.opt.norm,
+                                             not self.opt.no_dropout, self.opt.init_type, self.opt.init_gain, self.opt.gpu_ids, self.opt.padding)
+            print('***************************************')
+
+        if self.is_train:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
+            self.netD = [None for _ in range(self.mod_gen_no)]
+            for i in range(self.mod_gen_no):
+                self.netD[i] = networks.define_D(self.opt.input_nc * self.opt.input_no + self.opt.output_nc, self.opt.ndf, self.opt.net_d,
+                                                 self.opt.n_layers_D, self.opt.norm, self.opt.init_type, self.opt.init_gain,
+                                                 self.opt.gpu_ids)
+
+        if self.is_train:
+            # define loss functions
+            self.criterionGAN_mod = networks.GANLoss(self.opt.gan_mode).to(self.device)
+            self.criterionGAN_seg = networks.GANLoss(self.opt.gan_mode_s).to(self.device)
+
+            self.criterionSmoothL1 = torch.nn.SmoothL1Loss()
+
+            self.criterionVGG = networks.VGGLoss().to(self.device)
+
+            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
+            params = []
+            for i in range(len(self.netG)):
+                params += list(self.netG[i].parameters())
+            self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+            params = []
+            for i in range(len(self.netD)):
+                params += list(self.netD[i].parameters())
+            self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+
+    def set_input(self, input):
+        """
+        Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        :param input (dict): include the input image and the output modalities
+        """
+        self.real_A_array = input['A']
+        As = [A.to(self.device) for A in self.real_A_array]
+        self.real_A = torch.cat(As, dim=1) # shape: 1, (3 x input_no), 512, 512
+
+        self.real_B_array = input['B']
+        self.real_B = []
+        for i in range(len(self.real_B_array)):
+            self.real_B.append(self.real_B_array[i].to(self.device))
+
+        self.real_concatenated = []
+        self.image_paths = input['A_paths']
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        self.fake_B = []
+        for i in range(self.mod_gen_no):
+            self.fake_B.append(self.netG[i](self.real_A))
+
+
+    def backward_D(self):
+        """Calculate GAN loss for the discriminators"""
+
+        pred_fake = []
+        for i in range(self.mod_gen_no):
+            pred_fake.append(self.netD[i](torch.cat((self.real_A, self.fake_B[i]), 1).detach()))
+
+        self.loss_D_fake = []
+        for i in range(self.mod_gen_no):
+            self.loss_D_fake.append(self.criterionGAN_mod(pred_fake[i], False))
+
+        pred_real = []
+        for i in range(self.mod_gen_no):
+            pred_real.append(self.netD[i](torch.cat((self.real_A, self.real_B[i]), 1)))
+
+        self.loss_D_real = []
+        for i in range(self.mod_gen_no):
+            self.loss_D_real.append(self.criterionGAN_mod(pred_real[i], True))
+
+        # combine losses and calculate gradients
+        # self.loss_D = (self.loss_D_fake[0] + self.loss_D_real[0]) * 0.5 * self.loss_D_weights[0]
+        self.loss_D = torch.tensor(0., device=self.device)
+        for i in range(0, self.mod_gen_no):
+            self.loss_D += (self.loss_D_fake[i] + self.loss_D_real[i]) * 0.5 * self.loss_D_weights[i]
+        self.loss_D.backward()
+
+    def backward_G(self):
+        """Calculate GAN and L1 loss for the generator"""
+        pred_fake = []
+        for i in range(self.mod_gen_no):
+            pred_fake.append(self.netD[i](torch.cat((self.real_A, self.fake_B[i]), 1)))
+
+        self.loss_G_GAN = []
+        self.loss_GS_GAN = []
+        for i in range(self.mod_gen_no):
+            self.loss_G_GAN.append(self.criterionGAN_mod(pred_fake[i], True))
+
+        # Second, G(A) = B
+        self.loss_G_L1 = []
+        self.loss_GS_L1 = []
+        for i in range(self.mod_gen_no):
+            self.loss_G_L1.append(self.criterionSmoothL1(self.fake_B[i], self.real_B[i]) * self.opt.lambda_L1)
+
+        #self.loss_G_VGG = []
+        #for i in range(self.mod_gen_no):
+        #    self.loss_G_VGG.append(self.criterionVGG(self.fake_B[i], self.real_B[i]) * self.opt.lambda_feat)
+
+        # self.loss_G = (self.loss_G_GAN[0] + self.loss_G_L1[0]) * self.loss_G_weights[0]
+        self.loss_G = torch.tensor(0., device=self.device)
+        for i in range(0, self.mod_gen_no):
+            self.loss_G += (self.loss_G_GAN[i] + self.loss_G_L1[i]) * self.loss_G_weights[i]
+        #    self.loss_G += (self.loss_G_GAN[i] + self.loss_G_L1[i] + self.loss_G_VGG[i]) * self.loss_G_weights[i]
+        self.loss_G.backward()
+
+    def optimize_parameters(self):
+        self.forward()  # compute fake images: G(A)
+        # update D
+        for i in range(self.mod_gen_no):
+            self.set_requires_grad(self.netD[i], True)  # enable backprop for D1
+
+        self.optimizer_D.zero_grad()  # set D's gradients to zero
+        self.backward_D()  # calculate gradients for D
+        self.optimizer_D.step()  # update D's weights
+
+        # update G
+        for i in range(self.mod_gen_no):
+            self.set_requires_grad(self.netD[i], False)
+
+        self.optimizer_G.zero_grad()  # set G's gradients to zero
+        self.backward_G()  # calculate graidents for G
+        self.optimizer_G.step()  # udpate G's weights