main.py

# -*- coding: utf-8 -*-
import os
import sys
import time
import numpy as np
import models
from keras.utils import generic_utils
import general_utils
from data_utils import *
from image_history_buffer import *
from IPython import display
from models import *
from additional_models import *
from collections import deque
from ImageDataGenerator import *
from time import sleep

class _GAN:
    def __init__(self, gen, disc, disc_entropy,DCGAN,GenClass,classificator, batch_size, img_source_dim,
                  img_dest_dim, noise_dim, noise_scale, lr_D, lr_G, deterministic, inject_noise, model, lsmooth,
                 img_buffer, datagen, disc_type, data_aug, n_classes, disc_iters,name, dir ):
        self.generator_model = gen
        self.discriminator_model = disc
        self.discriminator2 = disc_entropy
        self.DCGAN_model = DCGAN
        self.GenClass_model = GenClass
        self.classificator_model = classificator
        self.batch_size = batch_size
        self.img_source_dim = img_source_dim
        self.img_dest_dim = img_dest_dim
        self.bn_mode = 2
        self.noise_dim = noise_dim
        self.noise_scale = noise_scale
        self.lr_D = lr_D
        self.lr_G = lr_G
        self.deterministic = deterministic
        self.inject_noise = inject_noise
        self.img_buffer = img_buffer
        self.datagen = datagen
        self.model = model
        self.lsmooth = lsmooth
        self.disc_type = disc_type
        self.dir=dir
        self.disc_iters=disc_iters
        self.data_aug = data_aug
        self.img_dim = img_dest_dim
        self.n_classes = n_classes
        self.name = name

def load_data(img_dim, image_dim_ordering, dset):
    # Load and normalize data
    if dset == "mnistM":
        X_source_train, Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(
            img_dim, image_dim_ordering, dset='mnist')
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(
            img_dim, image_dim_ordering, dset='mnistM')
    elif dset == "svhn_to_mnist":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist')
    elif dset == "svhn_to_mnist32":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
    elif dset == "svhn_to_mnist32gray":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32gray')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
    elif dset == "mnist_to_svhn32gray":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32gray')
    elif dset == "mnist_to_svhn32":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32')
    elif dset == "mnist_to_usps":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='usps')
    elif dset == "usps_to_mnist":
        X_source_train,Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(img_dim, image_dim_ordering,dset='usps')
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist')
    elif dset == "MnistMtoMnist":
        X_source_train, Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(
            img_dim, image_dim_ordering, dset='mnist')
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(
            img_dim, image_dim_ordering, dset='mnistM')
    elif dset == "OfficeAmazonToDslr":
        X_source_train, Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(
            img_dim, image_dim_ordering, dset='OfficeAmazon')
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(
            img_dim, image_dim_ordering, dset='OfficeDslr')
    elif dset == "OfficeAmazonToWebcam":
        X_source_train, Y_source_train, X_source_test, Y_source_test, n_classes1 = load_image_dataset(
            img_dim, image_dim_ordering, dset='OfficeAmazon')
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(
            img_dim, image_dim_ordering, dset='OfficeWebcam')
    else:
        print "dataset not supported"
    if n_classes1 != n_classes2:  # sanity check
        print "number of classes mismatch between source and dest domains"
    n_classes = n_classes1
    img_source_dim = X_source_train.shape[-3:]  # is it backend agnostic?
    img_dest_dim = X_dest_train.shape[-3:]
    if (dset == "mnist_to_usps") or (dset == "usps_to_mnist"):
        X_source=X_source_train
        Y_source=Y_source_train
        X_dest=X_dest_train
        Y_dest=Y_dest_train
    else:
        X_dest = np.concatenate([X_dest_train, X_dest_test], axis=0)
        Y_dest = np.concatenate([Y_dest_train, Y_dest_test], axis=0)
        X_source = np.concatenate([X_source_train, X_source_test], axis=0)
        Y_source = np.concatenate([Y_source_train, Y_source_test], axis=0)
    return X_source, Y_source, X_dest, Y_dest, n_classes, img_source_dim, img_dest_dim

