import utils, torch, time, os, pickle import numpy as np import torch.nn as nn import torch.cuda as cu import torch.optim as optim import pickle from torchvision import transforms from torchvision.utils import save_image from utils import augmentData, RGBtoL, LtoRGB from PIL import Image from dataloader import dataloader from torch.autograd import Variable import matplotlib.pyplot as plt import random from datetime import date from statistics import mean from architectures import depth_generator_UNet, \ depth_discriminator_noclass_UNet class WiggleGAN(object): def __init__(self, args): # parameters self.epoch = args.epoch self.sample_num = 100 self.nCameras = args.cameras self.batch_size = args.batch_size self.save_dir = args.save_dir self.result_dir = args.result_dir self.dataset = args.dataset self.log_dir = args.log_dir self.gpu_mode = args.gpu_mode self.model_name = args.gan_type self.input_size = args.input_size self.class_num = (args.cameras - 1) * 2 # un calculo que hice en paint self.sample_num = self.class_num ** 2 self.imageDim = args.imageDim self.epochVentaja = args.epochV self.cantImages = args.cIm self.visdom = args.visdom self.lambdaL1 = args.lambdaL1 self.depth = args.depth self.name_wiggle = args.name_wiggle self.clipping = args.clipping self.WGAN = False if (self.clipping > 0): self.WGAN = True self.seed = str(random.randint(0, 99999)) self.seed_load = args.seedLoad self.toLoad = False if (self.seed_load != "-0000"): self.toLoad = True self.zGenFactor = args.zGF self.zDisFactor = args.zDF self.bFactor = args.bF self.CR = False if (self.zGenFactor > 0 or self.zDisFactor > 0 or self.bFactor > 0): self.CR = True self.expandGen = args.expandGen self.expandDis = args.expandDis self.wiggleDepth = args.wiggleDepth self.wiggle = False if (self.wiggleDepth > 0): self.wiggle = True # load dataset self.onlyGen = args.lrD <= 0 if not self.wiggle: self.data_loader = dataloader(self.dataset, self.input_size, self.batch_size, self.imageDim, split='train', trans=not self.CR) self.data_Validation = dataloader(self.dataset, self.input_size, self.batch_size, self.imageDim, split='validation') self.dataprint = self.data_Validation.__iter__().__next__() data = self.data_loader.__iter__().__next__().get('x_im') if not self.onlyGen: self.D = depth_discriminator_noclass_UNet(input_dim=3, output_dim=1, input_shape=data.shape, class_num=self.class_num, expand_net=self.expandDis, depth = self.depth, wgan = self.WGAN) self.D_optimizer = optim.Adam(self.D.parameters(), lr=args.lrD, betas=(args.beta1, args.beta2)) self.data_Test = dataloader(self.dataset, self.input_size, self.batch_size, self.imageDim, split='test') self.dataprint_test = self.data_Test.__iter__().__next__() # networks init self.G = depth_generator_UNet(input_dim=4, output_dim=3, class_num=self.class_num, expand_net=self.expandGen, depth = self.depth) self.G_optimizer = optim.Adam(self.G.parameters(), lr=args.lrG, betas=(args.beta1, args.beta2)) if self.gpu_mode: self.G.cuda() if not self.wiggle and not self.onlyGen: self.D.cuda() self.BCE_loss = nn.BCELoss().cuda() self.CE_loss = nn.CrossEntropyLoss().cuda() self.L1 = nn.L1Loss().cuda() self.MSE = nn.MSELoss().cuda() self.BCEWithLogitsLoss = nn.BCEWithLogitsLoss().cuda() else: self.BCE_loss = nn.BCELoss() self.CE_loss = nn.CrossEntropyLoss() self.MSE = nn.MSELoss() self.L1 = nn.L1Loss() self.BCEWithLogitsLoss = nn.BCEWithLogitsLoss() print('---------- Networks architecture -------------') utils.print_network(self.G) if not self.wiggle and not self.onlyGen: utils.print_network(self.D) print('-----------------------------------------------') temp = torch.zeros((self.class_num, 1)) for i in range(self.class_num): temp[i, 0] = i temp_y = torch.zeros((self.sample_num, 1)) for i in range(self.class_num): temp_y[i * self.class_num: (i + 1) * self.class_num] = temp self.sample_y_ = torch.zeros((self.sample_num, self.class_num)).scatter_(1, temp_y.type(torch.LongTensor), 1) if self.gpu_mode: self.sample_y_ = self.sample_y_.cuda() if (self.toLoad): self.load() def train(self): if self.visdom: random.seed(time.time()) today = date.today() vis = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visValidation = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visEpoch = utils.VisdomLineTwoPlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visImages = utils.VisdomImagePlotter(env_name='Cobo_depth_Images_' + str(today) + '_' + self.seed) visImagesTest = utils.VisdomImagePlotter(env_name='Cobo_depth_ImagesTest_' + str(today) + '_' + self.seed) visLossGTest = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visLossGValidation = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visLossDTest = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) visLossDValidation = utils.VisdomLinePlotter(env_name='Cobo_depth_Train-Plots_' + str(today) + '_' + self.seed) self.train_hist = {} self.epoch_hist = {} self.details_hist = {} self.train_hist['D_loss_train'] = [] self.train_hist['G_loss_train'] = [] self.train_hist['D_loss_Validation'] = [] self.train_hist['G_loss_Validation'] = [] self.train_hist['per_epoch_time'] = [] self.train_hist['total_time'] = [] self.details_hist['G_T_Comp_im'] = [] self.details_hist['G_T_BCE_fake_real'] = [] self.details_hist['G_T_Cycle'] = [] self.details_hist['G_zCR'] = [] self.details_hist['G_V_Comp_im'] = [] self.details_hist['G_V_BCE_fake_real'] = [] self.details_hist['G_V_Cycle'] = [] self.details_hist['D_T_BCE_fake_real_R'] = [] self.details_hist['D_T_BCE_fake_real_F'] = [] self.details_hist['D_zCR'] = [] self.details_hist['D_bCR'] = [] self.details_hist['D_V_BCE_fake_real_R'] = [] self.details_hist['D_V_BCE_fake_real_F'] = [] self.epoch_hist['D_loss_train'] = [] self.epoch_hist['G_loss_train'] = [] self.epoch_hist['D_loss_Validation'] = [] self.epoch_hist['G_loss_Validation'] = [] ##Para poder tomar el promedio por epoch iterIniTrain = 0 iterFinTrain = 0 iterIniValidation = 0 iterFinValidation = 0 maxIter = self.data_loader.dataset.__len__() // self.batch_size maxIterVal = self.data_Validation.dataset.__len__() // self.batch_size if (self.WGAN): one = torch.tensor(1, dtype=torch.float).cuda() mone = one * -1 else: self.y_real_ = torch.ones(self.batch_size, 1) self.y_fake_ = torch.zeros(self.batch_size, 1) if self.gpu_mode: self.y_real_, self.y_fake_ = self.y_real_.cuda(), self.y_fake_.cuda() print('training start!!') start_time = time.time() for epoch in range(self.epoch): if (epoch < self.epochVentaja): ventaja = True else: ventaja = False self.G.train() if not self.onlyGen: self.D.train() epoch_start_time = time.time() # TRAIN!!! for iter, data in enumerate(self.data_loader): x_im = data.get('x_im') x_dep = data.get('x_dep') y_im = data.get('y_im') y_dep = data.get('y_dep') y_ = data.get('y_') # x_im = imagenes normales # x_dep = profundidad de images # y_im = imagen con el angulo cambiado # y_ = angulo de la imagen = tengo que tratar negativos # Aumento mi data if (self.CR): x_im_aug, y_im_aug = augmentData(x_im, y_im) x_im_vanilla = x_im if self.gpu_mode: x_im_aug, y_im_aug = x_im_aug.cuda(), y_im_aug.cuda() if iter >= maxIter: break if self.gpu_mode: x_im, y_, y_im, x_dep, y_dep = x_im.cuda(), y_.cuda(), y_im.cuda(), x_dep.cuda(), y_dep.cuda() # update D network if not ventaja and not self.onlyGen: for p in self.D.parameters(): # reset requires_grad p.requires_grad = True # they are set to False below in netG update self.D_optimizer.zero_grad() # Real Images D_real, D_features_real = self.D(y_im, x_im, y_dep, y_) ## Es la funcion forward `` g(z) x # Fake Images G_, G_dep = self.G( y_, x_im, x_dep) D_fake, D_features_fake = self.D(G_, x_im, G_dep, y_) # Losses # GAN Loss if (self.WGAN): # de WGAN D_loss_real_fake_R = - torch.mean(D_real) D_loss_real_fake_F = torch.mean(D_fake) #D_loss_real_fake_R = - D_loss_real_fake_R_positive else: # de Gan normal D_loss_real_fake_R = self.BCEWithLogitsLoss(D_real, self.y_real_) D_loss_real_fake_F = self.BCEWithLogitsLoss(D_fake, self.y_fake_) D_loss = D_loss_real_fake_F + D_loss_real_fake_R if self.CR: # Fake Augmented Images bCR x_im_aug_bCR, G_aug_bCR = augmentData(x_im_vanilla, G_.data.cpu()) if self.gpu_mode: G_aug_bCR, x_im_aug_bCR = G_aug_bCR.cuda(), x_im_aug_bCR.cuda() D_fake_bCR, D_features_fake_bCR = self.D(G_aug_bCR, x_im_aug_bCR, G_dep, y_) D_real_bCR, D_features_real_bCR = self.D(y_im_aug, x_im_aug, y_dep, y_) # Fake Augmented Images zCR G_aug_zCR, G_dep_aug_zCR = self.G(y_, x_im_aug, x_dep) D_fake_aug_zCR, D_features_fake_aug_zCR = self.D(G_aug_zCR, x_im_aug, G_dep_aug_zCR, y_) # bCR Loss (*) D_loss_real = self.MSE(D_features_real, D_features_real_bCR) D_loss_fake = self.MSE(D_features_fake, D_features_fake_bCR) D_bCR = (D_loss_real + D_loss_fake) * self.bFactor # zCR Loss D_zCR = self.MSE(D_features_fake, D_features_fake_aug_zCR) * self.zDisFactor D_CR_losses = D_bCR + D_zCR #D_CR_losses.backward(retain_graph=True) D_loss += D_CR_losses self.details_hist['D_bCR'].append(D_bCR.detach().item()) self.details_hist['D_zCR'].append(D_zCR.detach().item()) else: self.details_hist['D_bCR'].append(0) self.details_hist['D_zCR'].append(0) self.train_hist['D_loss_train'].append(D_loss.detach().item()) self.details_hist['D_T_BCE_fake_real_R'].append(D_loss_real_fake_R.detach().item()) self.details_hist['D_T_BCE_fake_real_F'].append(D_loss_real_fake_F.detach().item()) if self.visdom: visLossDTest.plot('Discriminator_losses', ['D_T_BCE_fake_real_R','D_T_BCE_fake_real_F', 'D_bCR', 'D_zCR'], 'train', self.details_hist) #if self.WGAN: # D_loss_real_fake_F.backward(retain_graph=True) # D_loss_real_fake_R_positive.backward(mone) #else: # D_loss_real_fake.backward() D_loss.backward() self.D_optimizer.step() #WGAN if (self.WGAN): for p in self.D.parameters(): p.data.clamp_(-self.clipping, self.clipping) #Segun paper si el valor es muy chico lleva al banishing gradient # Si se aplicaria la mejora en las WGANs tendiramos que sacar los batch normalizations de la red # update G network self.G_optimizer.zero_grad() G_, G_dep = self.G(y_, x_im, x_dep) if not ventaja and not self.onlyGen: for p in self.D.parameters(): p.requires_grad = False # to avoid computation # Fake images D_fake, _ = self.D(G_, x_im, G_dep, y_) if (self.WGAN): G_loss_fake = -torch.mean(D_fake) #de WGAN else: G_loss_fake = self.BCEWithLogitsLoss(D_fake, self.y_real_) # loss between images (*) #G_join = torch.cat((G_, G_dep), 1) #y_join = torch.cat((y_im, y_dep), 1) G_loss_Comp = self.L1(G_, y_im) if self.depth: G_loss_Comp += self.L1(G_dep, y_dep) G_loss_Dif_Comp = G_loss_Comp * self.lambdaL1 reverse_y = - y_ + 1 reverse_G, reverse_G_dep = self.G(reverse_y, G_, G_dep) G_loss_Cycle = self.L1(reverse_G, x_im) if self.depth: G_loss_Cycle += self.L1(reverse_G_dep, x_dep) G_loss_Cycle = G_loss_Cycle * self.lambdaL1/2 if (self.CR): # Fake images augmented G_aug, G_dep_aug = self.G(y_, x_im_aug, x_dep) D_fake_aug, _ = self.D(G_aug, x_im, G_dep_aug, y_) if (self.WGAN): G_loss_fake = - (torch.mean(D_fake)+torch.mean(D_fake_aug))/2 else: G_loss_fake = ( self.BCEWithLogitsLoss(D_fake, self.y_real_) + self.BCEWithLogitsLoss(D_fake_aug,self.y_real_)) / 2 # loss between images (*) #y_aug_join = torch.cat((y_im_aug, y_dep), 1) #G_aug_join = torch.cat((G_aug, G_dep_aug), 1) G_loss_Comp_Aug = self.L1(G_aug, y_im_aug) if self.depth: G_loss_Comp_Aug += self.L1(G_dep_aug, y_dep) G_loss_Dif_Comp = (G_loss_Comp + G_loss_Comp_Aug)/2 * self.lambdaL1 G_loss = G_loss_fake + G_loss_Dif_Comp + G_loss_Cycle self.details_hist['G_T_BCE_fake_real'].append(G_loss_fake.detach().item()) self.details_hist['G_T_Comp_im'].append(G_loss_Dif_Comp.detach().item()) self.details_hist['G_T_Cycle'].append(G_loss_Cycle.detach().item()) self.details_hist['G_zCR'].append(0) else: G_loss = self.L1(G_, y_im) if self.depth: G_loss += self.L1(G_dep, y_dep) G_loss = G_loss * self.lambdaL1 self.details_hist['G_T_Comp_im'].append(G_loss.detach().item()) self.details_hist['G_T_BCE_fake_real'].append(0) self.details_hist['G_T_Cycle'].append(0) self.details_hist['G_zCR'].append(0) G_loss.backward() self.G_optimizer.step() self.train_hist['G_loss_train'].append(G_loss.detach().item()) if self.onlyGen: self.train_hist['D_loss_train'].append(0) iterFinTrain += 1 if self.visdom: visLossGTest.plot('Generator_losses', ['G_T_Comp_im', 'G_T_BCE_fake_real', 'G_zCR','G_T_Cycle'], 'train', self.details_hist) vis.plot('loss', ['D_loss_train', 'G_loss_train'], 'train', self.train_hist) ##################Validation#################################### with torch.no_grad(): self.G.eval() if not self.onlyGen: self.D.eval() for iter, data in enumerate(self.data_Validation): # Aumento mi data x_im = data.get('x_im') x_dep = data.get('x_dep') y_im = data.get('y_im') y_dep = data.get('y_dep') y_ = data.get('y_') # x_im = imagenes normales # x_dep = profundidad de images # y_im = imagen con el angulo cambiado # y_ = angulo de la imagen = tengo que tratar negativos # x_im = torch.Tensor(list(x_im)) # x_dep = torch.Tensor(x_dep) # y_im = torch.Tensor(y_im) # print(y_.shape[0]) if iter == maxIterVal: # print ("Break") break # print (y_.type(torch.LongTensor).unsqueeze(1)) # print("y_vec_", y_vec_) # print ("z_", z_) if self.gpu_mode: x_im, y_, y_im, x_dep, y_dep = x_im.cuda(), y_.cuda(), y_im.cuda(), x_dep.cuda(), y_dep.cuda() # D network if not ventaja and not self.onlyGen: # Real Images D_real, _ = self.D(y_im, x_im, y_dep,y_) ## Es la funcion forward `` g(z) x # Fake Images G_, G_dep = self.G(y_, x_im, x_dep) D_fake, _ = self.D(G_, x_im, G_dep, y_) # Losses # GAN Loss if (self.WGAN): # de WGAN D_loss_real_fake_R = - torch.mean(D_real) D_loss_real_fake_F = torch.mean(D_fake) else: # de Gan normal D_loss_real_fake_R = self.BCEWithLogitsLoss(D_real, self.y_real_) D_loss_real_fake_F = self.BCEWithLogitsLoss(D_fake, self.y_fake_) D_loss_real_fake = D_loss_real_fake_F + D_loss_real_fake_R D_loss = D_loss_real_fake self.train_hist['D_loss_Validation'].append(D_loss.item()) self.details_hist['D_V_BCE_fake_real_R'].append(D_loss_real_fake_R.item()) self.details_hist['D_V_BCE_fake_real_F'].append(D_loss_real_fake_F.item()) if self.visdom: visLossDValidation.plot('Discriminator_losses', ['D_V_BCE_fake_real_R','D_V_BCE_fake_real_F'], 'Validation', self.details_hist) # G network G_, G_dep = self.G(y_, x_im, x_dep) if not ventaja and not self.onlyGen: # Fake images D_fake,_ = self.D(G_, x_im, G_dep, y_) #Loss GAN if (self.WGAN): G_loss = -torch.mean(D_fake) # porWGAN else: G_loss = self.BCEWithLogitsLoss(D_fake, self.y_real_) #de GAN NORMAL self.details_hist['G_V_BCE_fake_real'].append(G_loss.item()) #Loss comparation #G_join = torch.cat((G_, G_dep), 1) #y_join = torch.cat((y_im, y_dep), 1) G_loss_Comp = self.L1(G_, y_im) if self.depth: G_loss_Comp += self.L1(G_dep, y_dep) G_loss_Comp = G_loss_Comp * self.lambdaL1 reverse_y = - y_ + 1 reverse_G, reverse_G_dep = self.G(reverse_y, G_, G_dep) G_loss_Cycle = self.L1(reverse_G, x_im) if self.depth: G_loss_Cycle += self.L1(reverse_G_dep, x_dep) G_loss_Cycle = G_loss_Cycle * self.lambdaL1/2 G_loss += G_loss_Comp + G_loss_Cycle self.details_hist['G_V_Comp_im'].append(G_loss_Comp.item()) self.details_hist['G_V_Cycle'].append(G_loss_Cycle.detach().item()) else: G_loss = self.L1(G_, y_im) if self.depth: G_loss += self.L1(G_dep, y_dep) G_loss = G_loss * self.lambdaL1 self.details_hist['G_V_Comp_im'].append(G_loss.item()) self.details_hist['G_V_BCE_fake_real'].append(0) self.details_hist['G_V_Cycle'].append(0) self.train_hist['G_loss_Validation'].append(G_loss.item()) if self.onlyGen: self.train_hist['D_loss_Validation'].append(0) iterFinValidation += 1 if self.visdom: visLossGValidation.plot('Generator_losses', ['G_V_Comp_im', 'G_V_BCE_fake_real','G_V_Cycle'], 'Validation', self.details_hist) visValidation.plot('loss', ['D_loss_Validation', 'G_loss_Validation'], 'Validation', self.train_hist) ##Vis por epoch if ventaja or self.onlyGen: self.epoch_hist['D_loss_train'].append(0) self.epoch_hist['D_loss_Validation'].append(0) else: #inicioTr = (epoch - self.epochVentaja) * (iterFinTrain - iterIniTrain) #inicioTe = (epoch - self.epochVentaja) * (iterFinValidation - iterIniValidation) self.epoch_hist['D_loss_train'].append(mean(self.train_hist['D_loss_train'][iterIniTrain: -1])) self.epoch_hist['D_loss_Validation'].append(mean(self.train_hist['D_loss_Validation'][iterIniValidation: -1])) self.epoch_hist['G_loss_train'].append(mean(self.train_hist['G_loss_train'][iterIniTrain:iterFinTrain])) self.epoch_hist['G_loss_Validation'].append( mean(self.train_hist['G_loss_Validation'][iterIniValidation:iterFinValidation])) if self.visdom: visEpoch.plot('epoch', epoch, ['D_loss_train', 'G_loss_train', 'D_loss_Validation', 'G_loss_Validation'], self.epoch_hist) self.train_hist['D_loss_train'] = self.train_hist['D_loss_train'][-1:] self.train_hist['G_loss_train'] = self.train_hist['G_loss_train'][-1:] self.train_hist['D_loss_Validation'] = self.train_hist['D_loss_Validation'][-1:] self.train_hist['G_loss_Validation'] = self.train_hist['G_loss_Validation'][-1:] self.train_hist['per_epoch_time'] = self.train_hist['per_epoch_time'][-1:] self.train_hist['total_time'] = self.train_hist['total_time'][-1:] self.details_hist['G_T_Comp_im'] = self.details_hist['G_T_Comp_im'][-1:] self.details_hist['G_T_BCE_fake_real'] = self.details_hist['G_T_BCE_fake_real'][-1:] self.details_hist['G_T_Cycle'] = self.details_hist['G_T_Cycle'][-1:] self.details_hist['G_zCR'] = self.details_hist['G_zCR'][-1:] self.details_hist['G_V_Comp_im'] = self.details_hist['G_V_Comp_im'][-1:] self.details_hist['G_V_BCE_fake_real'] = self.details_hist['G_V_BCE_fake_real'][-1:] self.details_hist['G_V_Cycle'] = self.details_hist['G_V_Cycle'][-1:] self.details_hist['D_T_BCE_fake_real_R'] = self.details_hist['D_T_BCE_fake_real_R'][-1:] self.details_hist['D_T_BCE_fake_real_F'] = self.details_hist['D_T_BCE_fake_real_F'][-1:] self.details_hist['D_zCR'] = self.details_hist['D_zCR'][-1:] self.details_hist['D_bCR'] = self.details_hist['D_bCR'][-1:] self.details_hist['D_V_BCE_fake_real_R'] = self.details_hist['D_V_BCE_fake_real_R'][-1:] self.details_hist['D_V_BCE_fake_real_F'] = self.details_hist['D_V_BCE_fake_real_F'][-1:] ##Para poder tomar el promedio por epoch iterIniTrain = 1 iterFinTrain = 1 iterIniValidation = 1 iterFinValidation = 1 self.train_hist['per_epoch_time'].append(time.time() - epoch_start_time) if epoch % 10 == 0: self.save(str(epoch)) with torch.no_grad(): if self.visdom: self.visualize_results(epoch, dataprint=self.dataprint, visual=visImages) self.visualize_results(epoch, dataprint=self.dataprint_test, visual=visImagesTest) else: imageName = self.model_name + '_' + 'Train' + '_' + str(self.seed) + '_' + str(epoch) self.visualize_results(epoch, dataprint=self.dataprint, name= imageName) self.visualize_results(epoch, dataprint=self.dataprint_test, name= imageName) self.train_hist['total_time'].append(time.time() - start_time) print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (np.mean(self.train_hist['per_epoch_time']), self.epoch, self.train_hist['total_time'][0])) print("Training finish!... save training results") self.save() #utils.generate_animation(self.result_dir + '/' + self.dataset + '/' + self.model_name + '/' + self.model_name, # self.epoch) #utils.loss_plot(self.train_hist, os.path.join(self.save_dir, self.dataset, self.model_name), self.model_name) def visualize_results(self, epoch, dataprint, visual="", name= "test"): with torch.no_grad(): self.G.eval() #if not os.path.exists(self.result_dir + '/' + self.dataset + '/' + self.model_name): # os.makedirs(self.result_dir + '/' + self.dataset + '/' + self.model_name) # print("sample z: ",self.sample_z_,"sample y:", self.sample_y_) ##Podria hacer un loop # .zfill(4) #newSample = None #print(dataprint.shape) #newSample = torch.tensor([]) #se que es ineficiente pero lo hago cada 10 epoch nomas newSample = [] iter = 1 for x_im,x_dep in zip(dataprint.get('x_im'), dataprint.get('x_dep')): if (iter > self.cantImages): break #x_im = (x_im + 1) / 2 #imgX = transforms.ToPILImage()(x_im) #imgX.show() x_im_input = x_im.repeat(2, 1, 1, 1) x_dep_input = x_dep.repeat(2, 1, 1, 1) sizeImage = x_im.shape[2] sample_y_ = torch.zeros((self.class_num, 1, sizeImage, sizeImage)) for i in range(self.class_num): if(int(i % self.class_num) == 1): sample_y_[i] = torch.ones(( 1, sizeImage, sizeImage)) if self.gpu_mode: sample_y_, x_im_input, x_dep_input = sample_y_.cuda(), x_im_input.cuda(), x_dep_input.cuda() G_im, G_dep = self.G(sample_y_, x_im_input, x_dep_input) newSample.append(x_im.squeeze(0)) newSample.append(x_dep.squeeze(0).expand(3, -1, -1)) if self.wiggle: im_aux, im_dep_aux = G_im, G_dep for i in range(0, 2): index = i for j in range(0, self.wiggleDepth): # print(i,j) if (j == 0 and i == 1): # para tomar el original im_aux, im_dep_aux = G_im, G_dep newSample.append(G_im.cpu()[0].squeeze(0)) newSample.append(G_im.cpu()[1].squeeze(0)) elif (i == 1): # por el problema de las iteraciones proximas index = 0 # imagen generada x = im_aux[index].unsqueeze(0) x_dep = im_dep_aux[index].unsqueeze(0) y = sample_y_[i].unsqueeze(0) if self.gpu_mode: y, x, x_dep = y.cuda(), x.cuda(), x_dep.cuda() im_aux, im_dep_aux = self.G(y, x, x_dep) newSample.append(im_aux.cpu()[0]) else: newSample.append(G_im.cpu()[0]) newSample.append(G_im.cpu()[1]) newSample.append(G_dep.cpu()[0].expand(3, -1, -1)) newSample.append(G_dep.cpu()[1].expand(3, -1, -1)) # sadadas iter+=1 if self.visdom: visual.plot(epoch, newSample, int(len(newSample) /self.cantImages)) else: utils.save_wiggle(newSample, self.cantImages, name) ##TENGO QUE HACER QUE SAMPLES TENGAN COMO MAXIMO self.class_num * self.class_num # utils.save_images(newSample[:, :, :, :], [image_frame_dim * cantidadIm , image_frame_dim * (self.class_num+2)], # self.result_dir + '/' + self.dataset + '/' + self.model_name + '/' + self.model_name + '_epoch%04d' % epoch + '.png') def show_plot_images(self, images, cols=1, titles=None): """Display a list of images in a single figure with matplotlib. Parameters --------- images: List of np.arrays compatible with plt.imshow. cols (Default = 1): Number of columns in figure (number of rows is set to np.ceil(n_images/float(cols))). titles: List of titles corresponding to each image. Must have the same length as titles. """ # assert ((titles is None) or (len(images) == len(titles))) n_images = len(images) if titles is None: titles = ['Image (%d)' % i for i in range(1, n_images + 1)] fig = plt.figure() for n, (image, title) in enumerate(zip(images, titles)): a = fig.add_subplot(np.ceil(n_images / float(cols)), cols, n + 1) # print(image) image = (image + 1) * 255.0 # print(image) # new_im = Image.fromarray(image) # print(new_im) if image.ndim == 2: plt.gray() # print("spi imshape ", image.shape) plt.imshow(image) a.set_title(title) fig.set_size_inches(np.array(fig.get_size_inches()) * n_images) plt.show() def joinImages(self, data): nData = [] for i in range(self.class_num): nData.append(data) nData = np.array(nData) nData = torch.tensor(nData.tolist()) nData = nData.type(torch.FloatTensor) return nData def save(self, epoch=''): save_dir = os.path.join(self.save_dir, self.dataset, self.model_name) if not os.path.exists(save_dir): os.makedirs(save_dir) torch.save(self.G.state_dict(), os.path.join(save_dir, self.model_name + '_' + self.seed + '_' + epoch + '_G.pkl')) if not self.onlyGen: torch.save(self.D.state_dict(), os.path.join(save_dir, self.model_name + '_' + self.seed + '_' + epoch + '_D.pkl')) with open(os.path.join(save_dir, self.model_name + '_history_ '+self.seed+'.pkl'), 'wb') as f: pickle.dump(self.train_hist, f) def load(self): save_dir = os.path.join(self.save_dir, self.dataset, self.model_name) map_loc=None if not torch.cuda.is_available(): map_loc='cpu' self.G.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_G.pkl'), map_location=map_loc)) if not self.wiggle: self.D.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_D.pkl'), map_location=map_loc)) def wiggleEf(self): seed, epoch = self.seed_load.split('_') if self.visdom: visWiggle = utils.VisdomImagePlotter(env_name='Cobo_depth_wiggle_' + seed) self.visualize_results(epoch=epoch, dataprint=self.dataprint_test, visual=visWiggle) else: self.visualize_results(epoch=epoch, dataprint=self.dataprint_test, visual=None, name = self.name_wiggle) def recreate(self): dataloader_recreate = dataloader(self.dataset, self.input_size, self.batch_size, self.imageDim, split='score') with torch.no_grad(): self.G.eval() accum = 0 for data_batch in dataloader_recreate.__iter__(): #{'x_im': x1, 'x_dep': x1_dep, 'y_im': x2, 'y_dep': x2_dep, 'y_': torch.ones(1, self.imageDim, self.imageDim)} left,left_depth,right,right_depth,direction = data_batch.values() if self.gpu_mode: left,left_depth,right,right_depth,direction = left.cuda(),left_depth.cuda(),right.cuda(),right_depth.cuda(),direction.cuda() G_right, G_right_dep = self.G( direction, left, left_depth) reverse_direction = direction * 0 G_left, G_left_dep = self.G( reverse_direction, right, right_depth) for index in range(0,self.batch_size): image_right = (G_right[index] + 1.0)/2.0 image_right_dep = (G_right_dep[index] + 1.0)/2.0 image_left = (G_left[index] + 1.0)/2.0 image_left_dep = (G_left_dep[index] + 1.0)/2.0 save_image(image_right, os.path.join("results","recreate_dataset","CAM1","n_{num:0{width}}.png".format(num = index+accum, width = 4))) save_image(image_right_dep, os.path.join("results","recreate_dataset","CAM1","d_{num:0{width}}.png".format(num = index+accum, width = 4))) save_image(image_left, os.path.join("results","recreate_dataset","CAM0","n_{num:0{width}}.png".format(num = index+accum, width = 4))) save_image(image_left_dep, os.path.join("results","recreate_dataset","CAM0","d_{num:0{width}}.png".format(num = index+accum, width = 4))) accum+= self.batch_size