All

2019-12-21 16:37:10 -05:00
commit 21970b730a
@@ -0,0 +1,6 @@
+data/*
+experiments/*
+__pycache__
+.vscode
+animations/eric_andre/pretrained_models/*
+animations/eric_andre/results/*
@@ -0,0 +1,102 @@
+# GANimation 
+
+This repository contains an implementation of [GANimation](https://arxiv.org/pdf/1807.09251.pdf) by Pumarola et al. based on [StarGAN code](https://github.com/yunjey/stargan) by @yunjey. With this model they are able to modify in a continuous way facial expressions of single images.
+
+[Pretrained models](https://www.dropbox.com/sh/108g19dk3gt1l7l/AAB4OJHHrMHlBDbNK8aFQVZSa?dl=0) and the [preprocessed CelebA dataset](https://www.dropbox.com/s/payjdk08292csra/celeba.zip?dl=0) are provided to facilitate the use of this model as well as the process for preparing other datasets for training this model.
+
+<p align="center">
+  <img width="170" height="170" src="https://github.com/vipermu/ganimation/blob/master/video_results/frida.gif">
+</p>
+
+<p align="center">
+  <img width="600" height="150" src="https://github.com/vipermu/ganimation/blob/master/video_results/eric_andre.gif">
+</p>
+
+
+## Setup
+
+#### Conda environment
+Create your conda environment by just running the following command:
+`conda env create -f environment.yml`
+
+
+## Datasets
+
+#### CelebA preprocessed dataset
+Download and unzip the *CelebA* preprocessed dataset uploaded to [this link](https://www.dropbox.com/s/payjdk08292csra/celeba.zip?dl=0) extracted from [MMLAB](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html). Here you can find a folder containing the aligned and resized 128x128 images as well as a _txt_ file containing their respective Action Units vectors computed using [OpenFace](https://github.com/TadasBaltrusaitis/OpenFace). By default, this code assumes that you have these two elements in _`./data/celeba/`_.
+
+#### Use your own dataset
+If you want to use other datasets you will need to detect and crop bounding boxes around the face of each image, compute their corresponding Action Unit vectors and resize them to 128x128px.
+
+You can perform all these steps using [OpenFace](https://github.com/TadasBaltrusaitis/OpenFace). First you will need to setup the project. They provide guides for [linux](https://github.com/TadasBaltrusaitis/OpenFace/wiki/Unix-Installation) and [windows](https://github.com/TadasBaltrusaitis/OpenFace/wiki/Windows-Installation). Once the models are compiled, read their [Action Unit wiki](https://github.com/TadasBaltrusaitis/OpenFace/wiki/Action-Units) and their [documentation](https://github.com/TadasBaltrusaitis/OpenFace/wiki/Command-line-arguments) on these models to find out which is the command that you need to execute.
+
+In my case the command was the following: `./build/bin/FaceLandmarkImg -fdir datasets/my-dataset/ -out_dir processed/my-processed-dataset/ -aus -simalign -au_static -nobadaligned -simsize 128 -format_aligned jpg -nomask`
+
+After computing these Action Units, depending on the command that you have used, you will obtain different output formats. With the command that I used, I obtained a _csv_ file for each image containing its corresponding Action Units vector among extra information, a folder for each image containing the resized and cropped image and a _txt_ file with extra details about each image. You can find in _openface_utils_ folder the code that I used to extract all the Action Unit information in a _txt_ file and to group all the images into a single folder.
+
+After having the Action Unit _txt_ file and the image folder you can move them to the directory of this project. By default, this code assumes that you have these two elements in _`./data/celeba/`_.
+
+## Generate animations
+Pretrained models can be downloaded from [this](https://www.dropbox.com/sh/108g19dk3gt1l7l/AAB4OJHHrMHlBDbNK8aFQVZSa?dl=0) link. This folder contains the weights of both models (the Generator and the Discriminator) after training the model for 37 epochs.
+
+By running `python main.py --mode animation` the default animation will be executed. There are two different types of animations already implemented which can be selected with the parameter 'animation_mode'. It is presuposed that the following folders are present:
+
+- **attribute_images**: images from which the Action Units that we want to use for the animation were computed.
+- **images_to_animate**: images that we want to animate.
+- **pretrained_models**: pretrained models (only the generator is needed, you can download it from [here](https://www.dropbox.com/home/data/pretrained_models)
+- **results**: folder where the resulting images will be stored.
+- **attributes.txt**: file with the action units from 'attribute_images' computed.
+
+The two options already implemented are the following:
+- **animate_image**: applies the expressions from 'attributes.txt' to the images in 'images_to_animate'.
+- **animate_random_batch**: applies the expressions from 'attributes.txt' to random batches of images from the training dataset.
+
+
+## Train the model
+
+#### Parameters
+
+You can either modify these parameters in `main.py` or by calling them as command line arguments.
+
+
+##### Lambdas
+
+- *lambda_cls*: classification lambda.
+- *lambda_rec*: lambda for the cycle consistency loss.
+- *lambda_gp*: gradient penalty lambda.
+- *lambda_sat*: lambda for attention saturation loss.
+- *lambda_smooth*: lambda for attention smoothing loss.
+
+##### Training parameters
+
+- *c_dim*: number of Action Units to use to train the model.
+- *batch_size*
+- *num_epochs*
+- *num_epochs_decay*: number of epochs to start decaying the learning rate.
+- *g_lr*: generator's learning rate.
+- *d_lr*: discriminator's learning rate.
+
+##### Pretrained models parameters
+The weights are stored in the following format: `<epoch>-<iteration>-<G/D>.ckpt` where G and D represent the Generator and the Discriminator respectively. We save the state of thoptimizers in the same format and extension but add '_optim'.
+
+- *resume_iters*: iteration numbre from which we want to start the training. Note that we will need to have a saved model corresponding to that exact iteration number.
+- *first_epoch*: initial epoch for when we train from pretrained models.
+
+##### Miscellaneous:
+- *mode*: train/test.
+- *image_dir*: path to your image folder.
+- *attr_path*: path to your attributes _txt_ folder.
+- *outputs_dir*: name for the output folder.
+
+#### Virtual
+- *use_virtual*: this flag activates the use of _cycle consistency loss_ during the training.
+
+## Virtual Cycle Consistency Loss
+The aim of this new component is to minimize the noise produced by the Action Unit regression. This idea was extracted from [Label-Noise Robust Multi-Domain Image-to-Image Translation](https://arxiv.org/abs/1905.02185) by Kaneko et al.. It is not proven that this new component improves the outcomes of the model but the masks seem to be darker when it is applied without losing realism on the output images.
+
+## TODOs
+
+- Clean Test function. (DONE)
+- Add an Action Units selector option for training.
+- Add multi-gpu support.
+- Smoother video generation.
@@ -0,0 +1,81 @@
+22.jpg 0.00 0.00 0.00 0.00 1.44 0.00 0.74 1.22 2.54 1.76 0.00 0.00 0.16 0.00 1.89 0.69 0.13
+66.jpg 0.00 0.00 0.00 0.00 1.37 0.37 1.35 1.01 2.37 1.36 0.20 0.00 0.25 0.00 1.95 0.94 1.54
+38.jpg 0.22 0.02 0.00 0.00 0.99 0.29 0.30 0.63 2.35 1.31 0.00 0.00 0.00 0.00 1.44 0.99 0.31
+33.jpg 0.00 0.00 0.00 0.00 1.37 0.00 1.00 1.73 2.52 1.16 0.00 0.00 0.00 0.00 2.25 0.73 0.00
+39.jpg 0.32 0.04 0.00 0.00 0.98 0.05 0.28 0.67 2.31 1.36 0.00 0.00 0.13 0.00 1.50 0.68 0.24
+30.jpg 0.00 0.00 0.00 0.00 1.31 0.00 0.76 1.15 2.51 1.53 0.00 0.00 0.00 0.00 1.68 0.76 0.01
+12.jpg 0.00 0.00 0.00 0.00 1.25 0.97 1.11 0.59 2.04 1.47 0.00 0.00 0.38 0.00 1.61 0.66 0.00
+35.jpg 0.00 0.00 0.00 0.00 1.27 0.00 0.68 1.52 2.52 1.18 0.00 0.00 0.00 0.00 2.08 0.58 0.20
+40.jpg 0.32 0.00 0.00 0.10 0.90 0.17 0.36 0.83 2.22 1.27 0.00 0.00 0.29 0.00 1.59 0.65 0.24
+11.jpg 0.00 0.00 0.00 0.00 1.09 0.69 1.13 0.64 1.97 1.32 0.04 0.00 0.30 0.00 1.63 0.59 0.00
+50.jpg 0.35 0.05 0.00 0.00 0.75 0.68 0.09 0.53 2.10 1.71 0.00 0.00 0.36 0.06 0.55 0.48 0.02
+74.jpg 0.00 0.00 0.00 0.25 0.00 0.23 1.00 0.00 0.45 0.96 1.33 0.00 0.90 0.68 0.00 1.79 0.00
+42.jpg 0.16 0.00 0.00 0.00 1.00 0.27 0.74 1.22 2.12 1.26 0.11 0.00 0.00 0.00 1.68 0.41 0.09
+3.jpg 0.00 0.00 0.00 0.00 0.92 0.19 0.90 0.64 1.94 1.41 0.00 0.00 0.66 0.00 1.14 0.34 0.00
+77.jpg 0.00 0.00 0.00 0.67 0.00 0.00 0.84 0.00 0.53 0.83 1.44 0.00 0.42 0.68 0.00 1.20 0.10
+14.jpg 0.00 0.00 0.00 0.00 1.45 0.75 1.14 0.49 2.17 1.34 0.00 0.00 0.22 0.00 1.43 0.33 0.48
+25.jpg 0.00 0.00 0.00 0.00 1.44 0.09 0.91 1.33 2.48 1.54 0.00 0.00 0.00 0.00 1.88 0.57 0.00
+32.jpg 0.00 0.00 0.00 0.00 1.37 0.00 0.91 1.58 2.53 1.32 0.00 0.00 0.00 0.00 2.10 0.89 0.00
+81.jpg 0.00 0.00 0.00 0.04 0.52 0.13 1.99 0.97 0.67 1.12 1.27 0.00 0.27 0.00 0.03 0.66 0.00
+55.jpg 0.21 0.37 0.00 0.05 0.50 0.25 0.02 0.70 1.97 1.72 0.00 0.00 0.50 0.18 0.08 0.40 0.06
+49.jpg 0.32 0.00 0.00 0.00 0.76 0.74 0.28 0.97 2.02 1.74 0.00 0.00 0.29 0.00 1.46 0.51 0.00
+21.jpg 0.00 0.00 0.00 0.00 1.44 0.21 0.87 1.09 2.49 1.60 0.00 0.00 0.00 0.00 1.78 0.59 0.19
+52.jpg 0.29 0.00 0.00 0.21 0.61 0.97 0.21 0.75 1.75 1.79 0.00 0.00 0.34 0.13 0.94 0.90 0.00
+15.jpg 0.00 0.00 0.00 0.00 1.63 0.73 1.22 0.90 2.30 1.27 0.00 0.00 0.00 0.00 1.53 0.55 0.66
+48.jpg 0.46 0.00 0.00 0.00 1.04 0.82 0.55 1.22 2.25 1.84 0.00 0.00 0.48 0.00 1.95 0.69 0.00
+80.jpg 0.00 0.00 0.00 0.28 0.44 0.00 1.29 0.59 0.80 1.06 1.28 0.00 0.04 0.12 0.00 0.89 0.00
+67.jpg 0.00 0.01 0.00 0.00 0.54 0.03 1.55 0.00 1.20 0.80 0.67 0.00 0.64 0.00 0.84 0.76 1.16
+16.jpg 0.00 0.00 0.00 0.00 1.60 0.45 1.25 0.96 2.35 1.38 0.00 0.00 0.00 0.00 1.58 0.52 0.78
+26.jpg 0.00 0.00 0.00 0.00 1.32 0.00 0.95 1.46 2.45 1.76 0.00 0.00 0.00 0.00 1.68 0.80 0.00
+78.jpg 0.00 0.00 0.00 0.43 0.00 0.00 1.04 0.00 0.55 0.63 1.45 0.00 0.52 0.58 0.00 1.35 0.08
+37.jpg 0.25 0.01 0.00 0.13 1.16 0.35 0.51 1.09 2.38 1.31 0.13 0.00 0.00 0.00 1.70 0.95 0.29
+43.jpg 0.24 0.00 0.00 0.00 0.99 0.40 0.59 1.20 2.18 0.86 0.00 0.00 0.00 0.00 1.75 0.72 0.08
+79.jpg 0.00 0.00 0.00 0.30 0.12 0.00 1.11 0.00 0.62 0.42 1.37 0.00 0.14 0.37 0.00 0.84 0.10
+8.jpg 0.00 0.00 0.00 0.00 1.14 0.17 1.03 0.70 1.98 1.42 0.00 0.00 0.45 0.00 1.55 0.34 0.00
+17.jpg 0.00 0.00 0.00 0.00 1.60 0.59 1.23 1.04 2.35 1.42 0.00 0.00 0.11 0.00 1.67 0.81 0.57
+34.jpg 0.00 0.00 0.00 0.00 1.40 0.00 0.81 1.58 2.54 1.12 0.00 0.00 0.00 0.00 2.28 0.88 0.04
+65.jpg 0.00 0.00 0.00 0.00 1.53 0.08 1.19 1.40 2.45 1.66 0.13 0.00 0.29 0.00 2.36 1.03 1.34
+53.jpg 0.27 0.00 0.00 0.00 0.63 0.88 0.29 0.76 1.83 1.90 0.00 0.00 0.49 0.05 0.96 0.91 0.00
+69.jpg 0.00 0.00 0.00 0.07 0.26 0.39 1.29 0.00 0.31 0.77 0.87 0.00 1.37 0.30 0.38 1.19 0.58
+51.jpg 0.30 0.04 0.00 0.00 0.57 0.47 0.12 0.57 1.89 1.62 0.00 0.00 0.23 0.12 0.48 0.53 0.00
+6.jpg 0.00 0.00 0.00 0.00 0.89 0.03 0.98 0.60 1.90 1.30 0.00 0.00 0.54 0.00 1.18 0.01 0.00
+4.jpg 0.00 0.00 0.00 0.00 0.87 0.00 0.90 0.59 1.91 1.52 0.00 0.00 0.73 0.00 1.12 0.20 0.00
+44.jpg 0.34 0.00 0.00 0.00 0.94 0.26 0.46 1.06 2.19 1.45 0.00 0.00 0.00 0.23 1.09 0.76 0.02
+13.jpg 0.00 0.00 0.00 0.00 1.32 0.81 1.14 0.67 2.12 1.16 0.00 0.00 0.22 0.00 1.51 0.39 0.11
+1.jpg 0.00 0.00 0.00 0.00 0.97 0.01 0.90 0.76 1.91 1.36 0.00 0.00 0.49 0.00 1.22 0.22 0.00
+2.jpg 0.00 0.00 0.00 0.00 1.02 0.00 1.07 0.74 1.96 1.33 0.00 0.00 0.54 0.00 1.18 0.26 0.00
+56.jpg 0.64 0.34 0.00 0.17 0.55 0.72 0.41 0.78 1.81 1.58 0.01 0.00 0.57 0.00 0.99 0.88 0.00
+27.jpg 0.00 0.00 0.00 0.00 1.38 0.00 0.76 1.24 2.52 1.70 0.00 0.00 0.00 0.00 1.61 0.88 0.03
+57.jpg 0.42 0.20 0.00 0.00 0.69 0.57 0.69 1.24 1.59 2.05 0.30 0.00 1.15 0.00 1.60 1.14 0.00
+60.jpg 0.01 0.39 0.00 0.00 1.21 0.46 1.21 1.35 1.60 2.23 0.62 0.00 1.31 0.00 1.56 0.88 0.07
+71.jpg 0.00 0.38 0.00 0.04 0.20 1.14 1.17 0.00 0.26 0.36 0.76 0.00 0.92 0.08 0.52 1.82 0.29
+62.jpg 0.00 0.00 0.00 0.00 1.44 0.54 1.02 1.70 2.07 2.06 0.00 0.00 0.83 0.00 3.09 0.94 0.04
+46.jpg 0.27 0.00 0.00 0.00 0.87 0.60 0.39 0.89 2.03 1.58 0.00 0.00 0.14 0.08 1.34 0.88 0.04
+76.jpg 0.00 0.00 0.00 0.36 0.00 0.23 0.91 0.00 0.65 1.21 1.39 0.00 0.79 0.63 0.00 1.46 0.08
+70.jpg 0.00 0.05 0.00 0.27 0.19 1.08 1.04 0.00 0.31 0.50 0.85 0.00 1.27 0.21 0.48 1.27 0.44
+24.jpg 0.00 0.00 0.00 0.00 1.47 0.00 0.90 1.43 2.61 1.57 0.00 0.00 0.00 0.00 2.03 0.75 0.05
+59.jpg 0.45 0.81 0.00 0.35 0.95 0.67 0.99 1.12 1.52 2.26 0.64 0.00 1.43 0.00 1.60 1.06 0.00
+7.jpg 0.00 0.00 0.00 0.00 1.06 0.27 1.02 0.70 1.90 1.44 0.00 0.00 0.48 0.00 1.14 0.00 0.00
+64.jpg 0.00 0.00 0.00 0.00 1.65 0.68 1.18 1.76 2.49 1.66 0.03 0.00 0.52 0.00 2.73 0.75 1.21
+9.jpg 0.00 0.00 0.00 0.00 1.20 0.68 1.18 0.86 1.90 1.41 0.00 0.00 0.52 0.00 1.92 0.53 0.00
+72.jpg 0.00 0.02 0.00 0.24 0.00 0.98 1.13 0.00 0.29 0.97 1.06 0.00 1.33 0.36 0.26 1.79 0.03
+63.jpg 0.00 0.00 0.00 0.00 1.74 0.47 1.28 1.66 2.37 1.88 0.00 0.00 0.53 0.00 2.92 1.01 0.35
+19.jpg 0.00 0.00 0.00 0.00 1.62 0.44 1.25 1.17 2.43 1.65 0.00 0.00 0.04 0.00 1.70 0.85 0.41
+20.jpg 0.00 0.00 0.00 0.00 1.49 0.19 1.17 1.24 2.48 1.56 0.00 0.00 0.23 0.00 1.77 0.79 0.19
+47.jpg 0.25 0.00 0.00 0.00 0.90 0.47 0.41 0.96 2.13 1.66 0.05 0.00 0.40 0.03 1.59 0.88 0.03
+29.jpg 0.00 0.00 0.00 0.00 1.33 0.00 0.85 1.25 2.38 1.48 0.00 0.00 0.11 0.00 1.70 1.05 0.00
+36.jpg 0.00 0.00 0.00 0.08 1.44 0.45 0.68 1.14 2.44 1.51 0.11 0.00 0.09 0.00 1.95 0.88 0.37
+68.jpg 0.00 0.00 0.00 0.00 0.30 0.00 1.14 0.00 0.76 0.95 0.91 0.00 1.21 0.46 0.57 1.26 1.01
+73.jpg 0.00 0.00 0.00 0.27 0.00 0.28 1.16 0.00 0.43 1.39 1.30 0.00 1.16 0.68 0.00 1.90 0.07
+18.jpg 0.00 0.00 0.00 0.00 1.50 0.55 1.10 0.99 2.42 1.33 0.00 0.00 0.00 0.00 1.63 0.84 0.49
+58.jpg 0.86 1.24 0.00 0.34 0.76 0.23 0.63 1.08 1.65 2.15 0.51 0.00 1.34 0.00 1.89 1.51 0.00
+54.jpg 0.30 0.37 0.00 0.00 0.37 0.08 0.00 0.52 1.99 1.69 0.07 0.00 0.42 0.25 0.13 0.17 0.03
+10.jpg 0.00 0.00 0.00 0.00 1.11 0.59 1.15 0.80 1.87 1.36 0.01 0.00 0.43 0.00 1.87 0.81 0.00
+5.jpg 0.00 0.00 0.00 0.00 0.95 0.00 1.05 0.70 1.90 1.48 0.00 0.00 0.69 0.00 1.16 0.02 0.00
+28.jpg 0.00 0.00 0.00 0.00 1.31 0.00 0.69 1.17 2.54 1.48 0.00 0.00 0.07 0.00 1.76 0.94 0.00
+45.jpg 0.06 0.00 0.00 0.00 0.91 0.32 0.26 0.79 2.17 1.62 0.00 0.00 0.00 0.10 1.13 0.90 0.13
+61.jpg 0.00 0.00 0.00 0.00 1.36 0.25 1.01 1.49 1.80 2.20 0.08 0.00 1.29 0.00 2.84 1.11 0.00
+23.jpg 0.00 0.00 0.00 0.00 1.53 0.00 0.91 1.27 2.60 1.70 0.00 0.00 0.02 0.00 2.05 0.66 0.00
+31.jpg 0.00 0.00 0.00 0.00 1.35 0.00 0.74 1.03 2.49 1.49 0.00 0.00 0.06 0.00 1.77 0.85 0.05
+41.jpg 0.19 0.00 0.00 0.00 1.03 0.41 0.68 0.96 2.17 1.29 0.04 0.00 0.00 0.00 1.57 0.56 0.07
+75.jpg 0.00 0.00 0.00 0.35 0.00 0.21 1.00 0.00 0.57 1.41 1.27 0.00 0.84 0.88 0.00 1.99 0.13
@@ -0,0 +1,128 @@
+import copy
+import numpy as np
+from collections import Iterable
+from scipy.stats import truncnorm
+
+import torch
+import torch.nn as nn
+
+class LinfPGDAttack(object):
+    def __init__(self, model=None, device=None, epsilon=0.03, k=40, a=0.01):
+        self.model = model
+        self.epsilon = epsilon
+        self.k = k
+        self.a = a
+        self.loss_fn = nn.MSELoss().to(device)
+        self.device = device
+
+    def perturb(self, X_nat, y, c_trg):
+        """
+        Given examples (X_nat, y), returns adversarial
+        examples within epsilon of X_nat in l_infinity norm.
+        """
+        X = X_nat.clone().detach_()
+
+        for i in range(self.k):
+            # print(i)
+            X.requires_grad = True
+            output_att, output_img = self.model(X, c_trg)
+
+            out = imFromAttReg(output_att, output_img, X)
+
+            self.model.zero_grad()
+            # loss = -self.loss_fn(output_att, y) + self.loss_fn(output_img, y)
+            # loss = self.loss_fn(output_att, y)
+            loss = self.loss_fn(output_att, y)
+            loss.backward()
+            grad = X.grad
+
+            X_adv = X + self.a * grad.sign()
+
+            eta = torch.clamp(X_adv - X_nat, min=-self.epsilon, max=self.epsilon)
+            X = torch.clamp(X_nat + eta, min=-1, max=1).detach_()
+
+        return X, eta
+
+    def perturb_iter_data(self, X_nat, X_all, y, c_trg):
+        """
+        X_nat is a tensor with several different images.
+        This does not work at all yet..
+        """
+        X = X_nat.clone().detach_()
+        # X_all_local = X_all.clone().detach_()
+
+        j = 0
+        J = X_all.size(0)
+        J = 1
+
+        for i in range(self.k):
+            # print(i,j)
+            X_j = X_all[j].unsqueeze(0)
+            X_j.requires_grad = True
+            output_att, output_img = self.model(X_j, c_trg)
+
+            out = imFromAttReg(output_att, output_img, X_j)
+
+            self.model.zero_grad()
+            loss = -self.loss_fn(out, y)
+            loss.backward()
+            grad = X_j.grad
+
+            X_adv = X + self.a * grad.sign()
+
+            eta = torch.clamp(X_adv - X_nat, min=-self.epsilon, max=self.epsilon)
+            X = torch.clamp(X_nat + eta, min=-1, max=1).detach_()
+
+            j += 1
+            if j == J:
+                j = 0
+
+        return X, eta
+
+    def perturb_iter_class(self, X_nat, y, c_trg):
+        """
+        Given examples (X_nat, y), returns adversarial
+        examples within epsilon of X_nat in l_infinity norm.
+        """
+        X = X_nat.clone().detach_()
+
+        j = 0
+        J = c_trg.size(0)
+
+        for i in range(self.k):
+            # print(i)
+            X.requires_grad = True
+            output_att, output_img = self.model(X, c_trg[j,:].unsqueeze(0))
+
+            out = imFromAttReg(output_att, output_img, X)
+
+            self.model.zero_grad()
+
+            # Away from black
+            loss = self.loss_fn(output_att, y)
+            loss.backward()
+            grad = X.grad
+
+            X_adv = X + self.a * grad.sign()
+
+            eta = torch.clamp(X_adv - X_nat, min=-self.epsilon, max=self.epsilon)
+            X = torch.clamp(X_nat + eta, min=-1, max=1).detach_()
+
+            j += 1
+            if j == J:
+                j = 0
+
+        return X, eta
+
+def clip_tensor(X, Y, Z):
+    # Clip X with Y min and Z max
+    X_np = X.data.cpu().numpy()
+    Y_np = Y.data.cpu().numpy()
+    Z_np = Z.data.cpu().numpy()
+    X_clipped = np.clip(X_np, Y_np, Z_np)
+    X_res = torch.FloatTensor(X_clipped)
+    return X_res
+
+def imFromAttReg(att, reg, x_real):
+    """Mixes attention, color and real images"""
+    return (1-att)*reg + att*x_real
@@ -0,0 +1,98 @@
+import argparse
+
+
+def get_config():
+    parser = argparse.ArgumentParser()
+
+    # Model configuration.
+    parser.add_argument('--c_dim', type=int, default=17,
+                        help='dimension of domain labels')
+
+    parser.add_argument('--image_size', type=int,
+                        default=128, help='image resolution')
+    parser.add_argument('--g_conv_dim', type=int, default=64,
+                        help='number of conv filters in the first layer of G')
+    parser.add_argument('--d_conv_dim', type=int, default=64,
+                        help='number of conv filters in the first layer of D')
+    parser.add_argument('--g_repeat_num', type=int, default=6,
+                        help='number of residual blocks in G')
+    parser.add_argument('--d_repeat_num', type=int, default=6,
+                        help='number of strided conv layers in D')
+    parser.add_argument('--lambda_cls', type=float, default=160,
+                        help='weight for domain classification loss')
+    parser.add_argument('--lambda_rec', type=float, default=10,
+                        help='weight for reconstruction loss')
+    parser.add_argument('--lambda_gp', type=float, default=10,
+                        help='weight for gradient penalty')
+    parser.add_argument('--lambda_sat', type=float, default=0.1,
+                        help='weight for attention saturation loss')
+    parser.add_argument('--lambda_smooth', type=float, default=1e-4,
+                        help='weight for the attention smoothing loss')
+
+    # Training configuration.
+    parser.add_argument('--batch_size', type=int,
+                        default=1, help='mini-batch size')
+    parser.add_argument('--num_epochs', type=int, default=30,
+                        help='number of total epochs for training D')
+    parser.add_argument('--num_epochs_decay', type=int, default=20,
+                        help='number of epochs for start decaying lr')
+    parser.add_argument('--g_lr', type=float, default=0.0001,
+                        help='learning rate for G')
+    parser.add_argument('--d_lr', type=float, default=0.0001,
+                        help='learning rate for D')
+    parser.add_argument('--n_critic', type=int, default=5,
+                        help='number of D updates per each G update')
+    parser.add_argument('--beta2', type=float, default=0.999,
+                        help='beta2 for Adam optimizer')
+    parser.add_argument('--beta1', type=float, default=0.5,
+                        help='beta1 for Adam optimizer')
+    parser.add_argument('--resume_iters', type=int,
+                        default=None, help='resume training from this step')
+    parser.add_argument('--first_epoch', type=int,
+                        default=0, help='First epoch')
+    parser.add_argument('--gpu_id', type=int, default=0, help='GPU id')
+    parser.add_argument('--use_virtual', type=str2bool, default=False,
+                        help='Boolean to decide if we should use the virtual cycle concistency loss')
+    # Miscellaneous.
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--mode', type=str, default='train',
+                        choices=['train', 'animation'])
+    parser.add_argument('--use_tensorboard', type=str2bool, default=False)
+    parser.add_argument('--num_sample_targets', type=int, default=4,
+                        help="number of targets to use in the samples visualization")
+
+    # Directories.
+    parser.add_argument('--image_dir', type=str,
+                        default='data/celeba/images_aligned')
+    parser.add_argument('--attr_path', type=str,
+                        default='data/celeba/list_attr_celeba.txt')
+    parser.add_argument('--outputs_dir', type=str, default='experiment1')
+    parser.add_argument('--log_dir', type=str, default='logs')
+    parser.add_argument('--model_save_dir', type=str, default='models')
+    parser.add_argument('--sample_dir', type=str, default='samples')
+    parser.add_argument('--result_dir', type=str, default='results')
+
+    parser.add_argument('--animation_images_dir', type=str,
+                        default='animations/eric_andre/images_to_animate')
+    parser.add_argument('--animation_attribute_images_dir', type=str,
+                        default='animations/eric_andre/attribute_images')
+    parser.add_argument('--animation_attributes_path', type=str,
+                        default='animations/eric_andre/attributes.txt')
+    parser.add_argument('--animation_models_dir', type=str,
+                        default='models')
+    parser.add_argument('--animation_results_dir', type=str,
+                        default='animations/eric_andre/results')
+    parser.add_argument('--animation_mode', type=str, default='animate_image',
+                        choices=['animate_image', 'animate_random_batch'])
+
+    # Step size.
+    parser.add_argument('--log_step', type=int, default=10)
+    parser.add_argument('--sample_step', type=int, default=200)
+    parser.add_argument('--model_save_step', type=int, default=1000)
+
+    config = parser.parse_args()
+    return config
+
+
+def str2bool(v):
+    return v.lower() in ('true')
@@ -0,0 +1,86 @@
+from torch.utils import data
+from torchvision import transforms as T
+from torchvision.datasets import ImageFolder
+from PIL import Image
+import torch
+import os
+import random
+import numpy as np
+
+
+class CelebA(data.Dataset):
+
+    def __init__(self, image_dir, attr_path, transform, mode, c_dim):
+
+        self.image_dir = image_dir
+        self.attr_path = attr_path
+        self.transform = transform
+        self.mode = mode
+        self.c_dim = c_dim
+
+        self.train_dataset = []
+        self.test_dataset = []
+
+        # Fills train_dataset and test_dataset --> [filename, boolean attribute vector]
+        self.preprocess()
+
+        if mode == 'train':
+            self.num_images = len(self.train_dataset)
+        else:
+            self.num_images = len(self.test_dataset)
+
+        print("------------------------------------------------")
+        print("Training images: ", len(self.train_dataset))
+        print("Testing images: ", len(self.test_dataset))
+
+    def preprocess(self):
+        lines = [line.rstrip() for line in open(self.attr_path, 'r')]
+        lines = lines[2:]
+
+        random.seed(1234)
+        random.shuffle(lines)
+
+        # Extract the info from each line
+        for idx, line in enumerate(lines):
+            split = line.split()
+            filename = split[0]
+            values = split[1:]
+            label = []  # Vector representing the presence of each attribute in each image
+
+            for n in range(self.c_dim):
+                label.append(float(values[n])/5.)
+
+            if idx < 100:
+                self.test_dataset.append([filename, label])
+            else:
+                self.train_dataset.append([filename, label])
+
+        print('Dataset ready!...')
+
+    def __getitem__(self, index):
+        dataset = self.train_dataset if self.mode == 'train' else self.test_dataset
+        filename, label = dataset[index]
+        image = Image.open(os.path.join(self.image_dir, filename))
+        return self.transform(image), torch.FloatTensor(label)
+
+    def __len__(self):
+        return self.num_images
+
+
+def get_loader(image_dir, attr_path, c_dim, image_size=128,
+               batch_size=25, mode='train', num_workers=1):
+
+    transform = []
+    transform.append(T.ToTensor())
+    transform.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
+    transform = T.Compose(transform)
+
+    dataset = CelebA(image_dir, attr_path, transform, mode, c_dim)
+
+    data_loader = data.DataLoader(dataset=dataset,
+                                  batch_size=batch_size,
+                                  shuffle=True,
+                                  num_workers=num_workers,
+                                  drop_last=True)
+
+    return data_loader
@@ -0,0 +1,19 @@
+name: ganimation
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - numpy=1.16.3
+  - opencv=4.1.0
+  - openssl=1.1.1c
+  - pillow=6.0.0
+  - pip=19.1
+  - python=3.6.8
+  - pytorch=1.1.0
+  - scipy=1.3.0
+  - tensorboard=1.13.1
+  - tensorboardx=1.7
+  - tensorflow=1.13.1
+  - torchvision=0.3.0
+
@@ -0,0 +1,34 @@
+import tensorflow as tf
+import numpy as np
+
+
+class Logger(object):
+    """Tensorboard logger."""
+
+    def __init__(self, log_dir):
+        """Initialize summary writer."""
+        self.writer = tf.summary.FileWriter(log_dir)
+
+    def scalar_summary(self, tag, value, step):
+        """Add scalar summary."""
+        summary = tf.Summary(
+            value=[tf.Summary.Value(tag=tag, simple_value=value)])
+        self.writer.add_summary(summary, step)
+
+    def image_summary(self, name, x, step):
+        x = x.numpy()[0, :, :, :]
+        x = np.moveaxis(x, 0, -1)
+        x = np.expand_dims(x, 0)
+
+        tensor = tf.convert_to_tensor(
+            x,
+            dtype=tf.float32,
+            name=None,
+            preferred_dtype=None
+        )
+
+        print(tensor.value)
+
+        summary = tf.summary.image(name=name, tensor=tensor)
+
+        self.writer.add_summary(summary, step).eval()
@@ -0,0 +1,60 @@
+import os
+
+from config import get_config
+from solver import Solver
+from data_loader import get_loader
+from torch.backends import cudnn
+
+
+def main(config):
+    cudnn.benchmark = True  # Improves runtime if the input size is constant
+
+    config.outputs_dir = os.path.join('experiments', config.outputs_dir)
+
+    config.log_dir = os.path.join(config.outputs_dir, config.log_dir)
+    config.model_save_dir = os.path.join(
+        config.outputs_dir, config.model_save_dir)
+    config.sample_dir = os.path.join(config.outputs_dir, config.sample_dir)
+    config.result_dir = os.path.join(config.outputs_dir, config.result_dir)
+
+    data_loader = get_loader(config.image_dir, config.attr_path, config.c_dim,
+                             config.image_size, config.batch_size, config.mode,
+                             config.num_workers)
+
+    config_dict = vars(config)
+    solver = Solver(data_loader, config_dict)
+
+    if config.mode == 'train':
+        initialize_train_directories(config)
+        solver.train()
+    elif config.mode == 'animation':
+        initialize_animation_directories(config)
+        solver.animation()
+
+
+def initialize_train_directories(config):
+    if not os.path.exists('experiments'):
+        os.makedirs('experiments')
+    if not os.path.exists(config.outputs_dir):
+        os.makedirs(config.outputs_dir)
+    if not os.path.exists(config.log_dir):
+        os.makedirs(config.log_dir)
+    if not os.path.exists(config.model_save_dir):
+        os.makedirs(config.model_save_dir)
+    if not os.path.exists(config.sample_dir):
+        os.makedirs(config.sample_dir)
+    if not os.path.exists(config.result_dir):
+        os.makedirs(config.result_dir)
+
+
+def initialize_animation_directories(config):
+    if not os.path.exists(config.animation_results_dir):
+        os.makedirs(config.animation_results_dir)
+
+
+if __name__ == '__main__':
+
+    config = get_config()
+    print(config)
+
+    main(config)
@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import random
+
+from PIL import Image
+
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def init_weights(self):
+        self.apply(self._weights_init_fn)
+
+    def _weights_init_fn(self, m):
+        classname = m.__class__.__name__
+        if classname.find('Conv') != -1:
+            m.weight.data.normal_(0.0, 0.02)
+            if hasattr(m.bias, 'data'):
+                m.bias.data.fill_(0)
+        elif classname.find('BatchNorm2d') != -1:
+            m.weight.data.normal_(1.0, 0.02)
+            m.bias.data.fill_(0)
+
+
+class ResidualBlock(BaseNetwork):
+    """Residual Block with instance normalization."""
+
+    def __init__(self, dim_in, dim_out):
+        super(ResidualBlock, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(dim_in, dim_out, kernel_size=3,
+                      stride=1, padding=1, bias=False),
+            nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(dim_out, dim_out, kernel_size=3,
+                      stride=1, padding=1, bias=False),
+            nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True))
+
+        self.init_weights()
+
+    def forward(self, x):
+        return x + self.main(x)
+
+
+class Generator(BaseNetwork):
+    """Generator network."""
+
+    def __init__(self, conv_dim=64, c_dim=5, repeat_num=6):
+        super(Generator, self).__init__()
+
+        layers = []
+        layers.append(nn.Conv2d(3+c_dim, conv_dim, kernel_size=7,
+                                stride=1, padding=3, bias=False))
+        layers.append(nn.InstanceNorm2d(
+            conv_dim, affine=True, track_running_stats=True))
+        layers.append(nn.ReLU(inplace=True))
+
+        self.debug1 = nn.Sequential(*layers)
+
+        # Down-sampling layers.
+        curr_dim = conv_dim
+        for i in range(2):
+            layers.append(nn.Conv2d(curr_dim, curr_dim*2,
+                                    kernel_size=4, stride=2, padding=1, bias=False))
+            layers.append(nn.InstanceNorm2d(
+                curr_dim*2, affine=True, track_running_stats=True))
+            layers.append(nn.ReLU(inplace=True))
+            curr_dim = curr_dim * 2
+
+        self.debug2 = nn.Sequential(*layers)
+
+        # Bottleneck layers.
+        for i in range(repeat_num):
+            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))
+
+        self.debug3 = nn.Sequential(*layers)
+
+        # Up-sampling layers.
+        for i in range(2):
+            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim //
+                                             2, kernel_size=4, stride=2, padding=1, bias=False))
+            layers.append(nn.InstanceNorm2d(
+                curr_dim//2, affine=True, track_running_stats=True))
+            layers.append(nn.ReLU(inplace=True))
+            curr_dim = curr_dim // 2
+
+        self.main = nn.Sequential(*layers)
+
+        self.debug4 = nn.Sequential(*layers)
+
+        # Same architecture for the color regression
+        layers = []
+        layers.append(nn.Conv2d(curr_dim, 3, kernel_size=7,
+                                stride=1, padding=3, bias=False))
+        layers.append(nn.Tanh())
+        self.im_reg = nn.Sequential(*layers)
+
+        # One Channel output and Sigmoid function for the attention layer
+        layers = []
+        layers.append(nn.Conv2d(curr_dim, 1, kernel_size=7,
+                                stride=1, padding=3, bias=False))
+        layers.append(nn.Sigmoid())  # Values between 0 and 1
+        self.im_att = nn.Sequential(*layers)
+
+        self.init_weights()
+
+    def forward(self, x, c):
+        # Replicate spatially and concatenate domain information.
+
+        c = c.unsqueeze(2).unsqueeze(3)
+        c = c.expand(c.size(0), c.size(1), x.size(2), x.size(3))
+
+        x = torch.cat([x, c], dim=1)
+        features = self.main(x)
+
+        reg = self.im_reg(features)
+        att = self.im_att(features)
+
+        return att, reg
+
+
+class Discriminator(BaseNetwork):
+    """Discriminator network with PatchGAN."""
+
+    def __init__(self, image_size=128, conv_dim=64, c_dim=5, repeat_num=6):
+        super(Discriminator, self).__init__()
+
+        layers = []
+        layers.append(
+            nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
+        layers.append(nn.LeakyReLU(0.01))
+
+        curr_dim = conv_dim
+        for i in range(1, repeat_num):
+            layers.append(nn.Conv2d(curr_dim, curr_dim*2,
+                                    kernel_size=4, stride=2, padding=1))
+            layers.append(nn.LeakyReLU(0.01))
+            curr_dim = curr_dim * 2
+
+        kernel_size = int(image_size / np.power(2, repeat_num))
+        self.main = nn.Sequential(*layers)
+        self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(
+            curr_dim, c_dim, kernel_size=kernel_size, bias=False)
+
+        self.init_weights()
+
+    def forward(self, x):
+        h = self.main(x)
+        out_src = self.conv1(h)
+        out_cls = self.conv2(h)
+
+        # out_cls.view(out_cls.size(0), out_cls.size(1))
+        return out_src.squeeze(), out_cls.squeeze()
@@ -0,0 +1 @@
+/scratch2/ganimation/models
@@ -0,0 +1,11 @@
+import glob
+
+output_txt = 'list_attr_mydataset.txt'
+
+for idx, f in enumerate(glob.glob('./my-processed-dataset/*.csv')):
+    with open(f, 'r') as csv_file:
+        csv_file.readline()
+        csv_list = csv_file.readline().split(', ')
+        if float(csv_list[1]) >= 0.88:
+            aus = " ".join(csv_list[2:19])
+            open(output_txt, 'a').write(f.split('/')[-1].split('.')[0] + '.jpg ' + aus + '\n')
@@ -0,0 +1,12 @@
+import os
+import shutil
+
+output_dir = './images'
+
+os.mkdir(output_dir)
+
+for root, dirs, files in os.walk('./my-processed-dataset'):
+    for filename in files:
+        if 'jpg' in filename:
+            img_name = root.split('/')[-1].split('_')[0] + '.jpg'
+            shutil.copy2(os.path.join(root, filename), os.path.join(output_dir, img_name))
@@ -0,0 +1,455 @@
+import torch
+import torch.nn.functional as F
+from torchvision.utils import save_image
+
+from model import Generator, Discriminator
+from utils import Utils
+
+import numpy as np
+
+import os
+import time
+import datetime
+import random
+import glob
+
+import attacks
+
+
+class Solver(Utils):
+
+    def __init__(self, data_loader, config_dict):
+        # NOTE: the following line create new class arguments with the
+        # values in config_dict
+        self.__dict__.update(**config_dict)
+
+        self.data_loader = data_loader
+
+        self.device = 'cuda:' + \
+            str(self.gpu_id) if torch.cuda.is_available() else 'cpu'
+        print(f"Model running on {self.device}")
+
+        if self.use_tensorboard:
+            self.build_tensorboard()
+
+        self.loss_visualization = {}
+
+        self.build_model()
+
+    def train(self):
+        print('Training...')
+
+        self.global_counter = 0
+
+        if self.resume_iters:
+            self.first_iteration = self.resume_iters
+            self.restore_model(self.resume_iters)
+        else:
+            self.first_iteration = 0
+
+        self.start_time = time.time()
+
+        for epoch in range(self.first_epoch, self.num_epochs):
+            print(f"EPOCH {epoch} WITH {len(self.data_loader)} STEPS")
+
+            self.alpha_rec = 1
+            self.epoch = epoch
+
+            for iteration in range(self.first_iteration, len(self.data_loader)):
+                self.iteration = iteration
+                self.get_training_data()
+                self.train_discriminator()
+
+                if (self.iteration+1) % self.n_critic == 0:
+                    generation_outputs = self.train_generator()
+
+                if (self.iteration+1) % self.sample_step == 0:
+                    self.print_generations(generation_outputs)
+
+                if self.iteration % self.model_save_step == 0:
+                    self.save_models(self.iteration, self.epoch)
+
+                if self.iteration % self.log_step == 0:
+                    self.update_tensorboard()
+                self.global_counter += 1
+
+            # Decay learning rates.
+            if (self.epoch+1) > self.num_epochs_decay:
+                # float(self.num_epochs_decay))
+                self.g_lr -= (self.g_lr / 10.0)
+                # float(self.num_epochs_decay))
+                self.d_lr -= (self.d_lr / 10.0)
+                self.update_lr(self.g_lr, self.d_lr)
+                print('Decayed learning rates, self.g_lr: {}, self.d_lr: {}.'.format(
+                    self.g_lr, self.d_lr))
+
+            # Save the last model
+            self.save_models()
+
+            self.first_iteration = 0  # Next epochs start from 0
+
+    def get_training_data(self):
+        try:
+            self.x_real, self.label_org = next(self.data_iter)
+        except:
+            self.data_iter = iter(self.data_loader)
+            self.x_real, self.label_org = next(self.data_iter)
+
+        self.x_real = self.x_real.to(self.device)  # Input images.
+        # Labels for computing classification loss.
+        self.label_org = self.label_org.to(self.device)
+
+        # Get random targets for training
+        self.label_trg = self.get_random_labels_list()
+        self.label_trg = torch.FloatTensor(self.label_trg).clamp(0, 1)
+        # Labels for computing classification loss.
+        self.label_trg = self.label_trg.to(self.device)
+
+        if self.use_virtual:
+            self.label_trg_virtual = self.get_random_labels_list()
+            self.label_trg_virtual = torch.FloatTensor(
+                self.label_trg_virtual).clamp(0, 1)
+            # Labels for computing classification loss.
+            self.label_trg_virtual = self.label_trg_virtual.to(self.device)
+
+            assert not torch.equal(
+                self.label_trg_virtual, self.label_trg), "Target label and virtual label are the same"
+
+    def get_random_labels_list(self):
+        trg_list = []
+        for _ in range(self.batch_size):
+            random_num = random.randint(
+                0, len(self.data_loader)*self.batch_size-1)
+            # Select a random AU vector from the dataset
+            trg_list_aux = self.data_loader.dataset[random_num][1]
+            # Apply a variance of 0.1 to the vector
+            trg_list.append(trg_list_aux.numpy() +
+                            np.random.uniform(-0.1, 0.1, trg_list_aux.shape))
+        return trg_list
+
+    def train_discriminator(self):
+        # Compute loss with real images.
+        critic_output, classification_output = self.D(self.x_real)
+        d_loss_critic_real = -torch.mean(critic_output)
+        d_loss_classification = torch.nn.functional.mse_loss(
+            classification_output, self.label_org)
+
+        # Compute loss with fake images.
+        attention_mask, color_regression = self.G(self.x_real, self.label_trg)
+        x_fake = self.imFromAttReg(
+            attention_mask, color_regression, self.x_real)
+        critic_output, _ = self.D(x_fake.detach())
+        d_loss_critic_fake = torch.mean(critic_output)
+
+        # Compute loss for gradient penalty.
+        alpha = torch.rand(self.x_real.size(0), 1, 1, 1).to(self.device)
+        # Half of image info from fake and half from real
+        x_hat = (alpha * self.x_real.data + (1 - alpha)
+                 * x_fake.data).requires_grad_(True)
+        critic_output, _ = self.D(x_hat)
+        d_loss_gp = self.gradient_penalty(critic_output, x_hat)
+
+        # Backward and optimize.
+        d_loss = d_loss_critic_real + d_loss_critic_fake + self.lambda_cls * \
+            d_loss_classification + self.lambda_gp * d_loss_gp
+
+        self.reset_grad()
+        d_loss.backward()
+        self.d_optimizer.step()
+
+        # Logging.
+        self.loss_visualization['D/loss'] = d_loss.item()
+        self.loss_visualization['D/loss_real'] = d_loss_critic_real.item()
+        self.loss_visualization['D/loss_fake'] = d_loss_critic_fake.item()
+        self.loss_visualization['D/loss_cls'] = self.lambda_cls * \
+            d_loss_classification.item()
+        self.loss_visualization['D/loss_gp'] = self.lambda_gp * \
+            d_loss_gp.item()
+
+    def train_generator(self):
+        # Original-to-target domain.
+        attention_mask, color_regression = self.G(self.x_real, self.label_trg)
+        x_fake = self.imFromAttReg(
+            attention_mask, color_regression, self.x_real)
+
+        critic_output, classification_output = self.D(x_fake)
+        g_loss_fake = -torch.mean(critic_output)
+        g_loss_cls = torch.nn.functional.mse_loss(
+            classification_output, self.label_trg)
+
+        # Target-to-original domain.
+        if not self.use_virtual:
+            reconstructed_attention_mask, reconstructed_color_regression = self.G(
+                x_fake, self.label_org)
+            x_rec = self.imFromAttReg(
+                reconstructed_attention_mask, reconstructed_color_regression, x_fake)
+
+        else:
+            reconstructed_attention_mask, reconstructed_color_regression = self.G(
+                x_fake, self.label_org)
+            x_rec = self.imFromAttReg(
+                reconstructed_attention_mask, reconstructed_color_regression, x_fake)
+
+            reconstructed_attention_mask_2, reconstructed_color_regression_2 = self.G(
+                x_fake, self.label_trg_virtual)
+            x_fake_virtual = self.imFromAttReg(
+                reconstructed_attention_mask_2, reconstructed_color_regression_2, x_fake)
+
+            reconstructed_virtual_attention_mask, reconstructed_virtual_color_regression = self.G(
+                x_fake_virtual, self.label_trg)
+            x_rec_virtual = self.imFromAttReg(
+                reconstructed_virtual_attention_mask, reconstructed_virtual_color_regression, x_fake_virtual.detach())
+
+        # Compute losses
+        g_loss_saturation_1 = attention_mask.mean()
+        g_loss_smooth1 = self.smooth_loss(attention_mask)
+
+        if not self.use_virtual:
+            g_loss_rec = torch.nn.functional.l1_loss(self.x_real, x_rec)
+            g_loss_saturation_2 = reconstructed_attention_mask.mean()
+            g_loss_smooth2 = self.smooth_loss(reconstructed_attention_mask)
+
+        else:
+            g_loss_rec = (1-self.alpha_rec)*torch.nn.functional.l1_loss(self.x_real, x_rec) + \
+                self.alpha_rec * \
+                torch.nn.functional.l1_loss(x_fake, x_rec_virtual)
+
+            g_loss_saturation_2 = (1-self.alpha_rec) * reconstructed_attention_mask.mean() + \
+                self.alpha_rec * reconstructed_virtual_attention_mask.mean()
+
+            g_loss_smooth2 = (1-self.alpha_rec) * self.smooth_loss(reconstructed_virtual_attention_mask) + \
+                self.alpha_rec * self.smooth_loss(reconstructed_attention_mask)
+
+        g_attention_loss = self.lambda_smooth * g_loss_smooth1 + self.lambda_smooth * g_loss_smooth2 \
+            + self.lambda_sat * g_loss_saturation_1 + self.lambda_sat * g_loss_saturation_2
+
+        g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + \
+            self.lambda_cls * g_loss_cls + g_attention_loss
+
+        self.reset_grad()
+        g_loss.backward()
+        self.g_optimizer.step()
+
+        # Logging.
+        self.loss_visualization['G/loss'] = g_loss.item()
+        self.loss_visualization['G/loss_fake'] = g_loss_fake.item()
+        self.loss_visualization['G/loss_rec'] = self.lambda_rec * \
+            g_loss_rec.item()
+        self.loss_visualization['G/loss_cls'] = self.lambda_cls * \
+            g_loss_cls.item()
+        self.loss_visualization['G/attention_loss'] = g_attention_loss.item()
+        self.loss_visualization['G/loss_smooth1'] = self.lambda_smooth * \
+            g_loss_smooth1.item()
+        self.loss_visualization['G/loss_smooth2'] = self.lambda_smooth * \
+            g_loss_smooth2.item()
+        self.loss_visualization['G/loss_sat1'] = self.lambda_sat * \
+            g_loss_saturation_1.item()
+        self.loss_visualization['G/loss_sat2'] = self.lambda_sat * \
+            g_loss_saturation_2.item()
+        self.loss_visualization['G/alpha'] = self.alpha_rec
+
+        if not self.use_virtual:
+            return {
+                "color_regression": color_regression,
+                "x_fake": x_fake,
+                "attention_mask": attention_mask,
+                "x_rec": x_rec,
+                "reconstructed_attention_mask": reconstructed_attention_mask,
+                "reconstructed_attention_mask": reconstructed_attention_mask,
+                "reconstructed_color_regression": reconstructed_color_regression,
+            }
+
+        else:
+            return {
+                "color_regression": color_regression,
+                "x_fake": x_fake,
+                "attention_mask": attention_mask,
+                "x_rec": x_rec,
+                "reconstructed_attention_mask": reconstructed_attention_mask,
+                "reconstructed_attention_mask": reconstructed_attention_mask,
+                "reconstructed_color_regression": reconstructed_color_regression,
+                "reconstructed_virtual_attention_mask": reconstructed_virtual_attention_mask,
+                "reconstructed_virtual_color_regression": reconstructed_virtual_color_regression,
+                "x_rec_virtual": x_rec_virtual,
+            }
+
+    def print_generations(self, generator_outputs_dict):
+        print_epoch_images = False
+        save_image(self.denorm(self.x_real), self.sample_dir +
+                   '/{}_4real_.png'.format(self.epoch))
+        save_image((generator_outputs_dict["color_regression"]+1)/2,
+                   self.sample_dir + '/{}_2reg_.png'.format(self.epoch))
+        save_image(self.denorm(
+            generator_outputs_dict["x_fake"]), self.sample_dir + '/{}_3res_.png'.format(self.epoch))
+        save_image(generator_outputs_dict["attention_mask"],
+                   self.sample_dir + '/{}_1attention_.png'.format(self.epoch))
+        save_image(self.denorm(
+            generator_outputs_dict["x_rec"]), self.sample_dir + '/{}_5rec_.png'.format(self.epoch))
+
+        if not self.use_virtual:
+            save_image(generator_outputs_dict["reconstructed_attention_mask"],
+                       self.sample_dir + '/{}_6rec_attention.png'.format(self.epoch))
+            save_image(self.denorm(
+                generator_outputs_dict["reconstructed_color_regression"]), self.sample_dir + '/{}_7rec_reg.png'.format(self.epoch))
+
+        else:
+            save_image(generator_outputs_dict["reconstructed_attention_mask"],
+                       self.sample_dir + '/{}_6rec_attention_.png'.format(self.epoch))
+            save_image(self.denorm(
+                generator_outputs_dict["reconstructed_color_regression"]), self.sample_dir + '/{}_7rec_reg.png'.format(self.epoch))
+
+            save_image(generator_outputs_dict["reconstructed_virtual_attention_mask"],
+                       self.sample_dir + '/{}_8rec_virtual_attention.png'.format(self.epoch))
+            save_image(self.denorm(generator_outputs_dict["reconstructed_virtual_color_regression"]),
+                       self.sample_dir + '/{}_91rec_virtual_reg.png'.format(self.epoch))
+            save_image(self.denorm(
+                generator_outputs_dict["x_rec_virtual"]), self.sample_dir + '/{}_92rec_epoch_.png'.format(self.epoch))
+
+    def update_tensorboard(self):
+        # Print out training information.
+        et = time.time() - self.start_time
+        et = str(datetime.timedelta(seconds=et))[:-7]
+        log = "Elapsed [{}],  [{}/{}], Epoch [{}/{}]".format(
+            et, self.iteration+1, len(self.data_loader), self.epoch+1, self.num_epochs)
+        for tag, value in self.loss_visualization.items():
+            log += ", {}: {:.4f}".format(tag, value)
+        print(log)
+
+        if self.use_tensorboard:
+            for tag, value in self.loss_visualization.items():
+                self.writer.add_scalar(
+                    tag, value, global_step=self.global_counter)
+
+    def animation(self, mode='animate_image'):
+        from PIL import Image
+        from torchvision import transforms as T
+
+        regular_image_transform = []
+        regular_image_transform.append(T.ToTensor())
+        regular_image_transform.append(T.Normalize(
+            mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
+        regular_image_transform = T.Compose(regular_image_transform)
+
+        G_path = sorted(glob.glob(os.path.join(
+            self.animation_models_dir, '*G.ckpt')), key=self.numericalSort)[0]
+        self.G.load_state_dict(torch.load(G_path, map_location=f'cuda:{self.gpu_id}'))
+        self.G = self.G.cuda(0)
+
+        reference_expression_images = []
+
+        with torch.no_grad():
+            with open(self.animation_attributes_path, 'r') as txt_file:
+                csv_lines = txt_file.readlines()
+
+                targets = torch.zeros(len(csv_lines), self.c_dim)
+                input_images = torch.zeros(len(csv_lines), 3, 128, 128)
+
+                for idx, line in enumerate(csv_lines):
+                    splitted_lines = line.split(' ')
+                    image_path = os.path.join(
+                        self.animation_attribute_images_dir, splitted_lines[0])
+                    input_images[idx, :] = regular_image_transform(
+                        Image.open(image_path)).cuda()
+                    reference_expression_images.append(splitted_lines[0])
+                    targets[idx, :] = torch.Tensor(
+                        np.array(list(map(lambda x: float(x)/5., splitted_lines[1::]))))
+
+        if mode == 'animate_random_batch':
+            animation_batch_size = 7
+
+            self.data_iter = iter(self.data_loader)
+            images_to_animate, _ = next(self.data_iter)
+            images_to_animate = images_to_animate[0:animation_batch_size].cuda(
+            )
+
+            for target_idx in range(targets.size(0)):
+                targets_au = targets[target_idx, :].unsqueeze(
+                    0).repeat(animation_batch_size, 1).cuda()
+                resulting_images_att, resulting_images_reg = self.G(
+                    images_to_animate, targets_au)
+
+                resulting_images = self.imFromAttReg(
+                    resulting_images_att, resulting_images_reg, images_to_animate).cuda()
+
+                save_images = - \
+                    torch.ones((animation_batch_size + 1)
+                               * 2, 3, 128, 128).cuda()
+
+                save_images[1:animation_batch_size+1] = images_to_animate
+                save_images[animation_batch_size+1] = input_images[target_idx]
+                save_images[animation_batch_size +
+                            2:(animation_batch_size + 1)*2] = resulting_images
+
+                save_image((save_images+1)/2, os.path.join(self.animation_results_dir,
+                                                           reference_expression_images[target_idx]))
+
+        if mode == 'animate_image':
+
+            black = np.zeros((1,3,128,128))
+            black = torch.FloatTensor(black).to(self.device)
+
+            pgd_attack = attacks.LinfPGDAttack(model=self.G, device=self.device)
+
+            images_to_animate_path = sorted(glob.glob(
+                self.animation_images_dir + '/*'))
+
+            for idx, image_path in enumerate(images_to_animate_path):
+                image_to_animate = regular_image_transform(
+                    Image.open(image_path)).unsqueeze(0).cuda()
+
+                all_images = torch.cat([regular_image_transform(Image.open(path)).unsqueeze(0) for path in images_to_animate_path], dim=0).cuda()
+
+                for target_idx in range(targets.size(0)):
+                    if target_idx == 0:
+                        img = regular_image_transform(Image.open(images_to_animate_path[idx])).unsqueeze(0).cuda()
+                        # x_adv, perturb = pgd_attack.perturb(img, black, targets[0, :].unsqueeze(0).cuda())
+                        x_adv, perturb = pgd_attack.perturb_iter_class(image_to_animate, black, targets[:, :].cuda())
+                        # _, perturb = pgd_attack.perturb_iter_data(image_to_animate, all_images, black, targets[68, :].unsqueeze(0).cuda())
+                        
+                    targets_au = targets[target_idx, :].unsqueeze(0).cuda()
+                    # x_adv, perturb = pgd_attack.perturb(image_to_animate, black, targets_au)
+                    x_adv = image_to_animate
+                    # print(image_to_animate.shape, x_adv.shape)
+                    resulting_images_att, resulting_images_reg = self.G(
+                        x_adv, targets_au)
+                    resulting_image = self.imFromAttReg(
+                        resulting_images_att, resulting_images_reg, x_adv).cuda()
+
+                    save_image((resulting_image+1)/2, os.path.join(self.animation_results_dir,
+                                                                   image_path.split('/')[-1].split('.')[0]
+                                                                   + '_' + reference_expression_images[target_idx]))
+                    if target_idx == 0:
+                        save_image((x_adv+1)/2, os.path.join(self.animation_results_dir,
+                                                                   image_path.split('/')[-1].split('.')[0]
+                                                                   + '_ref.jpg'))
+
+        # """ Code to modify single Action Units """
+
+        # Set data loader.
+        # self.data_loader = self.data_loader
+
+        # with torch.no_grad():
+        #     for i, (self.x_real, c_org) in enumerate(self.data_loader):
+
+        #         # Prepare input images and target domain labels.
+        #         self.x_real = self.x_real.to(self.device)
+        #         c_org = c_org.to(self.device)
+
+        #         # c_trg_list = self.create_labels(self.data_loader)
+
+        #         crit, cl_regression = self.D(self.x_real)
+        #         # print(crit)
+        #         print("ORIGINAL", c_org[0])
+        #         print("REGRESSION", cl_regression[0])
+
+        #         for au in range(17):
+        #             alpha = np.linspace(-0.3,0.3,10)
+        #             for j, a in enumerate(alpha):
+        #                 new_emotion = c_org.clone()
+        #                 new_emotion[:,au]=torch.clamp(new_emotion[:,au]+a, 0, 1)
+        #                 attention, reg = self.G(self.x_real, new_emotion)
+        #                 x_fake = self.imFromAttReg(attention, reg, self.x_real)
+        #                 save_image((x_fake+1)/2, os.path.join(self.result_dir, '{}-{}-{}-images.jpg'.format(i,au,j)))
+
+        #         if i >= 3:
+        #             break
@@ -0,0 +1,151 @@
+import torch
+import torch.nn.functional as F
+from tensorboardX import SummaryWriter
+
+from model import Generator
+from model import Discriminator
+
+import os
+import re
+
+
+class Utils:
+
+    def build_model(self):
+        """Create a generator and a discriminator."""
+        self.G = Generator(self.g_conv_dim, self.c_dim,
+                           self.g_repeat_num).to(self.device)
+        self.D = Discriminator(
+            self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num).to(self.device)
+
+        self.g_optimizer = torch.optim.Adam(
+            self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
+        self.d_optimizer = torch.optim.Adam(
+            self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
+
+        # TODO: implement data parallelization for multiple gpus
+        # self.gpu_ids = torch.cuda.device_count()
+        # print("GPUS AVAILABLE: ", self.gpu_ids)
+        # if self.gpu_ids > 1:
+        #     torch.nn.DataParallel(self.D, device_ids=list(range(self.gpu_ids)))
+        #     torch.nn.DataParallel(self.G, device_ids=list(range(self.gpu_ids)))
+
+    def build_tensorboard(self):
+        """Build a tensorboard logger."""
+        from logger import Logger
+        self.logger = Logger(self.log_dir)
+        self.writer = SummaryWriter(logdir=self.log_dir)
+
+    def smooth_loss(self, att):
+        return torch.mean(torch.mean(torch.abs(att[:, :, :, :-1] - att[:, :, :, 1:])) +
+                          torch.mean(torch.abs(att[:, :, :-1, :] - att[:, :, 1:, :])))
+
+    def print_network(self, model, name):
+        """Print out the network information."""
+        num_params = 0
+        for p in model.parameters():
+            num_params += p.numel()
+        print(model)
+        print(name)
+        print("The number of parameters: {}".format(num_params))
+
+    def update_lr(self, g_lr, d_lr):
+        """Decay learning rates of the generator and discriminator."""
+        for param_group in self.g_optimizer.param_groups:
+            param_group['lr'] = g_lr
+        for param_group in self.d_optimizer.param_groups:
+            param_group['lr'] = d_lr
+
+    def reset_grad(self):
+        """Reset the gradient buffers."""
+        self.g_optimizer.zero_grad()
+        self.d_optimizer.zero_grad()
+
+    def denorm(self, x):
+        """Convert the range from [-1, 1] to [0, 1]."""
+        out = (x + 1) / 2
+        return out.clamp_(0, 1)
+
+    def gradient_penalty(self, y, x):
+        """Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
+        weight = torch.ones(y.size()).to(self.device)
+        dydx = torch.autograd.grad(outputs=y,
+                                   inputs=x,
+                                   grad_outputs=weight,
+                                   retain_graph=True,
+                                   create_graph=True,
+                                   only_inputs=True)[0]
+
+        dydx = dydx.view(dydx.size(0), -1)
+        dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
+        return torch.mean((dydx_l2norm-1)**2)
+
+    def imFromAttReg(self, att, reg, x_real):
+        """Mixes attention, color and real images"""
+        return (1-att)*reg + att*x_real
+
+    def create_labels(self, data_iter):
+        """Return samples for visualization"""
+        x, c = [], []
+        x_data, c_data = data_iter.next()
+
+        for i in range(self.num_sample_targets):
+            x.append(x_data[i].repeat(
+                self.batch_size, 1, 1, 1).to(self.device))
+            c.append(c_data[i].repeat(self.batch_size, 1).to(self.device))
+
+        return x, c
+
+    def save_models(self, iteration, epoch):
+        try:  # To avoid crashing on the first step
+            os.remove(os.path.join(self.model_save_dir,
+                                   '{}-{}-G.ckpt'.format(iteration+1-self.model_save_step, epoch)))
+            os.remove(os.path.join(self.model_save_dir,
+                                   '{}-{}-D.ckpt'.format(iteration+1-self.model_save_step, epoch)))
+            os.remove(os.path.join(self.model_save_dir,
+                                   '{}-{}-G_optim.ckpt'.format(iteration+1-self.model_save_step, epoch)))
+            os.remove(os.path.join(self.model_save_dir,
+                                   '{}-{}-D_optim.ckpt'.format(iteration+1-self.model_save_step, epoch)))
+        except:
+            pass
+
+        G_path = os.path.join(self.model_save_dir,
+                              '{}-{}-G.ckpt'.format(iteration+1, epoch))
+        D_path = os.path.join(self.model_save_dir,
+                              '{}-{}-D.ckpt'.format(iteration+1, epoch))
+        torch.save(self.G.state_dict(), G_path)
+        torch.save(self.D.state_dict(), D_path)
+
+        G_path_optim = os.path.join(
+            self.model_save_dir, '{}-{}-G_optim.ckpt'.format(iteration+1, epoch))
+        D_path_optim = os.path.join(
+            self.model_save_dir, '{}-{}-D_optim.ckpt'.format(iteration+1, epoch))
+        torch.save(self.g_optimizer.state_dict(), G_path_optim)
+        torch.save(self.d_optimizer.state_dict(), D_path_optim)
+
+        print(f'Saved model checkpoints in {self.model_save_dir}...')
+
+    def restore_model(self, resume_iters):
+        """Restore the trained generator and discriminator."""
+        print('Loading the trained models from step {}-{}...'.format(resume_iters, self.first_epoch))
+        G_path = os.path.join(
+            self.model_save_dir, '{}-{}-G.ckpt'.format(resume_iters, self.first_epoch))
+        D_path = os.path.join(
+            self.model_save_dir, '{}-{}-D.ckpt'.format(resume_iters, self.first_epoch))
+        self.G.load_state_dict(torch.load(
+            G_path, map_location=lambda storage, loc: storage))
+        self.D.load_state_dict(torch.load(
+            D_path, map_location=lambda storage, loc: storage))
+
+        G_optim_path = os.path.join(
+            self.model_save_dir, '{}-{}-G_optim.ckpt'.format(resume_iters, self.first_epoch))
+        D_optim_path = os.path.join(
+            self.model_save_dir, '{}-{}-D_optim.ckpt'.format(resume_iters, self.first_epoch))
+        self.d_optimizer.load_state_dict(torch.load(D_optim_path))
+        self.g_optimizer.load_state_dict(torch.load(G_optim_path))
+
+    def numericalSort(self, value):
+        numbers = re.compile(r'(\d+)')
+        parts = numbers.split(value)
+        parts[1::2] = map(int, parts[1::2])
+        return parts