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

deepliif/models/att_unet.py

@@ -0,0 +1,199 @@
+# adapted from https://github.com/LeeJunHyun/Image_Segmentation/blob/master/network.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+def init_weights(net, init_type='normal', gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)
+
+class conv_block(nn.Module):
+    def __init__(self,ch_in,ch_out,innermost=False,outermost=False):
+        super(conv_block,self).__init__()
+        if outermost:
+            self.conv = nn.Sequential(
+                nn.Conv2d(ch_in, ch_out, kernel_size=4,stride=2,padding=1,bias=True),
+                nn.LeakyReLU(0.2, True),
+            )
+        elif innermost:
+            self.conv = nn.Sequential(
+                nn.Conv2d(ch_in, ch_out, kernel_size=4,stride=2,padding=1,bias=True),
+                nn.ReLU(inplace=True),
+            )
+        else:
+            self.conv = nn.Sequential(
+                nn.Conv2d(ch_in, ch_out, kernel_size=4,stride=2,padding=1,bias=True),
+                nn.BatchNorm2d(ch_out),
+                nn.LeakyReLU(0.2, True),
+            )
+
+    def forward(self,x):
+        x = self.conv(x)
+        return x
+
+class up_conv(nn.Module):
+    def __init__(self,ch_in,ch_out,innermost=False,outermost=False):
+        super(up_conv,self).__init__()
+        use_bias=False
+        if outermost:
+            self.up = nn.Sequential(
+                          nn.ConvTranspose2d(ch_in * 2, ch_out,
+                                        kernel_size=4, stride=2,
+                                        padding=1),
+                          nn.Tanh())
+        elif innermost:
+            self.up = nn.Sequential(
+                          nn.ConvTranspose2d(ch_in, ch_out,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                          nn.BatchNorm2d(ch_out),
+                          nn.ReLU(True))
+        else:
+            self.up = nn.Sequential(
+                          nn.ConvTranspose2d(ch_in * 2, ch_out,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                          nn.BatchNorm2d(ch_out),
+                          nn.ReLU(True))
+                                        
+
+    def forward(self,x):
+        x = self.up(x)
+        return x
+
+
+class Attention_block(nn.Module):
+    def __init__(self,F_g,F_l,F_int):
+        super(Attention_block,self).__init__()
+        self.W_g = nn.Sequential(
+            nn.Conv2d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(F_int)
+            )
+        
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(F_int)
+        )
+
+        self.psi = nn.Sequential(
+            nn.Conv2d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(1),
+            nn.Sigmoid()
+        )
+        
+        self.relu = nn.ReLU(inplace=True)
+        
+    def forward(self,g,x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        psi = self.relu(g1+x1)
+        psi = self.psi(psi)
+
+        return x*psi
+
+
+
+class AttU_Net(nn.Module):
+    def __init__(self,img_ch=3,output_ch=1):
+        super(AttU_Net,self).__init__()
+        
+        self.Conv1 = conv_block(ch_in=img_ch,ch_out=64,outermost=True)
+        self.Conv2 = conv_block(ch_in=64,ch_out=128)
+        self.Conv3 = conv_block(ch_in=128,ch_out=256)
+        self.Conv4 = conv_block(ch_in=256,ch_out=512)
+        self.Conv5 = conv_block(ch_in=512,ch_out=512)
+        self.Conv6 = conv_block(ch_in=512,ch_out=512)
+        self.Conv7 = conv_block(ch_in=512,ch_out=512)
+        self.Conv8 = conv_block(ch_in=512,ch_out=512,innermost=True)
+        #self.Conv9 = conv_block(ch_in=512,ch_out=512,innermost=True)
+        
+        self.Up8 = up_conv(ch_in=512,ch_out=512,innermost=True)
+        self.Att8 = Attention_block(F_g=512,F_l=512,F_int=512)
+        
+        self.Up7 = up_conv(ch_in=512,ch_out=512)
+        self.Att7 = Attention_block(F_g=512,F_l=512,F_int=512)
+        
+        self.Up6 = up_conv(ch_in=512,ch_out=512)
+        self.Att6 = Attention_block(F_g=512,F_l=512,F_int=512)
+        
+        self.Up5 = up_conv(ch_in=512,ch_out=512)
+        self.Att5 = Attention_block(F_g=512,F_l=512,F_int=512)
+        
+        self.Up4 = up_conv(ch_in=512,ch_out=256)
+        self.Att4 = Attention_block(F_g=256,F_l=256,F_int=128)
+        
+        self.Up3 = up_conv(ch_in=256,ch_out=128)
+        self.Att3 = Attention_block(F_g=128,F_l=128,F_int=64)
+        
+        self.Up2 = up_conv(ch_in=128,ch_out=64)
+        self.Att2 = Attention_block(F_g=64,F_l=64,F_int=32)
+        
+        self.Up1 = up_conv(ch_in=64,ch_out=output_ch,outermost=True)
+
+    def forward(self,x):
+        # encoding path
+        x1 = self.Conv1(x)
+        x2 = self.Conv2(x1)
+        x3 = self.Conv3(x2)
+        x4 = self.Conv4(x3)
+        x5 = self.Conv5(x4)
+        x6 = self.Conv6(x5)
+        x7 = self.Conv7(x6)
+        x8 = self.Conv8(x7)
+        #x9 = self.Conv9(x8)
+        
+        #d9 = self.Up
+        d8 = self.Up8(x8)
+        x7 = self.Att8(g=d8,x=x7)
+        d8 = torch.cat((x7,d8),dim=1)
+
+        d7 = self.Up7(d8)
+        x6 = self.Att7(g=d7,x=x6)
+        d7 = torch.cat((x6,d7),dim=1)
+        
+        d6 = self.Up6(d7)
+        x5 = self.Att6(g=d6,x=x5)
+        d6 = torch.cat((x5,d6),dim=1)
+        
+        d5 = self.Up5(d6)
+        x4 = self.Att5(g=d5,x=x4)
+        d5 = torch.cat((x4,d5),dim=1)
+        
+        
+        d4 = self.Up4(d5) # x4: [2, 512, 4, 4], d4: [2, 256, 4, 4]
+        x3 = self.Att4(g=d4,x=x3)
+        d4 = torch.cat((x3,d4),dim=1)
+
+        d3 = self.Up3(d4)
+        x2 = self.Att3(g=d3,x=x2)
+        d3 = torch.cat((x2,d3),dim=1)
+
+        d2 = self.Up2(d3)
+        x1 = self.Att2(g=d2,x=x1)
+        d2 = torch.cat((x1,d2),dim=1)
+
+        d1 = self.Up1(d2)
+
+        return d1
+

deepliif/models/base_model.py

@@ -3,7 +3,7 @@ import torch
 from collections import OrderedDict
 from abc import ABC, abstractmethod
 from . import networks
-from ..util import disable_batchnorm_tracking_stats
+from ..util import disable_batchnorm_tracking_stats, enable_batchnorm_tracking_stats
 from deepliif.util import *
 import itertools
 
@@ -30,7 +30,7 @@ class BaseModel(ABC):
             -- self.loss_names (str list):          specify the training losses that you want to plot and save.
             -- self.model_names (str list):         define networks used in our training.
             -- self.visual_names (str list):        specify the images that you want to display and save.
-            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See CycleGAN_model.py for an example.
         """
         self.opt = opt
         self.gpu_ids = opt.gpu_ids
@@ -81,6 +81,7 @@ class BaseModel(ABC):
         Parameters:
             opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
         """
+        self.opt = opt
         if self.is_train:
             self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
         if not self.is_train or opt.continue_train:
@@ -88,6 +89,17 @@ class BaseModel(ABC):
             self.load_networks(load_suffix)
         self.print_networks(opt.verbose)
 
+    def train(self):
+        """Make models train mode """
+        for name in self.model_names:
+            if isinstance(name, str):
+                if '_' in name:
+                    net = getattr(self, 'net' + name.split('_')[0])[int(name.split('_')[-1]) - 1]
+                else:
+                    net = getattr(self, 'net' + name)
+                net.train()
+                net = enable_batchnorm_tracking_stats(net)
+
     def eval(self):
         """Make models eval mode during test time"""
         for name in self.model_names:
@@ -134,10 +146,21 @@ class BaseModel(ABC):
         for name in self.visual_names:
             if isinstance(name, str):
                 if not hasattr(self, name):
-                    if len(name.split('_')) == 2:
-                        visual_ret[name] = getattr(self, name.split('_')[0])[int(name.split('_')[-1]) -1]
+                    if len(name.split('_')) != 2:
+                        if self.opt.model == 'DeepLIIF':
+                            img_name = name[:-1] + '_' + name[-1]
+                            visual_ret[name] = getattr(self, img_name)
+                        else:
+                            if self.opt.model == 'CycleGAN':
+                                l_output = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])
+                                if len(l_output) > 0:
+                                    visual_ret[name] = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])[int(name.split('_')[-1]) - 1]
+                                else:
+                                    print('No output for',name)
+                            else:
+                                visual_ret[name] = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])[int(name.split('_')[-1]) - 1]
                     else:
-                        visual_ret[name] = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])[int(name.split('_')[-1]) - 1]
+                        visual_ret[name] = getattr(self, name.split('_')[0])[int(name.split('_')[-1]) -1]
                 else:
                     visual_ret[name] = getattr(self, name)
         return visual_ret
@@ -240,6 +263,7 @@ class BaseModel(ABC):
             epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
         """
         for name in self.model_names:
+            
             if isinstance(name, str):
                 load_filename = '%s_net_%s.pth' % (epoch, name)
                 load_path = os.path.join(self.save_dir, load_filename)
@@ -272,9 +296,9 @@ class BaseModel(ABC):
                 if hasattr(state_dict, '_metadata'):
                     del state_dict._metadata
 
-                # patch InstanceNorm checkpoints prior to 0.4
-                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
-                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                # # patch InstanceNorm checkpoints prior to 0.4
+                # for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                #     self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
                 net.load_state_dict(state_dict)
 
     def print_networks(self, verbose):

deepliif/models/networks.py

@@ -3,15 +3,19 @@ import torch.nn as nn
 from torch.nn import init
 import functools
 from torch.optim import lr_scheduler
+
 import os
 
 from torchvision import models
-
+from .att_unet import AttU_Net
 ###############################################################################
 # Helper Functions
 ###############################################################################
 from deepliif.util import util
 
+# as of pytorch 2.4, all optimizers start with an uppercase letter
+OPTIMIZER_MAPPING = {optimizer_name.lower():optimizer_name for optimizer_name in dir(torch.optim) if optimizer_name[0].isupper()} 
+
 
 class Identity(nn.Module):
     def forward(self, x):
@@ -31,12 +35,22 @@ def get_norm_layer(norm_type='instance'):
         norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
     elif norm_type == 'instance':
         norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
-    elif norm_type == 'none':
+    elif norm_type in ['none','spectral']:
         def norm_layer(x): return Identity()
+    # elif norm_type == 'spectral':
+    #     norm_layer = torch.nn.utils.parametrizations.spectral_norm # this needs to be called on nn modules individually, not on nn.Sequential
     else:
         raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
     return norm_layer
 
+def get_optimizer(optimizer_name):
+    try:
+        return getattr(torch.optim, optimizer_name)
+    except:
+        try:
+            return getattr(torch.optim, OPTIMIZER_MAPPING[optimizer_name])
+        except:
+            raise NotImplementedError('optimizer [%s] is not found' % optimizer_name)
 
 def get_scheduler(optimizer, opt):
     """Return a learning rate scheduler
@@ -125,7 +139,9 @@ def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
     return net
 
 
-def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], padding_type='reflect'):
+def define_G(
+  input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], padding_type='reflect', 
+             upsample='convtranspose'):
     """Create a generator
 
     Parameters:
@@ -154,17 +170,22 @@ def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, in
     """
     net = None
     norm_layer = get_norm_layer(norm_type=norm)
+    use_spectral_norm = norm == 'spectral'
 
     if netG == 'resnet_9blocks':
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9, padding_type=padding_type)
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9, 
+                              padding_type=padding_type, upsample=upsample, use_spectral_norm=use_spectral_norm)
     elif netG == 'resnet_6blocks':
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6, padding_type=padding_type)
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6, 
+                              padding_type=padding_type, upsample=upsample, use_spectral_norm=use_spectral_norm)
     elif netG == 'unet_128':
         net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
     elif netG == 'unet_256':
         net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
     elif netG == 'unet_512':
         net = UnetGenerator(input_nc, output_nc, 9, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_512_attention':
+        net = AttU_Net(img_ch=input_nc,output_ch=output_nc)
     else:
         raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
     return init_net(net, init_type, init_gain, gpu_ids)
@@ -202,11 +223,12 @@ def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal'
     """
     net = None
     norm_layer = get_norm_layer(norm_type=norm)
+    use_spectral_norm = norm == 'spectral'
 
     if netD == 'basic':  # default PatchGAN classifier
-        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_spectral_norm=use_spectral_norm)
     elif netD == 'n_layers':  # more options
-        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_spectral_norm=use_spectral_norm)
     elif netD == 'pixel':     # classify if each pixel is real or fake
         net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
     else:
@@ -224,7 +246,7 @@ class GANLoss(nn.Module):
     that has the same size as the input.
     """
 
-    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
+    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0, label_smoothing=0.0):
         """ Initialize the GANLoss class.
 
         Parameters:
@@ -239,6 +261,7 @@ class GANLoss(nn.Module):
         self.register_buffer('real_label', torch.tensor(target_real_label))
         self.register_buffer('fake_label', torch.tensor(target_fake_label))
         self.gan_mode = gan_mode
+        self.label_smoothing = label_smoothing
         if gan_mode == 'lsgan':
             self.loss = nn.MSELoss()
         elif gan_mode == 'vanilla':
@@ -261,9 +284,10 @@ class GANLoss(nn.Module):
 
         if target_is_real:
             target_tensor = self.real_label
+            return target_tensor.expand_as(prediction) * (1 - self.label_smoothing)
         else:
             target_tensor = self.fake_label
-        return target_tensor.expand_as(prediction)
+            return target_tensor.expand_as(prediction) * self.label_smoothing
 
     def __call__(self, prediction, target_is_real, epsilon=1.0):
         """Calculate loss given Discriminator's output and grount truth labels.
@@ -332,9 +356,13 @@ class ResnetGenerator(nn.Module):
     """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
 
     We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    
+    Resize-conv: optional replacement of ConvTranspose2d
+    https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/issues/190#issuecomment-358546675
     """
 
-    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='zero'):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='zero', 
+                 upsample='convtranspose', use_spectral_norm=False):
         """Construct a Resnet-based generator
 
         Parameters:
@@ -355,34 +383,54 @@ class ResnetGenerator(nn.Module):
 
         if padding_type == 'reflect':
             model = [nn.ReflectionPad2d(3),
-                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
-                 norm_layer(ngf),
-                 nn.ReLU(True)]
+                     SpectralNorm(nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                                  use_spectral_norm=use_spectral_norm),
+                     norm_layer(ngf),
+                     nn.ReLU(True)]
         else:
             model = [nn.ZeroPad2d(3),
-                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
-                 norm_layer(ngf),
-                 nn.ReLU(True)]
+                     SpectralNorm(nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                                  use_spectral_norm=use_spectral_norm),
+                     norm_layer(ngf),
+                     nn.ReLU(True)]
 
         n_downsampling = 2
         for i in range(n_downsampling):  # add downsampling layers
             mult = 2 ** i
-            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+            model += [SpectralNorm(nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),use_spectral_norm=use_spectral_norm),
                       norm_layer(ngf * mult * 2),
                       nn.ReLU(True)]
 
         mult = 2 ** n_downsampling
         for i in range(n_blocks):       # add ResNet blocks
 
-            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias,
+                                  use_spectral_norm=use_spectral_norm)]
 
         for i in range(n_downsampling):  # add upsampling layers
             mult = 2 ** (n_downsampling - i)
-            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
-                                         kernel_size=3, stride=2,
-                                         padding=1, output_padding=1,
-                                         bias=use_bias),
-                      norm_layer(int(ngf * mult / 2)),
+            if upsample == 'resize_conv':
+                upsample_layer = [#nn.Upsample(scale_factor = 2, mode='bilinear',align_corners=True),
+                                  nn.Upsample(scale_factor = 2, mode='nearest'),
+                                  nn.ReflectionPad2d(1),
+                                  SpectralNorm(nn.Conv2d(ngf * mult, int(ngf * mult / 2),
+                                            kernel_size=3, stride=1, padding=0),use_spectral_norm=use_spectral_norm)]
+            elif upsample == 'pixel_shuffle':
+                upsample_layer = [SpectralNorm(nn.Conv2d(ngf * mult, int(ngf * mult * 2),use_spectral_norm=use_spectral_norm), 
+                                            kernel_size=3, padding=1),
+                                  nn.PixelShuffle(2),
+                                  nn.ReLU()]
+            elif upsample == 'convtranspose':
+                upsample_layer = [SpectralNorm(nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                                                 kernel_size=3, stride=2,
+                                                                 padding=1, output_padding=1,
+                                                                 bias=use_bias),
+                                              use_spectral_norm=use_spectral_norm)]
+            else:
+                raise Exception(f'upsample layer type {upsample} not implemented')
+                
+            model += upsample_layer
+            model += [norm_layer(int(ngf * mult / 2)),
                       nn.ReLU(True)]
 
         if padding_type == 'reflect':
@@ -390,7 +438,7 @@ class ResnetGenerator(nn.Module):
         else:
             model += [nn.ZeroPad2d(3)]
 
-        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [SpectralNorm(nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),use_spectral_norm=use_spectral_norm)]
         model += [nn.Tanh()]
 
         self.model = nn.Sequential(*model)
@@ -403,7 +451,7 @@ class ResnetGenerator(nn.Module):
 class ResnetBlock(nn.Module):
     """Define a Resnet block"""
 
-    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias, use_spectral_norm=False):
         """Initialize the Resnet block
 
         A resnet block is a conv block with skip connections
@@ -412,9 +460,9 @@ class ResnetBlock(nn.Module):
         Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
         """
         super(ResnetBlock, self).__init__()
-        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias, use_spectral_norm)
 
-    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias, use_spectral_norm):
         """Construct a convolutional block.
 
         Parameters:
@@ -437,7 +485,9 @@ class ResnetBlock(nn.Module):
         else:
             raise NotImplementedError('padding [%s] is not implemented' % padding_type)
 
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        conv_block += [SpectralNorm(nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),use_spectral_norm=use_spectral_norm), 
+                       norm_layer(dim), 
+                       nn.ReLU(True)]
         if use_dropout:
             conv_block += [nn.Dropout(0.5)]
 
@@ -450,7 +500,8 @@ class ResnetBlock(nn.Module):
             p = 1
         else:
             raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+        conv_block += [SpectralNorm(nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),use_spectral_norm=use_spectral_norm), 
+                       norm_layer(dim)]
 
         return nn.Sequential(*conv_block)
 
@@ -565,7 +616,7 @@ class UnetSkipConnectionBlock(nn.Module):
 class NLayerDiscriminator(nn.Module):
     """Defines a PatchGAN discriminator"""
 
-    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_spectral_norm=False):
         """Construct a PatchGAN discriminator
 
         Parameters:
@@ -582,14 +633,15 @@ class NLayerDiscriminator(nn.Module):
 
         kw = 4
         padw = 1
-        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        sequence = [SpectralNorm(nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),use_spectral_norm=use_spectral_norm), 
+                    nn.LeakyReLU(0.2, True)]
         nf_mult = 1
         nf_mult_prev = 1
         for n in range(1, n_layers):  # gradually increase the number of filters
             nf_mult_prev = nf_mult
             nf_mult = min(2 ** n, 8)
             sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                SpectralNorm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),use_spectral_norm=use_spectral_norm),
                 norm_layer(ndf * nf_mult),
                 nn.LeakyReLU(0.2, True)
             ]
@@ -597,12 +649,12 @@ class NLayerDiscriminator(nn.Module):
         nf_mult_prev = nf_mult
         nf_mult = min(2 ** n_layers, 8)
         sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            SpectralNorm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),use_spectral_norm=use_spectral_norm),
             norm_layer(ndf * nf_mult),
             nn.LeakyReLU(0.2, True)
         ]
 
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        sequence += [SpectralNorm(nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw),use_spectral_norm=use_spectral_norm)]  # output 1 channel prediction map
         self.model = nn.Sequential(*sequence)
 
     def forward(self, input):
@@ -687,3 +739,25 @@ class VGGLoss(nn.Module):
         for i in range(len(x_vgg)):
             loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
         return loss
+
+
+class TotalVariationLoss(nn.Module):
+    """
+    Absolute difference for neighbouring pixels (i,j to i+1,j, then i,j to i,j+1), averaged on pixel level
+    """ 
+    def __init__(self):
+        super(TotalVariationLoss, self).__init__()
+    
+    def forward(self, x):
+        tv = torch.abs(x[:,:,1:,:]-x[:,:,:-1,:]).sum() + torch.abs(x[:,:,:,1:]-x[:,:,:,:-1]).sum()
+        return tv / torch.prod(torch.tensor(x.size()))
+
+def SpectralNorm(x, use_spectral_norm=False):
+    """
+    A custom wrapper for nn.utils.parametrizations.spectral_norm,
+    with a flag to turn it on or off.
+    """
+    if use_spectral_norm:
+        return nn.utils.parametrizations.spectral_norm(x)
+    else:
+        return x