def load_testset(img_dim, image_dim_ordering, dset):
    # Load and normalize data
    if dset == "mnistM":
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(
            img_dim, image_dim_ordering, dset='mnistM')
    elif dset == "svhn_to_mnist32":
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
    elif dset == "svhn_to_mnist32gray":
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='mnist32')
    elif dset == "mnist_to_svhn32gray":
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32gray')
    elif dset == "mnist_to_svhn32":
        X_dest_train,Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering,dset='svhn32')
    elif dset == "OfficeAmazonToDslr":
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering, dset='OfficeDslr')
    elif dset == "OfficeAmazonToWebcam":
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering, dset='OfficeWebcam')
    elif dset == "usps_to_mnist":
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering, dset='mnist')
    elif dset == "mnist_to_usps":
        X_dest_train, Y_dest_train, X_dest_test, Y_dest_test, n_classes2 = load_image_dataset(img_dim, image_dim_ordering, dset='usps')
    else:
        print "dataset not supported in load_testset function!"
    test_data=X_dest_test
    test_labels=Y_dest_test

    return test_data, test_labels

def build_opt(opt_D, opt_G, lr_D, lr_G,lr_rec=None,opt_rec=None):
    _opt_D = get_optimizer(opt_D, lr_D)
    _opt_G = get_optimizer(opt_G, lr_G)
    _opt_C = get_optimizer('SGD', 0.01)
    _opt_Z = get_optimizer('Adam', lr_G)
    if opt_rec is None:
        return _opt_D, _opt_G, _opt_C, _opt_Z
    else:
        _opt_rec = get_optimizer(opt_rec, lr_rec)
        return _opt_D, _opt_G, _opt_C, _opt_Z, _opt_rec

def load_compile_reconstructions(generator_model1, generator_model2,noise_dim,img_source_dim1,img_source_dim2,opt_G, opt_rec, classificator_model2=None):
    rec1 = models.reconstructor(generator_model1, generator_model2, noise_dim, img_source_dim1)
    rec2 = models.reconstructor(generator_model2, generator_model1, noise_dim, img_source_dim2)
    rec1.compile(loss='mse',  optimizer=opt_rec)
    rec2.compile(loss='mse',  optimizer=opt_rec)

    if classificator_model2 is not None:
        models.make_trainable(generator_model1, False) #because generator_model1 is already trained by a classificator in a supervised setting
        models.make_trainable(generator_model2, True)
        models.make_trainable(classificator_model2, True)
        recClass = models.reconstructorClass(generator_model1, generator_model2, classificator_model2, noise_dim, img_source_dim1)
        recClass.compile(loss='categorical_crossentropy',  optimizer=opt_rec)
        return rec1,rec2,recClass
    else:
        return rec1,rec2

def load_compile_models(noise_dim, img_source_dim, img_dest_dim, deterministic, pureGAN, wd, loss1, loss2, disc_type, n_classes, opt_D, opt_G, opt_C, opt_Z,suffix=None,pretrained=False):
    # LOAD MODELS:
    generator_model = models.generator_google_mnistM(
        noise_dim, img_source_dim, img_dest_dim, deterministic, pureGAN, wd,suffix)
    discriminator_model,discriminator2 = models.discriminator_dcgan_doubled(img_dest_dim, wd,n_classes,disc_type)
#    classificator_model = models.classificator_svhn(img_dest_dim, n_classes, wd)
    if pretrained:
        classificator_model = models.resnet50classifier(img_dest_dim, n_classes, wd)  
    else:
        classificator_model = models.classificator_google_mnistM(img_dest_dim, n_classes, wd)
    DCGAN_model = models.DCGAN_naive(generator_model, discriminator_model, noise_dim, img_source_dim)
    GenClass_model = models.DCGAN_naive2(generator_model, classificator_model, noise_dim, img_source_dim)
    if not deterministic:
        zclass_model = z_coerence(generator_model, img_source_dim, bn_mode=2, wd=wd,
                                  inject_noise=False, n_classes=n_classes, noise_dim=noise_dim, model_name="zClass")

    # COMPILE MODELS:
    generator_model.compile(loss=loss1, optimizer=opt_G)
    models.make_trainable(discriminator_model, False)
    models.make_trainable(discriminator2, False)
    models.make_trainable(classificator_model, False)
    if disc_type == "simple_disc":
        DCGAN_model.compile(loss=[loss1], optimizer=opt_G)
        models.make_trainable(discriminator_model, True)
        discriminator_model.compile(loss=[loss1], optimizer=opt_D)
    elif disc_type == "nclass_disc":
        DCGAN_model.compile(loss=loss1, optimizer=opt_G)
        GenClass_model.compile(loss=['categorical_crossentropy'], optimizer=opt_G)
        models.make_trainable(discriminator_model, True)
        models.make_trainable(discriminator2, True)
        discriminator_model.compile(loss=loss1, optimizer=opt_D)
        discriminator2.compile(loss=loss2,  optimizer=opt_D)
    models.make_trainable(classificator_model, True)
    classificator_model.compile(loss=loss2, metrics=['accuracy'], optimizer=opt_C)
    if not deterministic:
        zclass_model.compile(loss=[loss1], optimizer=opt_Z)
        return generator_model, discriminator_model,discriminator2, classificator_model, DCGAN_model,GenClass_model, zclass_model
    else:
        return generator_model, discriminator_model,discriminator2,  classificator_model, DCGAN_model,GenClass_model, None


def load_pretrained_weights(generator_model, discriminator_model,discriminator2, DCGAN_model, name, data, labels, noise_scale, classificator_model=None, resume=False):
    if resume:  # loading previous saved model weights and checking actual performance
        load_model_weights(generator_model, discriminator_model, DCGAN_model, name, classificator_model, discriminator2=discriminator2)
        #loss4, acc4 = classificator_model.evaluate(data, labels, batch_size=512, verbose=0)
        #print('\n Classifier Accuracy on full target domain:  %.2f%%' % (100 * acc4))


def load_buffer_and_augmentation(history_size, batch_size, img_source_dim, n_classes):

    max_history_size = int(history_size * batch_size)
    img_buffer = ImageHistoryBuffer(
        (0,) + img_source_dim, max_history_size, batch_size, n_classes)
    datagen = ImageDataGenerator(rotation_range=0.45,
                                 width_shift_range=0.1,
                                 height_shift_range=0.1,
                                 fill_mode='nearest')
#    datagen = ImageDataGenerator(elastic_distortion=True)
    return img_buffer, datagen


def get_loss_list():
    list_disc_loss_real = deque(10 * [0], 10)
    list_disc_loss_gen = deque(10 * [0], 10)
    list_gen_loss = deque(10 * [0], 10)
    list_zclass_loss = deque(10 * [0], 10)
    list_classifier_loss = deque(10 * [0], 10)
    list_GenClass_loss = deque(10 * [0], 10)
    l_rec = deque(10 * [0], 10)
    l_recClass = deque(10 * [0], 10)
    return list_disc_loss_real, list_disc_loss_gen, list_gen_loss, list_zclass_loss, list_classifier_loss, l_rec, list_GenClass_loss, l_recClass

def get_batch(A_data, A_labels, B_data, B_labels, batch_size):
    A_data_batch, A_labels_batch, _ = next(gen_batch(A_data, A_labels, batch_size))
    B_data_batch, B_labels_batch, _ = next(gen_batch(B_data, B_labels, batch_size))
    return A_data_batch, A_labels_batch, B_data_batch, B_labels_batch

def train_gan(GAN, disc_iters, A_data, A_labels, B_data, B_labels, batch_counter, l_disc_real, l_disc_gen, l_gen,l_GenClass,class_weight):
    if GAN.dir == 'BtoA':
        for disc_it in range(disc_iters):
            A_data_batch, A_labels_batch, B_data_batch, B_labels_batch = get_batch(A_data, A_labels, B_data, B_labels, GAN.batch_size)
            X_source_batch = B_data_batch
            #Y_source_batch = B_labels_batch
            X_dest_batch = A_data_batch
            Y_dest_batch = A_labels_batch

            ##########
            # Create a batch to feed the discriminator model
            #########
            X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
            gen_output = GAN.generator_model.predict([X_noise,X_source_batch])
            #X_disc_real, X_disc_gen = get_disc_batch(X_dest_batch, GAN.generator_model, batch_counter, GAN.batch_size,
            #                                                    GAN.noise_dim, X_source_batch, noise_scale=GAN.noise_scale)
            if GAN.disc_type == "simple_disc":
                current_labels_real = np.ones(GAN.batch_size)
                current_labels_gen = np.zeros(GAN.batch_size)
            if GAN.disc_type == ("nclass_disc"):
                current_labels_real = np.ones(GAN.batch_size) 
                current_labels_gen = np.zeros(GAN.batch_size) 
            ##############
            # Train the disc on gen-buffered samples and on current real samples
            ##############
            disc_loss_real = GAN.discriminator_model.train_on_batch(X_dest_batch, current_labels_real)
            GAN.img_buffer.add_to_buffer(gen_output, current_labels_gen, GAN.batch_size)
            bufferImages, bufferLabels = GAN.img_buffer.get_from_buffer(GAN.batch_size)
            disc_loss_gen = GAN.discriminator_model.train_on_batch(bufferImages, bufferLabels)
            disc2_loss = GAN.discriminator2.train_on_batch(X_dest_batch,Y_dest_batch * 1.0) #GAN.lsmooth) #training the discriminator_classifier model

            l_disc_real.appendleft(disc_loss_real)
            l_disc_gen.appendleft(disc_loss_gen)
        #Train the GENERATOR, it is the same on both AtoB and BtoA:
        X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
        if GAN.disc_type == "simple_disc":                
            gen_loss =  GAN.DCGAN_model.train_on_batch([X_noise,X_source_batch], np.ones(GAN.batch_size))
        elif GAN.disc_type == "nclass_disc":
            gen_loss =  GAN.DCGAN_model.train_on_batch([X_noise,X_source_batch], np.ones(GAN.batch_size)) 
            class_p = GAN.GenClass_model.predict([X_noise,X_source_batch]) #AUTOLABELLING
            idx = np.argmax(class_p, axis=1) #
            virtual_labels = (idx[:, None]) == np.arange(GAN.n_classes) * 1.0 #
            GenClass_loss =  GAN.GenClass_model.train_on_batch([X_noise,X_source_batch], virtual_labels) #
        l_GenClass.appendleft(GenClass_loss) #
        l_gen.appendleft(gen_loss)

    elif GAN.dir == 'AtoB':
        for disc_it in range(disc_iters):
            A_data_batch, A_labels_batch, B_data_batch, B_labels_batch = get_batch(A_data, A_labels, B_data, B_labels, GAN.batch_size)
            X_source_batch = A_data_batch
            Y_source_batch = A_labels_batch
            X_dest_batch = B_data_batch
            #Y_dest_batch = B_labels_batch
            X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
            gen_output = GAN.generator_model.predict([X_noise,X_source_batch])
            #X_disc_real, X_disc_gen = get_disc_batch(X_dest_batch, GAN.generator_model, batch_counter, GAN.batch_size,
            #                                                    GAN.noise_dim, X_source_batch, noise_scale=GAN.noise_scale)
            if GAN.disc_type == "simple_disc":
                current_labels_real = np.ones(GAN.batch_size)
                current_labels_gen = np.zeros(GAN.batch_size)
            if GAN.disc_type == ("nclass_disc"):
                current_labels_real = np.ones(GAN.batch_size) 
                current_labels_gen = np.zeros(GAN.batch_size) 
            ##############
            #Train the disc on gen-buffered samples and on current real samples
            ##############
            disc_loss_real = GAN.discriminator_model.train_on_batch(X_dest_batch, current_labels_real)
            GAN.img_buffer.add_to_buffer(gen_output,current_labels_gen, GAN.batch_size)
            bufferImages, bufferLabels = GAN.img_buffer.get_from_buffer(GAN.batch_size)
            disc_loss_gen = GAN.discriminator_model.train_on_batch(bufferImages, bufferLabels)
            disc2_loss = GAN.discriminator2.train_on_batch(gen_output,Y_source_batch * 1.0) #GAN.lsmooth) #training the discriminator_classifier model

            l_disc_real.appendleft(disc_loss_real)
            l_disc_gen.appendleft(disc_loss_gen)

        #Train the GENERATOR, it is the same on both AtoB and BtoA:
        X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
        if GAN.disc_type == "simple_disc":                
            gen_loss =  GAN.DCGAN_model.train_on_batch([X_noise,X_source_batch], np.ones(GAN.batch_size)) #TRYING SAME BATCH OF DISC
        elif GAN.disc_type == "nclass_disc":
            gen_loss =  GAN.DCGAN_model.train_on_batch([X_noise,X_source_batch], np.ones(GAN.batch_size)) #TRYING SAME BATCH OF DISC
            GenClass_loss =  GAN.GenClass_model.train_on_batch([X_noise,X_source_batch], Y_source_batch,sample_weight=np.ones(GAN.batch_size)*class_weight)
            #gen_loss = gen_loss[0]
        l_gen.appendleft(gen_loss)
        l_GenClass.appendleft(GenClass_loss)
    return A_data_batch, A_labels_batch, B_data_batch, B_labels_batch

def train_class(GAN, l_class,  A_data_batch, A_labels_batch):
    if GAN.dir == 'AtoB':
        X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
        if GAN.data_aug:
            x_dest_batch = GAN.generator_model.predict([X_noise,datagen.output(A_data_batch)])
        else:
            x_dest_batch = GAN.generator_model.predict([X_noise,A_data_batch])
        # NO LABEL SMOOTHING!!!! inverted training w.r.t. to AtoB, because I
        # have labels of A
        class_loss = GAN.classificator_model.train_on_batch(x_dest_batch, A_labels_batch)
    elif GAN.dir == 'BtoA':
        class_loss = GAN.classificator_model.train_on_batch(A_data_batch, A_labels_batch)
    l_class.appendleft(class_loss[0])
    return l_class

def train_rec(GAN,rec1, rec2, A_data_batch, B_data_batch, l_rec1, l_rec2,rec_weight):
    X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
    X_noise2 = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
    rec_loss = rec1.train_on_batch([X_noise, A_data_batch,X_noise2],A_data_batch,sample_weight=np.ones(GAN.batch_size)*rec_weight)
    rec_loss2 = rec2.train_on_batch([X_noise, B_data_batch,X_noise2],B_data_batch,sample_weight=np.ones(GAN.batch_size)*rec_weight)
    l_rec1.appendleft(rec_loss)
    l_rec2.appendleft(rec_loss2)
    return l_rec1, l_rec2

def train_recClass(GAN,recClass, A_data_batch, A_labels_batch,  l_recClass, rec_weight):
    X_noise = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
    X_noise2 = sample_noise(GAN.noise_scale, GAN.batch_size, GAN.noise_dim)
    recClass_loss = recClass.train_on_batch([X_noise, A_data_batch,X_noise2],A_labels_batch,sample_weight=np.ones(GAN.batch_size)*rec_weight)
    l_recClass.appendleft(recClass_loss)
    return l_recClass


def train_gen_zclass(generator_model, DCGAN_model, zclass_model, disc_type, deterministic, noise_dim, noise_scale, batch_size, l_gen, l_zclass, X_source, Y_source, n_classes):
    X_gen = sample_noise(noise_scale, batch_size, noise_dim)
    X_source_batch2, Y_source_batch2, idx_source_batch2 = next(
        gen_batch(X_source, Y_source, batch_size))
    if disc_type == "simple_disc":
        gen_loss = DCGAN_model.train_on_batch([X_gen, X_source_batch2], np.ones(X_gen.shape[0]))  # TRYING SAME BATCH OF DISC
    elif disc_type == ("nclass_disc"):

        #(disc_p, class_p) = DCGAN_model.predict_on_batch(X_source_batch2)
        #idx = np.argmax(class_p, axis=1)
        #virtual_labels = (idx[:, None] == np.arange(n_classes)) * 1

        virtual_labels = np.zeros([GAN.batch_size, GAN.n_classes])
        gen_loss = DCGAN_model.train_on_batch([X_gen, X_source_batch2], [np.ones(X_gen.shape[0]), virtual_labels])  # FIX :((
        #gen_loss = gen_loss[0]
    l_gen.appendleft(gen_loss)
    if not deterministic:
        zclass_loss = zclass_model.train_on_batch(
            [X_gen, X_source_batch2], [X_gen])
    else:
        zclass_loss = 0.0
    l_zclass.appendleft(zclass_loss)
    return l_gen, l_zclass


def visualize_save_stuffs(GANs, progbar, gen_iterations, batch_counter, n_batch_per_epoch, l_disc_real1, l_disc_gen1, l_gen_loss1,
                          l_class_loss1, l_disc_real2, l_disc_gen2, l_gen_loss2,l_class_loss2, A_data, A_labels, B_data, B_labels,
                          start,e, l_rec1, l_rec2,l_GenClass1, l_GenClass2, l_recClass):
    gen_iterations += 1
    batch_counter += 1
    image_dim_ordering = 'th'
    progbar.add(GANs[0].batch_size, values=[("Loss_D_real1", np.mean(l_disc_real1)),
                                            ("Loss_D_gen1", np.mean(l_disc_gen1)),
                                            ("Loss_G1", np.mean(l_gen_loss1)),
                                            ("Loss_Classifier1",np.mean(l_class_loss1)),
                                            ("Loss_D_real2", np.mean(l_disc_real2)),
                                            ("Loss_D_gen2", np.mean(l_disc_gen2)),
                                            ("Loss_G2", np.mean(l_gen_loss2)),
                                            ("Loss_Classifier2",np.mean(l_class_loss2)),
                                            ("Loss_Rec1", np.mean(l_rec1)),
                                            ("Loss_Rec2", np.mean(l_rec2)),
                                            ("Loss_GenClass1", np.mean(l_GenClass1)),
                                            ("Loss_AutoLabel", np.mean(l_GenClass2)),
                                            ("Loss_RecClass", np.mean(l_recClass))
                                       ])

    for GAN in GANs:
        # plot images 1 times per epochs        if GAN.dir == 'BtoA':
        if GAN.dir == 'BtoA':
            X_source=B_data
            Y_source=B_labels
            X_dest = A_data
            Y_dest = A_labels
        elif GAN.dir == 'AtoB':
            X_source=A_data
            Y_source=A_labels
            X_dest = B_data
            Y_dest = B_labels

        if batch_counter == n_batch_per_epoch:
        #if batch_counter % (n_batch_per_epoch) == 0:
            X_source_batch_plot, Y_source_batch_plot, idx_source_plot = next(gen_batch(X_source, Y_source, batch_size=GAN.batch_size))
            returned_idx = plot_generated_batch(X_dest, X_source, GAN.generator_model, GAN.noise_dim, image_dim_ordering, idx_source_plot,
                                        batch_size=GAN.batch_size,different_idx=True, datagen=GAN.datagen, data_aug=GAN.data_aug)
            print ("Dest labels:") 
            print (Y_dest[returned_idx].argmax(1))
            print ("Source labels:") 
            print (Y_source_batch_plot.argmax(1))
            print('\nEpoch %s, Time: %s' % (e + 1, time.time() - start))
        else:
            idx_source_plot = 0
            Y_source_batch_plot = 0

         #Save model weights (by default, every 5 epochs)
        if batch_counter == n_batch_per_epoch:
            save_model_weights(GAN.generator_model, GAN.discriminator_model,
                                      GAN.DCGAN_model, e, GAN.name, GAN.classificator_model, discriminator2=GAN.discriminator2)
    return batch_counter, gen_iterations

def pretrain_disc( GAN, A_data, A_labels,B_data, B_labels,class_weight, pretrain_iters=100, resume=False):
    l_real = deque(10 * [0], 10)
    l_gen = deque(10 * [0], 10)
    l_genclass = deque(10 * [0], 10)
    if not resume:
        _, _, _, _ = train_gan(GAN, pretrain_iters, A_data, A_labels, B_data, B_labels, 1, l_real, l_gen,l_gen,l_genclass,class_weight)
        print "Pretrain of discriminator finished."
    else:
        print "resumed previous training."


def testing_class_accuracy(GANs,classificator_model, generator_model, vis_samples, noise_dim, noise_scale, data, labels):
    acc=[]
    loss=[]
    for GAN in GANs:
        if GAN.dir == 'BtoA':
            # testing accuracy of trained classifier
            X_noise = sample_noise(GAN.noise_scale, vis_samples, GAN.noise_dim)
            Xsource_dataset_mapped = GAN.generator_model.predict(
                [X_noise, data[:vis_samples]], batch_size=1000)
            true_labels = labels[:vis_samples]
            p1 = GAN.classificator_model.predict(Xsource_dataset_mapped, batch_size=1000, verbose=1)
            score1 = np.sum(np.argmax(true_labels,axis=1) == np.argmax(p1, axis=1)) / float(true_labels.shape[0]) 
            print('\n Classifier Accuracy and loss on full target domain:  %.2f%%  ' %
                  ((100 * score1)))

        if GAN.dir == 'AtoB':
            X_noise = sample_noise(GAN.noise_scale, vis_samples, GAN.noise_dim)
            Xsource_dataset_mapped = data[:vis_samples]
            true_labels = labels[:vis_samples]
            p2 = GAN.classificator_model.predict(Xsource_dataset_mapped, batch_size=1000, verbose=1)
            score2 = np.sum(np.argmax(true_labels,axis=1) == np.argmax(p2, axis=1)) / float(true_labels.shape[0]) 
            print('\n Classifier Accuracy and loss on full target domain:  %.2f%%  ' %
                  ((100 * score2)))
    
    res = []
    for x in np.arange(0, 1.1, 0.1):
        res.append((x, np.sum(np.argmax(true_labels,axis=1) == np.argmax(p1*x + p2*(1-x), axis=1)) / float(true_labels.shape[0])))
    for (x, score) in res:
        print("\n Coeff: %f - score: %.2f" % (x, score*100))



def train(**kwargs):
    """
    Train standard DCGAN model

    args: **kwargs (dict) keyword arguments that specify the model hyperparameters
    """

    # Roll out the parameters
    generator = kwargs["generator"]
    discriminator = kwargs["discriminator"]
    dset = kwargs["dset"]
    img_dim = kwargs["img_dim"]
    nb_epoch = kwargs["nb_epoch"]
    batch_size = kwargs["batch_size"]
    n_batch_per_epoch = kwargs["n_batch_per_epoch"]
    bn_mode = kwargs["bn_mode"]
    noise_dim = kwargs["noise_dim"]
    noise_scale = kwargs["noise_scale"]
    lr_rec = kwargs["lr_D"]    
    opt_rec = kwargs["opt_rec"]
    lr_G = kwargs["lr_G"]
    lr_D = kwargs["lr_D"]
    opt_D = kwargs["opt_D"]
    opt_G = kwargs["opt_G"]
    use_mbd = kwargs["use_mbd"]
    image_dim_ordering = kwargs["image_dim_ordering"]
    epoch_size = n_batch_per_epoch * batch_size
    deterministic1 = kwargs["deterministic1"]
    deterministic2 = kwargs["deterministic2"]
    inject_noise = kwargs["inject_noise"]
    model = kwargs["model"]
    no_supertrain = kwargs["no_supertrain"]
    pureGAN = kwargs["pureGAN"]
    lsmooth = kwargs["lsmooth"]
    disc_type = kwargs["disc_type"]
    resume = kwargs["resume"]
    name = kwargs["name"]
    wd = kwargs["wd"]
    history_size = kwargs["history_size"]
    monsterClass = kwargs["monsterClass"]
    data_aug = kwargs["data_aug"]
    disc_iters = kwargs["disc_iterations"]
    class_weight = kwargs["class_weight"]
    reconst_w= kwargs["reconst_w"]
    rec = kwargs["rec"]
    reconstClass = kwargs["reconstClass"]
    pretrained = kwargs["pretrained"]
    print("\nExperiment parameters:")
    for key in kwargs.keys():
        print key, kwargs[key]
    print("\n")
    #####some extra parameters:
    
    noise_dim = (noise_dim,)
    name1 = name + '1'
    name2 = name + '2'
    # Setup environment (logging directory etc)
    general_utils.setup_logging("DCGAN")
    gen_iterations = 0
    # Loading data
    A_data, A_labels, B_data, B_labels, n_classes, img_A_dim, img_B_dim = load_data(img_dim, image_dim_ordering, dset)
    test_data, test_labels  = load_testset(img_dim, image_dim_ordering, dset)
    if deterministic1 is None:
        deterministic1 = False   
    if deterministic2 is None:
        deterministic2 = False

    opt_D1, opt_G1, opt_C1, opt_Z1, opt_rec = build_opt(opt_D, opt_G, lr_D, lr_G, lr_rec, opt_rec)
    generator_model1, discriminator_model1,discriminator_class1, classificator_model1, DCGAN_model1, GenClass_model1, zclass_model1 = load_compile_models(noise_dim, img_A_dim, img_B_dim,
                                                     deterministic1, pureGAN, wd, 'mse', 'categorical_crossentropy', disc_type, n_classes, opt_D1, opt_G1, opt_C1, opt_Z1, suffix=None, pretrained=pretrained)
    load_pretrained_weights(generator_model1, discriminator_model1,discriminator_class1, DCGAN_model1, name1, B_data, B_labels, noise_scale, classificator_model1, resume=resume)
    img_buffer1, datagen1 = load_buffer_and_augmentation(history_size, batch_size, img_A_dim, n_classes)

    GAN1=_GAN(generator_model1, discriminator_model1, discriminator_class1,DCGAN_model1,GenClass_model1,classificator_model1, batch_size, img_A_dim,img_B_dim, noise_dim, noise_scale,
               lr_D, lr_G, deterministic1, inject_noise, model, lsmooth, img_buffer1, datagen1, disc_type, data_aug, n_classes, disc_iters,name1, dir='AtoB' )
    pretrain_disc( GAN1, A_data,A_labels, B_data, B_labels,class_weight, pretrain_iters=500, resume=resume)
    #####################

   ##### Setup GAN2
    opt_D2, opt_G2, opt_C2, opt_Z2 = build_opt(opt_D, opt_G, lr_D, lr_G)
    generator_model2, discriminator_model2, discriminator_class2, classificator_model2, DCGAN_model2, GenClass_model2, zclass_model2 = load_compile_models(noise_dim, img_B_dim, img_A_dim,
                                                       deterministic2, pureGAN, wd, 'mse', 'categorical_crossentropy', disc_type, n_classes, opt_D2, opt_G2, opt_C2, opt_Z2, suffix=True)
    load_pretrained_weights(generator_model2, discriminator_model2,discriminator_class2, DCGAN_model2, name2, B_data, B_labels, noise_scale, classificator_model2, resume=resume)
    img_buffer2, datagen2 = load_buffer_and_augmentation(history_size, batch_size, img_B_dim, n_classes)

    gen_entropy2=None
    GAN2=_GAN(generator_model2, discriminator_model2, discriminator_class2, DCGAN_model2,GenClass_model2,classificator_model2, batch_size, img_B_dim,img_A_dim, noise_dim, noise_scale,
               lr_D, lr_G, deterministic2, inject_noise, model, lsmooth, img_buffer2, datagen2, disc_type, data_aug, n_classes, disc_iters, name2, dir='BtoA' )
    pretrain_disc( GAN2, A_data,A_labels, B_data, B_labels,class_weight, pretrain_iters=500, resume=resume)
    ################

    ####Reconstruction losses between Gen1 and Gen2:
    rec1, rec2, recClass = load_compile_reconstructions(generator_model1, generator_model2,noise_dim,img_A_dim,img_B_dim,opt_G,opt_rec, classificator_model2)  

    if resume:
        testing_class_accuracy([GAN1,GAN2],GAN1.classificator_model, GAN1.generator_model,
                               test_data.shape[0], GAN1.noise_dim, GAN1.noise_scale, test_data, test_labels)

#    if two_gans:
#        GANs=[GAN1,GAN2]
#    else:
#        GANs=[GAN1]
    ################
    ##################
    for e in range(1, nb_epoch + 1):
        # Initialize progbar and batch counter
        progbar = generic_utils.Progbar(epoch_size,interval=0.2)
        batch_counter = 1
        start = time.time()
        while batch_counter < n_batch_per_epoch:
            l_disc_real1, l_disc_gen1, l_gen1, l_z1, l_class1,l_rec1,l_GenClass1,_ = get_loss_list()
            A_data_batch, A_labels_batch, B_data_batch, B_labels_batch = train_gan(GAN1, GAN1.disc_iters, A_data, A_labels, B_data, B_labels, batch_counter, l_disc_real1, l_disc_gen1, l_gen1,l_GenClass1, class_weight)

            l_disc_real2, l_disc_gen2, l_gen2, l_z2, l_class2, l_rec2, l_GenClass2, l_recClass = get_loss_list()
            A_data_batch, A_labels_batch, B_data_batch, B_labels_batch = train_gan(GAN2, GAN2.disc_iters, A_data, A_labels, B_data, B_labels, batch_counter, l_disc_real2, l_disc_gen2,l_gen2,l_GenClass2, class_weight)
            if rec:
                train_rec(GAN1, rec1, rec2, A_data_batch, B_data_batch,l_rec1, l_rec2,reconst_w) #BRINGING US TO L.A.? :)
            if reconstClass > 0.0:
                train_recClass(GAN1,recClass, A_data_batch, A_labels_batch,  l_recClass, reconstClass)
            l_class1 = train_class(GAN1, l_class1,  A_data_batch, A_labels_batch)
            l_class2 = train_class(GAN2, l_class2,  A_data_batch, A_labels_batch)

            batch_counter, gen_iterations = visualize_save_stuffs([GAN1,GAN2], progbar, gen_iterations, batch_counter, n_batch_per_epoch, 
                                                                              l_disc_real1, l_disc_gen1, l_gen1, l_class1, l_disc_real2, l_disc_gen2,
                                                                              l_gen2, l_class2, A_data, A_labels, B_data, B_labels,start,e,l_rec1, l_rec2,
                                                                              l_GenClass1, l_GenClass2, l_recClass)


        testing_class_accuracy([GAN1,GAN2],GAN1.classificator_model, GAN1.generator_model,
                               test_data.shape[0], GAN1.noise_dim, GAN1.noise_scale, test_data, test_labels)
#        testing_class_accuracy([GAN1],GAN1.classificator_model, GAN1.generator_model,
#                               5000, GAN1.noise_dim, GAN1.noise_scale, B_data, B_labels)