deepliif 1.1.12__tar.gz → 1.1.14__tar.gz

{deepliif-1.1.12/deepliif.egg-info → deepliif-1.1.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: deepliif
-Version: 1.1.12
+Version: 1.1.14
 Summary: DeepLIIF: Deep-Learning Inferred Multiplex Immunofluorescence for Immunohistochemical Image Quantification
 Home-page: https://github.com/nadeemlab/DeepLIIF
 Author: Parmida93
@@ -53,7 +53,7 @@ segmentation.*
 
 © This code is made available for non-commercial academic purposes.
 
-![Version](https://img.shields.io/static/v1?label=latest&message=v1.1.12&color=darkgreen)
+![Version](https://img.shields.io/static/v1?label=latest&message=v1.1.14&color=darkgreen)
 [![Total Downloads](https://static.pepy.tech/personalized-badge/deepliif?period=total&units=international_system&left_color=grey&right_color=blue&left_text=total%20downloads)](https://pepy.tech/project/deepliif?&left_text=totalusers)
 
 ![overview_image](./images/overview.png)*Overview of DeepLIIF pipeline and sample input IHCs (different 

{deepliif-1.1.12 → deepliif-1.1.14}/README.md

@@ -42,7 +42,7 @@ segmentation.*
 
 © This code is made available for non-commercial academic purposes.
 
-![Version](https://img.shields.io/static/v1?label=latest&message=v1.1.12&color=darkgreen)
+![Version](https://img.shields.io/static/v1?label=latest&message=v1.1.14&color=darkgreen)
 [![Total Downloads](https://static.pepy.tech/personalized-badge/deepliif?period=total&units=international_system&left_color=grey&right_color=blue&left_text=total%20downloads)](https://pepy.tech/project/deepliif?&left_text=totalusers)
 
 ![overview_image](./images/overview.png)*Overview of DeepLIIF pipeline and sample input IHCs (different 

{deepliif-1.1.12 → deepliif-1.1.14}/cli.py

@@ -14,6 +14,7 @@ from deepliif.data import create_dataset, transform
 from deepliif.models import init_nets, infer_modalities, infer_results_for_wsi, create_model, postprocess
 from deepliif.util import allowed_file, Visualizer, get_information, test_diff_original_serialized, disable_batchnorm_tracking_stats
 from deepliif.util.util import mkdirs
+from deepliif.util.checks import check_weights
 # from deepliif.util import infer_results_for_wsi
 from deepliif.options import Options, print_options
 
@@ -78,6 +79,7 @@ def cli():
 @click.option('--modalities-no', default=4, type=int, help='number of targets')
 # model parameters
 @click.option('--model', default='DeepLIIF', help='name of model class')
+@click.option('--model-dir-teacher', default='', help='the directory of the teacher model, only applicable if model is DeepLIIFKD')
 @click.option('--seg-weights', default='', type=str, help='weights used to aggregate modality images for the final segmentation image; numbers should add up to 1, and each number corresponds to the modality in order; example: 0.25,0.15,0.25,0.1,0.25')
 @click.option('--loss-weights-g', default='', type=str, help='weights used to aggregate modality-wise losses for the final loss; numbers should add up to 1, and each number corresponds to the modality in order; example: 0.2,0.2,0.2,0.2,0.2')
 @click.option('--loss-weights-d', default='', type=str, help='weights used to aggregate modality-wise losses for the final loss; numbers should add up to 1, and each number corresponds to the modality in order; example: 0.2,0.2,0.2,0.2,0.2')
@@ -193,7 +195,8 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
           verbose, lambda_l1, is_train, display_freq, display_ncols, display_id, display_server, display_env,
           display_port, update_html_freq, print_freq, no_html, save_latest_freq, save_epoch_freq, save_by_iter,
           continue_train, epoch_count, phase, lr_policy, n_epochs, n_epochs_decay, optimizer, beta1, lr_g, lr_d, lr_decay_iters,
-          remote, remote_transfer_cmd, seed, dataset_mode, padding, model, seg_weights, loss_weights_g, loss_weights_d,
+          remote, remote_transfer_cmd, seed, dataset_mode, padding, model, model_dir_teacher,
+          seg_weights, loss_weights_g, loss_weights_d,
           modalities_no, seg_gen, net_ds, net_gs, gan_mode, gan_mode_s, local_rank, with_val, debug, debug_data_size):
     """General-purpose training script for multi-task image-to-image translation.
 
@@ -206,8 +209,8 @@ 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','CycleGAN'], f'model class {model} is not implemented'
-    if model == 'DeepLIIF':
+    assert model in ['DeepLIIF','DeepLIIFExt','SDG','CycleGAN','DeepLIIFKD'], f'model class {model} is not implemented'
+    if model in ['DeepLIIF','DeepLIIFKD']:
         seg_no = 1
     elif model == 'DeepLIIFExt':
         if seg_gen:
@@ -221,6 +224,9 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
     if model == 'CycleGAN':
         dataset_mode = "unaligned"
     
+    if model == 'DeepLIIFKD':
+        assert len(model_dir_teacher) > 0 and os.path.isdir(model_dir_teacher), f'Teacher model directory {model_dir_teacher} is not valid.'
+    
     if optimizer != 'adam':
         print(f'Optimizer torch.optim.{optimizer} is not tested. Be careful about the parameters of the optimizer.')
     
@@ -310,7 +316,7 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
     
     net_gs = net_gs.split(',')
     assert len(net_gs) in [1,seg_no], f'net_gs should contain either 1 architecture for all segmentation generators or the same number of architectures as the number of segmentation generators ({seg_no})'
-    if len(net_gs) == 1 and model == 'DeepLIIF':
+    if len(net_gs) == 1 and model in ['DeepLIIF','DeepLIIFKD']:
         net_gs = net_gs*(modalities_no + seg_no)
     elif len(net_gs) == 1:
         net_gs = net_gs*seg_no
@@ -320,34 +326,21 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
     
     # check seg weights and loss weights
     if len(d_params['seg_weights']) == 0:
-        seg_weights = [0.25,0.15,0.25,0.1,0.25] if d_params['model'] == 'DeepLIIF' else [1 / modalities_no] * modalities_no
+        seg_weights = [0.25,0.15,0.25,0.1,0.25] if d_params['model'] in ['DeepLIIF','DeepLIIFKD'] else [1 / modalities_no] * modalities_no
     else:
         seg_weights = [float(x) for x in seg_weights.split(',')]
     
     if len(d_params['loss_weights_g']) == 0:
-        loss_weights_g = [0.2]*5 if d_params['model'] == 'DeepLIIF' else [1 / modalities_no] * modalities_no
+        loss_weights_g = [0.2]*5 if d_params['model'] in ['DeepLIIF','DeepLIIFKD'] else [1 / modalities_no] * modalities_no
     else:
         loss_weights_g = [float(x) for x in loss_weights_g.split(',')]
     
     if len(d_params['loss_weights_d']) == 0:
-        loss_weights_d = [0.2]*5 if d_params['model'] == 'DeepLIIF' else [1 / modalities_no] * modalities_no
+        loss_weights_d = [0.2]*5 if d_params['model'] in ['DeepLIIF','DeepLIIFKD'] else [1 / modalities_no] * modalities_no
     else:
         loss_weights_d = [float(x) for x in loss_weights_d.split(',')]
     
-    assert sum(seg_weights) == 1, 'seg weights should add up to 1'
-    assert sum(loss_weights_g) == 1, 'loss weights g should add up to 1'
-    assert sum(loss_weights_d) == 1, 'loss weights d should add up to 1'
-    
-    if model == 'DeepLIIF':
-        # +1 because input becomes an additional modality used in generating the final segmentation
-        assert len(seg_weights) == modalities_no+1, 'seg weights should have the same number of elements as number of modalities to be generated'
-        assert len(loss_weights_g) == modalities_no+1, 'loss weights g should have the same number of elements as number of modalities to be generated'
-        assert len(loss_weights_d) == modalities_no+1, 'loss weights d should have the same number of elements as number of modalities to be generated'
-
-    else:
-        assert len(seg_weights) == modalities_no, 'seg weights should have the same number of elements as number of modalities to be generated'
-        assert len(loss_weights_g) == modalities_no, 'loss weights g should have the same number of elements as number of modalities to be generated'
-        assert len(loss_weights_d) == modalities_no, 'loss weights d should have the same number of elements as number of modalities to be generated'
+    check_weights(model, modalities_no, seg_weights, loss_weights_g, loss_weights_d)
 
     d_params['seg_weights'] = seg_weights
     d_params['loss_G_weights'] = loss_weights_g
@@ -373,7 +366,7 @@ def train(dataroot, name, gpu_ids, checkpoints_dir, input_nc, output_nc, ngf, nd
         data_val = [batch for batch in dataset_val]
         click.echo('The number of validation images = %d' % len(dataset_val))
         
-        if model in ['DeepLIIF']: 
+        if model in ['DeepLIIF', 'DeepLIIFKD']: 
             metrics_val = json.load(open(os.path.join(dataset_val.dataset.dir_AB,'metrics.json')))
 
     # create a model given model and other options

{deepliif-1.1.12 → deepliif-1.1.14}/deepliif/data/aligned_dataset.py

@@ -50,7 +50,7 @@ class AlignedDataset(BaseDataset):
         AB = Image.open(AB_path).convert('RGB')
         # split AB image into A and B
         w, h = AB.size
-        if self.model == 'DeepLIIF':
+        if self.model in ['DeepLIIF','DeepLIIFKD']:
             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   
@@ -68,7 +68,7 @@ class AlignedDataset(BaseDataset):
 
         A = A_transform(A)
         B_Array = []
-        if self.model == 'DeepLIIF':
+        if self.model in ['DeepLIIF','DeepLIIFKD']:
             for i in range(1, num_img):
                 B = AB.crop((w2 * i, 0, w2 * (i + 1), h))
                 B = B_transform(B)

deepliif-1.1.14/deepliif/models/DeepLIIFKD_model.py

@@ -0,0 +1,409 @@
+import torch
+from .base_model import BaseModel
+from . import networks
+from .networks import get_optimizer
+from . import init_nets, run_dask, get_opt
+from torch import nn
+
+class DeepLIIFKDModel(BaseModel):
+    """ This class implements the DeepLIIF model, 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)
+        if not hasattr(opt,'net_gs'):
+            opt.net_gs = 'unet_512'
+        
+        self.seg_weights = opt.seg_weights
+        self.loss_G_weights = opt.loss_G_weights
+        self.loss_D_weights = opt.loss_D_weights
+        
+        if not opt.is_train:
+            self.gpu_ids = [] # avoid the models being loaded as DP
+        else:
+            self.gpu_ids = opt.gpu_ids
+
+        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.opt.modalities_no + 1 + 1):
+            self.loss_names.extend([f'G_GAN_{i}', f'G_L1_{i}', f'D_real_{i}', f'D_fake_{i}', f'G_KLDiv_{i}', f'G_KLDiv_5_{i}'])
+            self.visual_names.extend([f'fake_B_{i}', f'fake_B_5_{i}', f'fake_B_{i}_teacher', f'fake_B_5_{i}_teacher', f'real_B_{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.opt.modalities_no + 1):
+                self.model_names.extend(['G' + str(i), 'D' + str(i)])
+
+            for i in range(1, self.opt.modalities_no + 1 + 1):
+                self.model_names.extend(['G5' + str(i), 'D5' + str(i)])
+        else:  # during test time, only load G
+            self.model_names = []
+            for i in range(1, self.opt.modalities_no + 1):
+                self.model_names.extend(['G' + str(i)])
+
+            for i in range(1, self.opt.modalities_no + 1 + 1):
+                self.model_names.extend(['G5' + str(i)])
+
+        # define networks (both generator and discriminator)
+        if isinstance(opt.netG, str):
+            opt.netG = [opt.netG] * 4
+        if isinstance(opt.net_gs, str):
+            opt.net_gs = [opt.net_gs]*5
+
+            
+        self.netG1 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[0], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
+        self.netG2 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[1], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
+        self.netG3 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[2], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
+        self.netG4 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG[3], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, opt.padding)
+
+        # DeepLIIF model currently uses one gs arch because there is only one explicit seg mod output
+        self.netG51 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[0], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+        self.netG52 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[1], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+        self.netG53 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[2], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+        self.netG54 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[3], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+        self.netG55 = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.net_gs[4], opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+
+
+        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.netD1 = networks.define_D(opt.input_nc+opt.output_nc , opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD2 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD3 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD4 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+
+            self.netD51 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD52 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD53 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD54 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD55 = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+                                          
+            # load the teacher model
+            self.opt_teacher = get_opt(opt.model_dir_teacher, mode='test')
+            self.opt_teacher.gpu_ids = opt.gpu_ids # use student's gpu_ids
+            self.nets_teacher = init_nets(opt.model_dir_teacher, eager_mode=True, opt=self.opt_teacher, phase='test')
+            
+
+        if self.is_train:
+            # define loss functions
+            self.criterionGAN_BCE = networks.GANLoss('vanilla').to(self.device)
+            self.criterionGAN_lsgan = networks.GANLoss('lsgan').to(self.device)
+            self.criterionSmoothL1 = torch.nn.SmoothL1Loss()
+
+            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
+            params = list(self.netG1.parameters()) + list(self.netG2.parameters()) + list(self.netG3.parameters()) + list(self.netG4.parameters()) + list(self.netG51.parameters()) + list(self.netG52.parameters()) + list(self.netG53.parameters()) + list(self.netG54.parameters()) + list(self.netG55.parameters())
+            try:
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in generators')
+                self.optimizer_G = get_optimizer(opt.optimizer)(params, lr=opt.lr_g)
+
+            params = list(self.netD1.parameters()) + list(self.netD2.parameters()) + list(self.netD3.parameters()) + list(self.netD4.parameters()) + list(self.netD51.parameters()) + list(self.netD52.parameters()) + list(self.netD53.parameters()) + list(self.netD54.parameters()) + list(self.netD55.parameters())
+            try:
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d, betas=(opt.beta1, 0.999))
+            except:
+                print(f'betas are not used for optimizer torch.optim.{opt.optimizer} in discriminators')
+                self.optimizer_D = get_optimizer(opt.optimizer)(params, lr=opt.lr_d)
+
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+
+            self.criterionVGG = networks.VGGLoss().to(self.device)
+            self.criterionKLDiv = torch.nn.KLDivLoss(reduction='batchmean').to(self.device)
+            self.softmax = torch.nn.Softmax(dim=-1).to(self.device) # apply softmax on the last dim
+            self.logsoftmax = torch.nn.LogSoftmax(dim=-1).to(self.device) # apply log-softmax on the last dim
+
+    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 = input['A'].to(self.device)
+
+        self.real_B_array = input['B']
+        self.real_B_1 = self.real_B_array[0].to(self.device)
+        self.real_B_2 = self.real_B_array[1].to(self.device)
+        self.real_B_3 = self.real_B_array[2].to(self.device)
+        self.real_B_4 = self.real_B_array[3].to(self.device)
+        self.real_B_5 = self.real_B_array[4].to(self.device)
+        self.image_paths = input['A_paths']
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        self.fake_B_1 = self.netG1(self.real_A)   # Hematoxylin image generator
+        self.fake_B_2 = self.netG2(self.real_A)   # mpIF DAPI image generator
+        self.fake_B_3 = self.netG3(self.real_A)   # mpIF Lap2 image generator
+        self.fake_B_4 = self.netG4(self.real_A)   # mpIF Ki67 image generator
+
+        self.fake_B_5_1 = self.netG51(self.real_A)      # Segmentation mask generator from IHC input image
+        self.fake_B_5_2 = self.netG52(self.fake_B_1)    # Segmentation mask generator from Hematoxylin input image
+        self.fake_B_5_3 = self.netG53(self.fake_B_2)    # Segmentation mask generator from mpIF DAPI input image
+        self.fake_B_5_4 = self.netG54(self.fake_B_3)    # Segmentation mask generator from mpIF Lap2 input image
+        self.fake_B_5_5 = self.netG55(self.fake_B_4)    # Segmentation mask generator from mpIF Lap2 input image
+        self.fake_B_5 = torch.stack([torch.mul(self.fake_B_5_1, self.seg_weights[0]),
+                                     torch.mul(self.fake_B_5_2, self.seg_weights[1]),
+                                     torch.mul(self.fake_B_5_3, self.seg_weights[2]),
+                                     torch.mul(self.fake_B_5_4, self.seg_weights[3]),
+                                     torch.mul(self.fake_B_5_5, self.seg_weights[4])]).sum(dim=0)
+        
+        fakes_teacher = run_dask(img=self.real_A, nets=self.nets_teacher, opt=self.opt_teacher, use_dask=False, output_tensor=True)
+        for k,v in fakes_teacher.items():
+            suffix = k[1:] # starts with G
+            suffix = '_'.join(list(suffix)) # 51 -> 5_1
+            setattr(self,f'fake_B_{suffix}_teacher',v)
+
+    def backward_D(self):
+        """Calculate GAN loss for the discriminators"""
+        fake_AB_1 = torch.cat((self.real_A, self.fake_B_1), 1)  # Conditional GANs; feed IHC input and Hematoxtlin output to the discriminator
+        fake_AB_2 = torch.cat((self.real_A, self.fake_B_2), 1)  # Conditional GANs; feed IHC input and mpIF DAPI output to the discriminator
+        fake_AB_3 = torch.cat((self.real_A, self.fake_B_3), 1)  # Conditional GANs; feed IHC input and mpIF Lap2 output to the discriminator
+        fake_AB_4 = torch.cat((self.real_A, self.fake_B_4), 1)  # Conditional GANs; feed IHC input and mpIF Ki67 output to the discriminator
+
+        pred_fake_1 = self.netD1(fake_AB_1.detach())
+        pred_fake_2 = self.netD2(fake_AB_2.detach())
+        pred_fake_3 = self.netD3(fake_AB_3.detach())
+        pred_fake_4 = self.netD4(fake_AB_4.detach())
+
+        fake_AB_5_1 = torch.cat((self.real_A, self.fake_B_5), 1)    # Conditional GANs; feed IHC input and Segmentation mask output to the discriminator
+        fake_AB_5_2 = torch.cat((self.real_B_1, self.fake_B_5), 1)  # Conditional GANs; feed Hematoxylin input and Segmentation mask output to the discriminator
+        fake_AB_5_3 = torch.cat((self.real_B_2, self.fake_B_5), 1)  # Conditional GANs; feed mpIF DAPI input and Segmentation mask output to the discriminator
+        fake_AB_5_4 = torch.cat((self.real_B_3, self.fake_B_5), 1)  # Conditional GANs; feed mpIF Lap2 input and Segmentation mask output to the discriminator
+        fake_AB_5_5 = torch.cat((self.real_B_4, self.fake_B_5), 1)  # Conditional GANs; feed mpIF Lap2 input and Segmentation mask output to the discriminator
+
+        pred_fake_5_1 = self.netD51(fake_AB_5_1.detach())
+        pred_fake_5_2 = self.netD52(fake_AB_5_2.detach())
+        pred_fake_5_3 = self.netD53(fake_AB_5_3.detach())
+        pred_fake_5_4 = self.netD54(fake_AB_5_4.detach())
+        pred_fake_5_5 = self.netD55(fake_AB_5_5.detach())
+
+        pred_fake_5 = torch.stack(
+            [torch.mul(pred_fake_5_1, self.seg_weights[0]),
+             torch.mul(pred_fake_5_2, self.seg_weights[1]),
+             torch.mul(pred_fake_5_3, self.seg_weights[2]),
+             torch.mul(pred_fake_5_4, self.seg_weights[3]),
+             torch.mul(pred_fake_5_5, self.seg_weights[4])]).sum(dim=0)
+
+        self.loss_D_fake_1 = self.criterionGAN_BCE(pred_fake_1, False)
+        self.loss_D_fake_2 = self.criterionGAN_BCE(pred_fake_2, False)
+        self.loss_D_fake_3 = self.criterionGAN_BCE(pred_fake_3, False)
+        self.loss_D_fake_4 = self.criterionGAN_BCE(pred_fake_4, False)
+        self.loss_D_fake_5 = self.criterionGAN_lsgan(pred_fake_5, False)
+
+
+        real_AB_1 = torch.cat((self.real_A, self.real_B_1), 1)
+        real_AB_2 = torch.cat((self.real_A, self.real_B_2), 1)
+        real_AB_3 = torch.cat((self.real_A, self.real_B_3), 1)
+        real_AB_4 = torch.cat((self.real_A, self.real_B_4), 1)
+
+        pred_real_1 = self.netD1(real_AB_1)
+        pred_real_2 = self.netD2(real_AB_2)
+        pred_real_3 = self.netD3(real_AB_3)
+        pred_real_4 = self.netD4(real_AB_4)
+
+        real_AB_5_1 = torch.cat((self.real_A, self.real_B_5), 1)
+        real_AB_5_2 = torch.cat((self.real_B_1, self.real_B_5), 1)
+        real_AB_5_3 = torch.cat((self.real_B_2, self.real_B_5), 1)
+        real_AB_5_4 = torch.cat((self.real_B_3, self.real_B_5), 1)
+        real_AB_5_5 = torch.cat((self.real_B_4, self.real_B_5), 1)
+
+        pred_real_5_1 = self.netD51(real_AB_5_1)
+        pred_real_5_2 = self.netD52(real_AB_5_2)
+        pred_real_5_3 = self.netD53(real_AB_5_3)
+        pred_real_5_4 = self.netD54(real_AB_5_4)
+        pred_real_5_5 = self.netD55(real_AB_5_5)
+
+        pred_real_5 = torch.stack(
+            [torch.mul(pred_real_5_1, self.seg_weights[0]),
+             torch.mul(pred_real_5_2, self.seg_weights[1]),
+             torch.mul(pred_real_5_3, self.seg_weights[2]),
+             torch.mul(pred_real_5_4, self.seg_weights[3]),
+             torch.mul(pred_real_5_5, self.seg_weights[4])]).sum(dim=0)
+
+        self.loss_D_real_1 = self.criterionGAN_BCE(pred_real_1, True)
+        self.loss_D_real_2 = self.criterionGAN_BCE(pred_real_2, True)
+        self.loss_D_real_3 = self.criterionGAN_BCE(pred_real_3, True)
+        self.loss_D_real_4 = self.criterionGAN_BCE(pred_real_4, True)
+        self.loss_D_real_5 = self.criterionGAN_lsgan(pred_real_5, True)
+
+        # combine losses and calculate gradients
+        self.loss_D = (self.loss_D_fake_1 + self.loss_D_real_1) * 0.5 * self.loss_D_weights[0] + \
+                      (self.loss_D_fake_2 + self.loss_D_real_2) * 0.5 * self.loss_D_weights[1] + \
+                      (self.loss_D_fake_3 + self.loss_D_real_3) * 0.5 * self.loss_D_weights[2] + \
+                      (self.loss_D_fake_4 + self.loss_D_real_4) * 0.5 * self.loss_D_weights[3] + \
+                      (self.loss_D_fake_5 + self.loss_D_real_5) * 0.5 * self.loss_D_weights[4]
+
+        self.loss_D.backward()
+
+    def backward_G(self):
+        """Calculate GAN and L1 loss for the generator"""
+
+        fake_AB_1 = torch.cat((self.real_A, self.fake_B_1), 1)
+        fake_AB_2 = torch.cat((self.real_A, self.fake_B_2), 1)
+        fake_AB_3 = torch.cat((self.real_A, self.fake_B_3), 1)
+        fake_AB_4 = torch.cat((self.real_A, self.fake_B_4), 1)
+
+        fake_AB_5_1 = torch.cat((self.real_A, self.fake_B_5), 1)
+        fake_AB_5_2 = torch.cat((self.real_B_1, self.fake_B_5), 1)
+        fake_AB_5_3 = torch.cat((self.real_B_2, self.fake_B_5), 1)
+        fake_AB_5_4 = torch.cat((self.real_B_3, self.fake_B_5), 1)
+        fake_AB_5_5 = torch.cat((self.real_B_4, self.fake_B_5), 1)
+
+        pred_fake_1 = self.netD1(fake_AB_1)
+        pred_fake_2 = self.netD2(fake_AB_2)
+        pred_fake_3 = self.netD3(fake_AB_3)
+        pred_fake_4 = self.netD4(fake_AB_4)
+
+        pred_fake_5_1 = self.netD51(fake_AB_5_1)
+        pred_fake_5_2 = self.netD52(fake_AB_5_2)
+        pred_fake_5_3 = self.netD53(fake_AB_5_3)
+        pred_fake_5_4 = self.netD54(fake_AB_5_4)
+        pred_fake_5_5 = self.netD55(fake_AB_5_5)
+        pred_fake_5 = torch.stack(
+            [torch.mul(pred_fake_5_1, self.seg_weights[0]),
+             torch.mul(pred_fake_5_2, self.seg_weights[1]),
+             torch.mul(pred_fake_5_3, self.seg_weights[2]),
+             torch.mul(pred_fake_5_4, self.seg_weights[3]),
+             torch.mul(pred_fake_5_5, self.seg_weights[4])]).sum(dim=0)
+
+        self.loss_G_GAN_1 = self.criterionGAN_BCE(pred_fake_1, True)
+        self.loss_G_GAN_2 = self.criterionGAN_BCE(pred_fake_2, True)
+        self.loss_G_GAN_3 = self.criterionGAN_BCE(pred_fake_3, True)
+        self.loss_G_GAN_4 = self.criterionGAN_BCE(pred_fake_4, True)
+        self.loss_G_GAN_5 = self.criterionGAN_lsgan(pred_fake_5, True)
+
+        # Second, G(A) = B
+        self.loss_G_L1_1 = self.criterionSmoothL1(self.fake_B_1, self.real_B_1) * self.opt.lambda_L1
+        self.loss_G_L1_2 = self.criterionSmoothL1(self.fake_B_2, self.real_B_2) * self.opt.lambda_L1
+        self.loss_G_L1_3 = self.criterionSmoothL1(self.fake_B_3, self.real_B_3) * self.opt.lambda_L1
+        self.loss_G_L1_4 = self.criterionSmoothL1(self.fake_B_4, self.real_B_4) * self.opt.lambda_L1
+        self.loss_G_L1_5 = self.criterionSmoothL1(self.fake_B_5, self.real_B_5) * self.opt.lambda_L1
+
+        self.loss_G_VGG_1 = self.criterionVGG(self.fake_B_1, self.real_B_1) * self.opt.lambda_feat
+        self.loss_G_VGG_2 = self.criterionVGG(self.fake_B_2, self.real_B_2) * self.opt.lambda_feat
+        self.loss_G_VGG_3 = self.criterionVGG(self.fake_B_3, self.real_B_3) * self.opt.lambda_feat
+        self.loss_G_VGG_4 = self.criterionVGG(self.fake_B_4, self.real_B_4) * self.opt.lambda_feat
+        
+        
+        # .view(1,1,-1) reshapes the input (batch_size, 3, 512, 512) to (batch_size, 1, 3*512*512)
+        # softmax/log-softmax is then applied on the concatenated vector of size (1, 3*512*512)
+        # this normalizes the pixel values across all 3 RGB channels
+        # the resulting vectors are then used to compute KL divergence loss
+        self.loss_G_KLDiv_1 = self.criterionKLDiv(self.logsoftmax(self.fake_B_1.view(1,1,-1)), self.softmax(self.fake_B_1_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_2 = self.criterionKLDiv(self.logsoftmax(self.fake_B_2.view(1,1,-1)), self.softmax(self.fake_B_2_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_3 = self.criterionKLDiv(self.logsoftmax(self.fake_B_3.view(1,1,-1)), self.softmax(self.fake_B_3_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_4 = self.criterionKLDiv(self.logsoftmax(self.fake_B_4.view(1,1,-1)), self.softmax(self.fake_B_4_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5.view(1,1,-1)), self.softmax(self.fake_B_5_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5_1 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5_1.view(1,1,-1)), self.softmax(self.fake_B_5_1_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5_2 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5_2.view(1,1,-1)), self.softmax(self.fake_B_5_2_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5_3 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5_3.view(1,1,-1)), self.softmax(self.fake_B_5_3_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5_4 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5_4.view(1,1,-1)), self.softmax(self.fake_B_5_4_teacher.view(1,1,-1)))
+        self.loss_G_KLDiv_5_5 = self.criterionKLDiv(self.logsoftmax(self.fake_B_5_5.view(1,1,-1)), self.softmax(self.fake_B_5_5_teacher.view(1,1,-1)))
+
+        self.loss_G = (self.loss_G_GAN_1 + self.loss_G_L1_1 + self.loss_G_VGG_1) * self.loss_G_weights[0] + \
+                      (self.loss_G_GAN_2 + self.loss_G_L1_2 + self.loss_G_VGG_2) * self.loss_G_weights[1] + \
+                      (self.loss_G_GAN_3 + self.loss_G_L1_3 + self.loss_G_VGG_3) * self.loss_G_weights[2] + \
+                      (self.loss_G_GAN_4 + self.loss_G_L1_4 + self.loss_G_VGG_4) * self.loss_G_weights[3] + \
+                      (self.loss_G_GAN_5 + self.loss_G_L1_5) * self.loss_G_weights[4] + \
+                      (self.loss_G_KLDiv_1 + self.loss_G_KLDiv_2 + self.loss_G_KLDiv_3 + self.loss_G_KLDiv_4 + \
+                      self.loss_G_KLDiv_5 + self.loss_G_KLDiv_5_1 + self.loss_G_KLDiv_5_2 + self.loss_G_KLDiv_5_3 + \
+                      self.loss_G_KLDiv_5_4 + self.loss_G_KLDiv_5_5) * 10
+
+        # combine loss and calculate gradients
+        # self.loss_G = (self.loss_G_GAN_1 + self.loss_G_L1_1) * self.loss_G_weights[0] + \
+        #               (self.loss_G_GAN_2 + self.loss_G_L1_2) * self.loss_G_weights[1] + \
+        #               (self.loss_G_GAN_3 + self.loss_G_L1_3) * self.loss_G_weights[2] + \
+        #               (self.loss_G_GAN_4 + self.loss_G_L1_4) * self.loss_G_weights[3] + \
+        #               (self.loss_G_GAN_5 + self.loss_G_L1_5) * self.loss_G_weights[4]
+        self.loss_G.backward()
+
+    def optimize_parameters(self):
+        self.forward()                   # compute fake images: G(A)
+        # update D
+        self.set_requires_grad(self.netD1, True)  # enable backprop for D1
+        self.set_requires_grad(self.netD2, True)  # enable backprop for D2
+        self.set_requires_grad(self.netD3, True)  # enable backprop for D3
+        self.set_requires_grad(self.netD4, True)  # enable backprop for D4
+        self.set_requires_grad(self.netD51, True)  # enable backprop for D51
+        self.set_requires_grad(self.netD52, True)  # enable backprop for D52
+        self.set_requires_grad(self.netD53, True)  # enable backprop for D53
+        self.set_requires_grad(self.netD54, True)  # enable backprop for D54
+        self.set_requires_grad(self.netD55, True)  # enable backprop for D54
+
+        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
+        self.set_requires_grad(self.netD1, False)  # D1 requires no gradients when optimizing G1
+        self.set_requires_grad(self.netD2, False)  # D2 requires no gradients when optimizing G2
+        self.set_requires_grad(self.netD3, False)  # D3 requires no gradients when optimizing G3
+        self.set_requires_grad(self.netD4, False)  # D4 requires no gradients when optimizing G4
+        self.set_requires_grad(self.netD51, False)  # D51 requires no gradients when optimizing G51
+        self.set_requires_grad(self.netD52, False)  # D52 requires no gradients when optimizing G52
+        self.set_requires_grad(self.netD53, False)  # D53 requires no gradients when optimizing G53
+        self.set_requires_grad(self.netD54, False)  # D54 requires no gradients when optimizing G54
+        self.set_requires_grad(self.netD55, False)  # D54 requires no gradients when optimizing G54
+
+        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
+    
+    def calculate_losses(self):
+        """
+        Calculate losses but do not optimize parameters. Used in validation loss calculation during training.
+        """
+        
+        self.forward()                   # compute fake images: G(A)
+        # update D
+        self.set_requires_grad(self.netD1, True)  # enable backprop for D1
+        self.set_requires_grad(self.netD2, True)  # enable backprop for D2
+        self.set_requires_grad(self.netD3, True)  # enable backprop for D3
+        self.set_requires_grad(self.netD4, True)  # enable backprop for D4
+        self.set_requires_grad(self.netD51, True)  # enable backprop for D51
+        self.set_requires_grad(self.netD52, True)  # enable backprop for D52
+        self.set_requires_grad(self.netD53, True)  # enable backprop for D53
+        self.set_requires_grad(self.netD54, True)  # enable backprop for D54
+        self.set_requires_grad(self.netD55, True)  # enable backprop for D54
+
+        self.optimizer_D.zero_grad()        # set D's gradients to zero
+        self.backward_D()                # calculate gradients for D
+
+        # update G
+        self.set_requires_grad(self.netD1, False)  # D1 requires no gradients when optimizing G1
+        self.set_requires_grad(self.netD2, False)  # D2 requires no gradients when optimizing G2
+        self.set_requires_grad(self.netD3, False)  # D3 requires no gradients when optimizing G3
+        self.set_requires_grad(self.netD4, False)  # D4 requires no gradients when optimizing G4
+        self.set_requires_grad(self.netD51, False)  # D51 requires no gradients when optimizing G51
+        self.set_requires_grad(self.netD52, False)  # D52 requires no gradients when optimizing G52
+        self.set_requires_grad(self.netD53, False)  # D53 requires no gradients when optimizing G53
+        self.set_requires_grad(self.netD54, False)  # D54 requires no gradients when optimizing G54
+        self.set_requires_grad(self.netD55, False)  # D54 requires no gradients when optimizing G54
+
+        self.optimizer_G.zero_grad()        # set G's gradients to zero
+        self.backward_G()                   # calculate graidents for G
+            

deepliif-1.1.12/deepliif/models/__init__.py → deepliif-1.1.14/deepliif/models/__init__ - different weighted.py

@@ -273,6 +273,27 @@ def run_dask(img, model_path, eager_mode=False, opt=None, seg_only=False):
             'G54': 0.00, # Lap2
             'G55': 0.25, # Marker
         }
+        weights1 = {
+            'G51': 0.5, # IHC
+            'G52': 0.0, # Hema
+            'G53': 0.0, # DAPI
+            'G54': 0.0, # Lap2
+            'G55': 0.5, # Marker
+        }
+        weights2 = {
+            'G51': 0.34, # IHC
+            'G52': 0.00, # Hema
+            'G53': 0.33, # DAPI
+            'G54': 0.00, # Lap2
+            'G55': 0.33, # Marker
+        }
+        weights3 = {
+            'G51': 0.34, # IHC
+            'G52': 0.33, # Hema
+            'G53': 0.00, # DAPI
+            'G54': 0.00, # Lap2
+            'G55': 0.33, # Marker
+        }
 
         seg_map = {'G1': 'G52', 'G2': 'G53', 'G3': 'G54', 'G4': 'G55'}
         if seg_only:
@@ -289,6 +310,9 @@ def run_dask(img, model_path, eager_mode=False, opt=None, seg_only=False):
         segs = compute(lazy_segs)[0]
     
         seg = torch.stack([torch.mul(segs[k], weights[k]) for k in segs.keys()]).sum(dim=0)
+        seg1 = torch.stack([torch.mul(segs[k], weights1[k]) for k in segs.keys()]).sum(dim=0)
+        seg2 = torch.stack([torch.mul(segs[k], weights2[k]) for k in segs.keys()]).sum(dim=0)
+        seg3 = torch.stack([torch.mul(segs[k], weights3[k]) for k in segs.keys()]).sum(dim=0)
     
         if seg_only:
             res = {'G4': tensor_to_pil(gens['G4'])} if 'G4' in gens else {}
@@ -296,6 +320,9 @@ def run_dask(img, model_path, eager_mode=False, opt=None, seg_only=False):
             res = {k: tensor_to_pil(v) for k, v in gens.items()}
             res.update({k: tensor_to_pil(v) for k, v in segs.items()})
         res['G5'] = tensor_to_pil(seg)
+        res['G5a'] = tensor_to_pil(seg1)
+        res['G5b'] = tensor_to_pil(seg2)
+        res['G5c'] = tensor_to_pil(seg3)
     
         return res
     elif opt.model in ['DeepLIIFExt','SDG','CycleGAN']:
@@ -348,6 +375,9 @@ def run_wrapper(tile, run_fn, model_path, eager_mode=False, opt=None, seg_only=F
                     'G3': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
                     'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
                     'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                    'G5a': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                    'G5b': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                    'G5c': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
                     'G51': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
                     'G52': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
                     'G53': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
@@ -409,6 +439,9 @@ def inference(img, tile_size, overlap_size, model_path, use_torchserve=False,
                 'Lap2': results['G3'],
                 'Marker': results['G4'],
                 'Seg': results['G5'],
+                'Seg1': results['G5a'],
+                'Seg2': results['G5b'],
+                'Seg3': results['G5c'],
             }
         
         if return_seg_intermediate and not seg_only: