WiggleGAN.py

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