Add files via upload

Update README.md
Add files via upload
Delete epoch132_gt.png
Delete epoch132_img.png
Delete epoch132_pred.png
Add files via upload
Update README.md
Update README.md
Update README.md
Update README.md
Update README.md

Author: Nerdyvedi

modules/alphagan-matting/AlphaMatting/image_pool.py

@@ -0,0 +1,41 @@
+import random
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+
+class ImagePool():
+    """
+    This class implements an image buffer that stores previously generated images.
+    This buffer enables us to update discriminators using a history of generated images
+    rather than the ones produced by the latest generators.
+    """
+
+    def __init__(self, pool_size):
+        #pool_size is the size of image buffer , if pool_size = 0, no buffer will be created
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        if self.pool_size == 0:
+            return Variable(images)
+        return_images = []
+        for image in images:
+            image = torch.unsqueeze(image, 0)
+            if self.num_imgs < self.pool_size: #If the buffer is not full, keep inserting images into it
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:     #50% chance, buffer will return a previously stored image and insert the current image into buffer
+                    random_id = random.randint(0, self.pool_size-1)
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:
+                    return_images.append(image) #Another 50% chance, will return the current image 
+        return_images = Variable(torch.cat(return_images, 0)) # collect all images and return
+return return_images

modules/alphagan-matting/AlphaMatting/network.py

@@ -0,0 +1,233 @@
+import torch
+import torch.nn as nn
+
+
+def define_G(which_model_netG, norm ='batch', init_type = 'normal',gpu_ids = [], pretrain = True):
+      
+     netG       = None
+     use_gpu    = len(gpu_ids) > 0
+     norm_layer = get_norm_layer(norm_type = norm)
+
+     if use_gpu:
+        assert(torch.cuda.is_avaialable())
+    
+     netG       = ResnetX(id = 50, gpu_ids = gpu_ids , pretrain = pretrain)
+  
+     if len(gpu_ids) > 0:
+        netG.cuda(gpu_ids[0])
+    
+     if pretrain is True:
+        print('Using pretrained weights')
+   
+     else
+        print('Not using pretrained weights')
+  
+     init_weights(netG, init_type = init_type)
+
+     return netG
+
+def define_D(which_model_netD,norm = 'batch',use_sigmoid = False,init_type = 'normal', gpu_ids = []):
+  
+      netD       = None
+      use_gpu    = len(gpu_ids) > 0
+      norm_layer = get_norm_layer(norm_type = norm)
+     
+      if use_gpu:
+         assert(torch.cuda.is_available())
+
+      netD       = NLayerDiscriminator(4,64,n_layers = 3,norm_layer = norm_layer,use_sigmoid = use_sigmoid, gpu_ids = gpu_ids)
+
+      if use_gpu:
+         netD.cuda(gpu_ids[0])
+
+      init_weights(netD,init_type = init_type)
+      return netD
+
+
+class GANLoss(nn.Module):
+
+      def __init__(self,target_real_label = 1.0, target_fake_label = 0.0,tensor = torch.FloatTensor):
+                
+                  super(GANLoss,self).__init__()
+                  self.real_label = target_real_label
+                  self.fake_label = target_fake_label
+                  self.fake_label_var = None
+                  self.real_label_var = None
+
+                  self.Tensor         = tensor
+
+                  self.loss           = nn.BCELoss()
+
+      def get_target_tensor(self,input,target_is_real):
+                   
+                   target_tensor   = None
+                   if target_is_real:
+                      create_label = ((self.real_label_var is None) or self.real_label_var.numel() !=input.numel())
+
+                      if create_label :
+                          real_tensor = self.Tensor(input.size()).fill(self.real_label)   
+                          self.real_label_var = Variable(real_tensor,requires_grad = False)
+                      target_tensor = self.real_label_var
+
+                   else:
+           
+                      create_label = ((self.fake_label_var is None) or self.fake_label_var.numel() != input.numel())
+                      
+                      if create_label :
+                           fake_tensor = self.Tensor(input.size()).fill_(self.real_label)
+                           self.fake_label_var = Variable(fake_tensor,requires_grad = False)
+                      target_tensor = self.fake_label_var
+           
+                    return target_tensor
+
+class AlphaPredicitionLoss(nn.Module):
+    
+      def __init__(self):
+            
+                   super(AlphaPredictionLoss,self).__init__()
+
+      def forward(self,input,target,trimap):
+                 
+                   #trimap region , 1 in unknown region, 0 in known regions
+                   trimap_weights = torch.where(torch.eq(torch.ge(trimap,0.4),torch.le(trimap,0.6)),torch.ones_like(trimap),torch.ones_like(trimap))
+                   unknown_region_size = trimap_weights.sum()
+                   diff                = torch.sqrt(torch.add(torch.pow(input-target,2),1e-12))
+                   return torch.mul(diff,trimap_weights).sum() / unknown_region_size
+
+
+class CompLoss(nn.Module):
+       
+      def __init__(self):
+             
+                   super(CompLoss,self).__init__()
+
+      def forward(self,input,target,trimap,fg,bg):
+         
+                   trimap_weights = torch.where(torch.eq(torch.ge(trimap,0.4),torch.le(trimap,0.6)),torch.ones_like(trimap),torch.ones_like(trimap))
+                   unknown_region_size = trimap_weights.sum()
+
+                   comp_target         = torch.mul(target,fg) + torch.mul((1.0 - target),bg)
+                   comp_input          = torch.mul(input,fg) + torch.mul((1.0 - input),bg)
+
+                   diff = torch.sqrt(torch.add(torch.pow(comp_input - comp_target, 2), 1e-12))
+                   return torch.mul(diff, trimap_weights).sum() / unknown_region_size 
+
+
+class ResNetX(nn.Module):
+      
+      def __init__(self,gpu_ids,pretrain):
+     
+                  super(ResNetX,self).__init__()
+                  self.encoder  =   ResnetXEncoder(gpu_ids,pretrain)
+                  self.decoder  =   ResnetXDecoder(gpu_ids)
+
+      def forward(self,input):
+ 
+                   #Encoder also gives us the saved pooling indices
+                   x,ind   = self.encoder(input)
+                   x       = self.decoder(x,ind)
+
+
+
+class ResnetXEncoder(nn.Module):
+      #Encoder has the same structure as that of ResNet50, but the last 2 layers are removed
+      def __init__(self,pretrain):
+         super(ResnetXEncoder,self).__init__()
+         
+         resnet   = models.resnet50(pretrained = pretrain)
+         
+         #Removing the last 2 Layers
+         modules  = list(resnet.children())[:-2]
+
+         #to save the pooling indices
+         for m in modules:
+              if 'MaxPool' in m.__class__.__name__:
+                  m.return_indices = True
+
+        # Change input shape of the first convolutional layer
+        # Resnet had 3 channels, but for this task we need 4 channels, as we are also adding the trimap
+        conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        weights = torch.zeros(64, 4, 7, 7)
+        weights[:,:3,:,:] = modules[0].weight.data.view(64, 3, 7, 7)
+        conv1.weight.data.copy_(weights)
+        modules[0] = conv1
+
+        self.pool1  = nn.Sequential(*modules[: 4])
+        self.resnet = nn.Sequential(*modules[4:])
+
+
+        def forward(self,input):
+        
+            x,ind = self.pool1(input)
+            x     = self.resnet(x)
+
+            return x,ind
+
+
+
+class ResnetXDecoder(nn.Module):
+    def __init__(self, gpu_ids=[]):
+        super(ResnetXDecoder, self).__init__()
+        model = [nn.Conv2d(2048, 2048, kernel_size=1, padding=0),
+                 nn.BatchNorm2d(2048),
+                 nn.ReLU(True),
+                 nn.ConvTranspose2d(2048, 1024, kernel_size=1, stride=2, output_padding=1, bias=False),
+                 # nn.ConvTranspose2d(2048, 1024, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+                 nn.BatchNorm2d(1024),
+                 nn.ReLU(True)]
+        model += [nn.Conv2d(1024, 1024, kernel_size=5, padding=2),
+                  nn.BatchNorm2d(1024),
+                  nn.ReLU(True),
+                  nn.ConvTranspose2d(1024, 512, kernel_size=1, stride=2, output_padding=1, bias=False),
+                  # nn.ConvTranspose2d(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+                  nn.BatchNorm2d(512),
+                  nn.ReLU(True)]
+        model += [nn.Conv2d(512, 512, kernel_size=5, padding=2),
+                  nn.BatchNorm2d(512),
+                  nn.ReLU(True),
+                  nn.ConvTranspose2d(512, 256, kernel_size=1, stride=2, output_padding=1, bias=False),
+                  # nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+                  nn.BatchNorm2d(256),
+                  nn.ReLU(True)]
+        model += [nn.Conv2d(256, 256, kernel_size=5, padding=2),
+                  nn.BatchNorm2d(256),
+                  nn.ReLU(True),
+                  nn.Conv2d(256, 64, kernel_size=1, stride=1, bias=False),
+                  nn.BatchNorm2d(64),
+                  nn.ReLU(True)]  
+        model += [nn.Conv2d(64, 64, kernel_size=5, padding=2),
+                  nn.BatchNorm2d(64),
+                  nn.ReLU(True)]
+        self.model1 = nn.Sequential(*model)
+        self.unpool = nn.MaxUnpool2d(kernel_size=2, stride=2)
+        model = [nn.Conv2d(64, 64, kernel_size=5, padding=2),
+                 nn.BatchNorm2d(64),
+                 nn.ReLU(True),
+                 nn.ConvTranspose2d(64, 64, kernel_size=7, stride=2, padding=3, output_padding=1, bias=False),
+                 nn.BatchNorm2d(64),
+                 nn.ReLU(True)]
+        model += [nn.Conv2d(64, 64, kernel_size=5, padding=2),
+                  nn.BatchNorm2d(64),
+                  nn.ReLU(True),
+                  nn.Conv2d(64, 1, kernel_size=5, padding=2),
+                  nn.Sigmoid()]
+        self.model2 = nn.Sequential(*model)
+
+        init_weights(self.model1, 'xavier')
+        init_weights(self.model2, 'xavier')
+
+    def forward(self, input, ind):
+        x = self.model1(input)
+        x = self.unpool(x, ind)
+        x = self.model2(x)
+
+return x
+            
+                  
+
+                      
+      
+
+
+
+

modules/alphagan-matting/AlphaMatting/options/base_options.py

@@ -0,0 +1,68 @@
+import argparse
+import os
+from util import util
+import torch
+
+
+class BaseOptions():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+        self.initialized = False
+
+    def initialize(self):
+        self.parser.add_argument('--dataroot', required=True, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
+        self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
+        self.parser.add_argument('--loadSize', type=int, default=286, help='scale images to this size')
+        self.parser.add_argument('--fineSize', type=int, default=256, help='then crop to this size')
+
+
+        self.parser.add_argument('--which_model_netD', type=str, default='basic', help='selects model to use for netD')
+        self.parser.add_argument('--which_model_netG', type=str, default='resnet_9blocks', help='selects model to use for netG')
+        self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers')
+        self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        self.parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
+        self.parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single]')
+        self.parser.add_argument('--model', type=str, default='cycle_gan',
+                                 help='chooses which model to use. cycle_gan, pix2pix, test')
+        self.parser.add_argument('--which_direction', type=str, default='AtoB', help='AtoB or BtoA')
+        self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data')
+        self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+        self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
+        self.parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
+
+        self.initialized = True
+
+    def parse(self):
+        if not self.initialized:
+            self.initialize()
+        self.opt = self.parser.parse_args()
+        self.opt.isTrain = self.isTrain   # train or test
+
+        str_ids = self.opt.gpu_ids.split(',')
+        self.opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                self.opt.gpu_ids.append(id)
+
+        # set gpu ids
+        if len(self.opt.gpu_ids) > 0:
+            torch.cuda.set_device(self.opt.gpu_ids[0])
+
+        args = vars(self.opt)
+
+        print('------------ Options -------------')
+        for k, v in sorted(args.items()):
+            print('%s: %s' % (str(k), str(v)))
+        print('-------------- End ----------------')
+
+        # save to the disk
+        expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
+        util.mkdirs(expr_dir)
+        file_name = os.path.join(expr_dir, 'opt.txt')
+        with open(file_name, 'wt') as opt_file:
+            opt_file.write('------------ Options -------------\n')
+            for k, v in sorted(args.items()):
+                opt_file.write('%s: %s\n' % (str(k), str(v)))
+            opt_file.write('-------------- End ----------------\n')
+return self.opt

modules/alphagan-matting/AlphaMatting/options/train_options.py

@@ -0,0 +1,28 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+    def initialize(self):
+        BaseOptions.initialize(self)
+
+        self.parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
+        self.parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
+        self.parser.add_argument('--save_epoch_freq', type=int, default=50, help='frequency of saving checkpoints at the end of epochs')
+        self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
+        self.parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
+        self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        self.parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')
+        self.parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
+
+        self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
+
+
+        self.parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')
+
+        self.parser.add_argument('--lr_policy', type=str, default='lambda', help='learning rate policy: lambda|step|plateau')
+        self.parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
+
+        self.parser.add_argument('--no_pretrain', action='store_true', help='initialize weights with pretrained imagenet weights')
+
+