| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304 | 
							- import numpy as np
 
- import torch
 
- import os
 
- from torch.autograd import Variable
 
- from util.image_pool import ImagePool
 
- from .base_model import BaseModel
 
- from . import networks
 
- class Pix2PixHDModel(BaseModel):
 
-     def name(self):
 
-         return 'Pix2PixHDModel'
 
-     
 
-     def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
 
-         flags = (True, use_gan_feat_loss, use_vgg_loss, True, True)
 
-         def loss_filter(g_gan, g_gan_feat, g_vgg, d_real, d_fake):
 
-             return [l for (l,f) in zip((g_gan,g_gan_feat,g_vgg,d_real,d_fake),flags) if f]
 
-         return loss_filter
 
-     
 
-     def initialize(self, opt):
 
-         BaseModel.initialize(self, opt)
 
-         if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
 
-             torch.backends.cudnn.benchmark = True
 
-         self.isTrain = opt.isTrain
 
-         self.use_features = opt.instance_feat or opt.label_feat
 
-         self.gen_features = self.use_features and not self.opt.load_features
 
-         input_nc = opt.label_nc if opt.label_nc != 0 else opt.input_nc
 
-         ##### define networks        
 
-         # Generator network
 
-         netG_input_nc = input_nc        
 
-         if not opt.no_instance:
 
-             netG_input_nc += 1
 
-         if self.use_features:
 
-             netG_input_nc += opt.feat_num                  
 
-         self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG, 
 
-                                       opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 
 
-                                       opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)        
 
-         # Discriminator network
 
-         if self.isTrain:
 
-             use_sigmoid = opt.no_lsgan
 
-             netD_input_nc = input_nc + opt.output_nc
 
-             if not opt.no_instance:
 
-                 netD_input_nc += 1
 
-             self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid, 
 
-                                           opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
 
-         ### Encoder network
 
-         if self.gen_features:          
 
-             self.netE = networks.define_G(opt.output_nc, opt.feat_num, opt.nef, 'encoder', 
 
-                                           opt.n_downsample_E, norm=opt.norm, gpu_ids=self.gpu_ids)  
 
-         if self.opt.verbose:
 
-                 print('---------- Networks initialized -------------')
 
-         # load networks
 
-         if not self.isTrain or opt.continue_train or opt.load_pretrain:
 
-             pretrained_path = '' if not self.isTrain else opt.load_pretrain
 
-             self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)            
 
-             if self.isTrain:
 
-                 self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)  
 
-             if self.gen_features:
 
-                 self.load_network(self.netE, 'E', opt.which_epoch, pretrained_path)              
 
-         # set loss functions and optimizers
 
-         if self.isTrain:
 
-             if opt.pool_size > 0 and (len(self.gpu_ids)) > 1:
 
-                 raise NotImplementedError("Fake Pool Not Implemented for MultiGPU")
 
-             self.fake_pool = ImagePool(opt.pool_size)
 
-             self.old_lr = opt.lr
 
-             # define loss functions
 
-             self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
 
-             
 
-             self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)   
 
-             self.criterionFeat = torch.nn.L1Loss()
 
-             if not opt.no_vgg_loss:             
 
-                 self.criterionVGG = networks.VGGLoss(self.gpu_ids)
 
-                 
 
-         
 
-             # Names so we can breakout loss
 
-             self.loss_names = self.loss_filter('G_GAN','G_GAN_Feat','G_VGG','D_real', 'D_fake')
 
-             # initialize optimizers
 
-             # optimizer G
 
-             if opt.niter_fix_global > 0:                
 
-                 import sys
 
-                 if sys.version_info >= (3,0):
 
-                     finetune_list = set()
 
-                 else:
 
-                     from sets import Set
 
-                     finetune_list = Set()
 
-                 params_dict = dict(self.netG.named_parameters())
 
-                 params = []
 
-                 for key, value in params_dict.items():       
 
-                     if key.startswith('model' + str(opt.n_local_enhancers)):                    
 
-                         params += [value]
 
-                         finetune_list.add(key.split('.')[0])  
 
-                 print('------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global)
 
-                 print('The layers that are finetuned are ', sorted(finetune_list))                         
 
-             else:
 
-                 params = list(self.netG.parameters())
 
-             if self.gen_features:              
 
-                 params += list(self.netE.parameters())         
 
-             self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))                            
 
-             # optimizer D                        
 
-             params = list(self.netD.parameters())    
 
-             self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
 
-     def encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):             
 
-         if self.opt.label_nc == 0:
 
-             input_label = label_map.data.cuda()
 
-         else:
 
-             # create one-hot vector for label map 
 
-             size = label_map.size()
 
-             oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
 
-             input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
 
-             input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
 
-             if self.opt.data_type == 16:
 
-                 input_label = input_label.half()
 
-         # get edges from instance map
 
-         if not self.opt.no_instance:
 
-             inst_map = inst_map.data.cuda()
 
-             edge_map = self.get_edges(inst_map)
 
-             input_label = torch.cat((input_label, edge_map), dim=1)         
 
-         input_label = Variable(input_label, volatile=infer)
 
-         # real images for training
 
-         if real_image is not None:
 
-             real_image = Variable(real_image.data.cuda())
 
-         # instance map for feature encoding
 
-         if self.use_features:
 
-             # get precomputed feature maps
 
-             if self.opt.load_features:
 
-                 feat_map = Variable(feat_map.data.cuda())
 
-             if self.opt.label_feat:
 
-                 inst_map = label_map.cuda()
 
-         return input_label, inst_map, real_image, feat_map
 
-     def discriminate(self, input_label, test_image, use_pool=False):
 
-         input_concat = torch.cat((input_label, test_image.detach()), dim=1)
 
-         if use_pool:            
 
-             fake_query = self.fake_pool.query(input_concat)
 
-             return self.netD.forward(fake_query)
 
-         else:
 
-             return self.netD.forward(input_concat)
 
-     def forward(self, label, inst, image, feat, infer=False):
 
-         # Encode Inputs
 
-         input_label, inst_map, real_image, feat_map = self.encode_input(label, inst, image, feat)  
 
-         # Fake Generation
 
-         if self.use_features:
 
-             if not self.opt.load_features:
 
-                 feat_map = self.netE.forward(real_image, inst_map)                     
 
-             input_concat = torch.cat((input_label, feat_map), dim=1)                        
 
-         else:
 
-             input_concat = input_label
 
-         fake_image = self.netG.forward(input_concat)
 
-         # Fake Detection and Loss
 
-         pred_fake_pool = self.discriminate(input_label, fake_image, use_pool=True)
 
-         loss_D_fake = self.criterionGAN(pred_fake_pool, False)        
 
-         # Real Detection and Loss        
 
-         pred_real = self.discriminate(input_label, real_image)
 
-         loss_D_real = self.criterionGAN(pred_real, True)
 
-         # GAN loss (Fake Passability Loss)        
 
-         pred_fake = self.netD.forward(torch.cat((input_label, fake_image), dim=1))        
 
-         loss_G_GAN = self.criterionGAN(pred_fake, True)               
 
-         
 
-         # GAN feature matching loss
 
-         loss_G_GAN_Feat = 0
 
-         if not self.opt.no_ganFeat_loss:
 
-             feat_weights = 4.0 / (self.opt.n_layers_D + 1)
 
-             D_weights = 1.0 / self.opt.num_D
 
-             for i in range(self.opt.num_D):
 
-                 for j in range(len(pred_fake[i])-1):
 
-                     loss_G_GAN_Feat += D_weights * feat_weights * \
 
-                         self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
 
-                    
 
-         # VGG feature matching loss
 
-         loss_G_VGG = 0
 
-         if not self.opt.no_vgg_loss:
 
-             loss_G_VGG = self.criterionVGG(fake_image, real_image) * self.opt.lambda_feat
 
-         
 
-         # Only return the fake_B image if necessary to save BW
 
-         return [ self.loss_filter( loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_D_real, loss_D_fake ), None if not infer else fake_image ]
 
-     def inference(self, label, inst, image=None):
 
-         # Encode Inputs        
 
-         image = Variable(image) if image is not None else None
 
-         input_label, inst_map, real_image, _ = self.encode_input(Variable(label), Variable(inst), image, infer=True)
 
-         # Fake Generation
 
-         if self.use_features:
 
-             if self.opt.use_encoded_image:
 
-                 # encode the real image to get feature map
 
-                 feat_map = self.netE.forward(real_image, inst_map)
 
-             else:
 
-                 # sample clusters from precomputed features             
 
-                 feat_map = self.sample_features(inst_map)
 
-             input_concat = torch.cat((input_label, feat_map), dim=1)                        
 
-         else:
 
-             input_concat = input_label        
 
-            
 
-         if torch.__version__.startswith('0.4'):
 
-             with torch.no_grad():
 
-                 fake_image = self.netG.forward(input_concat)
 
-         else:
 
-             fake_image = self.netG.forward(input_concat)
 
-         return fake_image
 
-     def sample_features(self, inst): 
 
-         # read precomputed feature clusters 
 
-         cluster_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, self.opt.cluster_path)        
 
-         features_clustered = np.load(cluster_path, encoding='latin1').item()
 
-         # randomly sample from the feature clusters
 
-         inst_np = inst.cpu().numpy().astype(int)                                      
 
-         feat_map = self.Tensor(inst.size()[0], self.opt.feat_num, inst.size()[2], inst.size()[3])
 
-         for i in np.unique(inst_np):    
 
-             label = i if i < 1000 else i//1000
 
-             if label in features_clustered:
 
-                 feat = features_clustered[label]
 
-                 cluster_idx = np.random.randint(0, feat.shape[0]) 
 
-                                             
 
-                 idx = (inst == int(i)).nonzero()
 
-                 for k in range(self.opt.feat_num):                                    
 
-                     feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k]
 
-         if self.opt.data_type==16:
 
-             feat_map = feat_map.half()
 
-         return feat_map
 
-     def encode_features(self, image, inst):
 
-         image = Variable(image.cuda(), volatile=True)
 
-         feat_num = self.opt.feat_num
 
-         h, w = inst.size()[2], inst.size()[3]
 
-         block_num = 32
 
-         feat_map = self.netE.forward(image, inst.cuda())
 
-         inst_np = inst.cpu().numpy().astype(int)
 
-         feature = {}
 
-         for i in range(self.opt.label_nc):
 
-             feature[i] = np.zeros((0, feat_num+1))
 
-         for i in np.unique(inst_np):
 
-             label = i if i < 1000 else i//1000
 
-             idx = (inst == int(i)).nonzero()
 
-             num = idx.size()[0]
 
-             idx = idx[num//2,:]
 
-             val = np.zeros((1, feat_num+1))                        
 
-             for k in range(feat_num):
 
-                 val[0, k] = feat_map[idx[0], idx[1] + k, idx[2], idx[3]].data[0]            
 
-             val[0, feat_num] = float(num) / (h * w // block_num)
 
-             feature[label] = np.append(feature[label], val, axis=0)
 
-         return feature
 
-     def get_edges(self, t):
 
-         edge = torch.cuda.ByteTensor(t.size()).zero_()
 
-         edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
 
-         edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
 
-         edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
 
-         edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
 
-         if self.opt.data_type==16:
 
-             return edge.half()
 
-         else:
 
-             return edge.float()
 
-     def save(self, which_epoch):
 
-         self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
 
-         self.save_network(self.netD, 'D', which_epoch, self.gpu_ids)
 
-         if self.gen_features:
 
-             self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)
 
-     def update_fixed_params(self):
 
-         # after fixing the global generator for a number of iterations, also start finetuning it
 
-         params = list(self.netG.parameters())
 
-         if self.gen_features:
 
-             params += list(self.netE.parameters())           
 
-         self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
 
-         if self.opt.verbose:
 
-             print('------------ Now also finetuning global generator -----------')
 
-     def update_learning_rate(self):
 
-         lrd = self.opt.lr / self.opt.niter_decay
 
-         lr = self.old_lr - lrd        
 
-         for param_group in self.optimizer_D.param_groups:
 
-             param_group['lr'] = lr
 
-         for param_group in self.optimizer_G.param_groups:
 
-             param_group['lr'] = lr
 
-         if self.opt.verbose:
 
-             print('update learning rate: %f -> %f' % (self.old_lr, lr))
 
-         self.old_lr = lr
 
- class InferenceModel(Pix2PixHDModel):
 
-     def forward(self, inp):
 
-         label, inst = inp
 
-         return self.inference(label, inst)
 
-         
 
 
  |