{"cells":[{"metadata":{},"cell_type":"markdown","source":"This notebook code mostly relies on the implementation of UNIT found here:\nhttps://github.com/eriklindernoren/PyTorch-GAN/tree/master/implementations/unit"},{"metadata":{"trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nfrom torch.autograd import Variable\nimport numpy as np\nfrom torchvision.utils import make_grid\n\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\n\n\ndef weights_init_normal(m):\n    classname = m.__class__.__name__\n    if classname.find(\"Conv\") != -1:\n        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find(\"BatchNorm2d\") != -1:\n        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n        torch.nn.init.constant_(m.bias.data, 0.0)\n\n\nclass LambdaLR:\n    def __init__(self, n_epochs, offset, decay_start_epoch):\n        assert (n_epochs - decay_start_epoch) > 0, \"Decay must start before the training session ends!\"\n        self.n_epochs = n_epochs\n        self.offset = offset\n        self.decay_start_epoch = decay_start_epoch\n\n    def step(self, epoch):\n        return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)\n\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, features):\n        super(ResidualBlock, self).__init__()\n\n        conv_block = [\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(features, features, 3),\n            nn.InstanceNorm2d(features),\n            nn.ReLU(inplace=True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(features, features, 3),\n            nn.InstanceNorm2d(features),\n        ]\n\n        self.conv_block = nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        return x + self.conv_block(x)\n\n\nclass Encoder(nn.Module):\n    def __init__(self, in_channels=3, dim=64, n_downsample=2, shared_block=None):\n        super(Encoder, self).__init__()\n\n        # Initial convolution block\n        layers = [\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(in_channels, dim, 7),\n            nn.InstanceNorm2d(64),\n            nn.LeakyReLU(0.2, inplace=True),\n        ]\n\n        # Downsampling\n        for _ in range(n_downsample):\n            layers += [\n                nn.Conv2d(dim, dim * 2, 4, stride=2, padding=1),\n                nn.InstanceNorm2d(dim * 2),\n                nn.ReLU(inplace=True),\n            ]\n            dim *= 2\n\n        # Residual blocks\n        for _ in range(3):\n            layers += [ResidualBlock(dim)]\n\n        self.model_blocks = nn.Sequential(*layers)\n        self.shared_block = shared_block\n\n    def reparameterization(self, mu):\n        Tensor = torch.cuda.FloatTensor if mu.is_cuda else torch.FloatTensor\n        z = Variable(Tensor(np.random.normal(0, 1, mu.shape)))\n        return z + mu\n\n    def forward(self, x):\n        x = self.model_blocks(x)\n        mu = self.shared_block(x)\n        z = self.reparameterization(mu)\n        return mu, z\n\n\nclass Generator(nn.Module):\n    def __init__(self, out_channels=3, dim=64, n_upsample=2, shared_block=None):\n        super(Generator, self).__init__()\n\n        self.shared_block = shared_block\n\n        layers = []\n        dim = dim * 2 ** n_upsample\n        # Residual blocks\n        for _ in range(3):\n            layers += [ResidualBlock(dim)]\n\n        # Upsampling\n        for _ in range(n_upsample):\n            layers += [\n                nn.ConvTranspose2d(dim, dim // 2, 4, stride=2, padding=1),\n                nn.InstanceNorm2d(dim // 2),\n                nn.LeakyReLU(0.2, inplace=True),\n            ]\n            dim = dim // 2\n\n        # Output layer\n        layers += [nn.ReflectionPad2d(3), nn.Conv2d(dim, out_channels, 7), nn.Tanh()]\n\n        self.model_blocks = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.shared_block(x)\n        x = self.model_blocks(x)\n        return x\n\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, input_shape):\n        super(Discriminator, self).__init__()\n        channels, height, width = input_shape\n        # Calculate output of image discriminator (PatchGAN)\n        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)\n\n        def discriminator_block(in_filters, out_filters, normalize=True):\n            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n            if normalize:\n                layers.append(nn.InstanceNorm2d(out_filters))\n            layers.append(nn.LeakyReLU(0.2, inplace=True))\n            return layers\n\n        self.model = nn.Sequential(\n            *discriminator_block(channels, 64, normalize=False),\n            *discriminator_block(64, 128),\n            *discriminator_block(128, 256),\n            *discriminator_block(256, 512),\n            nn.Conv2d(512, 1, 3, padding=1)\n        )\n\n    def forward(self, img):\n        return self.model(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import glob\nimport random\nimport os\n\nfrom torch.utils.data import Dataset\nfrom PIL import Image\nimport torchvision.transforms as transforms\n\n\nclass ImageDataset(Dataset):\n    def __init__(self, root, transforms_=None, unaligned=True, mode=\"train\"):\n        self.transform = transforms.Compose(transforms_)\n        self.unaligned = unaligned\n        #print(os.path.join(root, \"monet_jpg\") + \"/*.*\")\n        self.files_A = sorted(glob.glob(os.path.join(root, \"monet_jpg\") + \"/*.*\"))\n        self.files_B = sorted(glob.glob(os.path.join(root, \"photo_jpg\") + \"/*.*\"))\n\n        #print(\"taille fichier max\", max(len(self.files_A), len(self.files_B)))\n        #print(sorted(glob.glob(os.path.join(root, \"monet_jpg\") + \"/*.*\")))\n    def __getitem__(self, index):\n        item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]))\n\n        if self.unaligned:\n            item_B = self.transform(Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)]))\n        else:\n            item_B = self.transform(Image.open(self.files_B[index % len(self.files_B)]))\n\n        return {\"monet_jpg\": item_A, \"photo_jpg\": item_B}\n\n    def __len__(self):\n        \n        return min(len(self.files_A), len(self.files_B))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import argparse\nimport os\nimport numpy as np\nimport math\nimport itertools\nimport datetime\nimport time\nimport sys\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import save_image\n\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torch.autograd import Variable\n\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\n\ncuda = True if torch.cuda.is_available() else False\n\ndataset_name ='../input/gan-getting-started'     # default apple2orange help=\"name of the dataset\" \nn_epochs = 15\ndim = 64\n\nb1 = 0.5  #\"adam: decay of first order momentum of gradient\")\nb2 = 0.999\nepoch = 0\n  \n\nsavings = \"/content/gdrive/My Drive/deep_project\"\nbatch_size=1\nlr = 0.005\nb1 = 0.5  #\"adam: decay of first order momentum of gradient\")\nb2 = 0.999\ndecay_epoch = 2000  \nn_cpu = 1     #\"number of cpu threads to use during batch generation\"\nimg_height = 256\nimg_width = 256\nchannels = 3\nsample_interval = 2000 #\"interval between saving generator samples\"\ncheckpoint_interval = 1000  #interval between saving model checkpoints\n\n\nLR = [0.0002] #learning rates\nn_downsample = 1 #change to 2 originally\n\ndef reverse_normalize(image, mean_=0.5, std_=0.5):\n    #credit: https://www.kaggle.com/adrianda/cyclegan-pytorch-style-transfer\n    if torch.is_tensor(image):\n        image = image.detach().numpy()\n    un_normalized_img = image * std_ + mean_\n    un_normalized_img = un_normalized_img * 255\n    return np.uint8(un_normalized_img)\n\n\ndef sample_images(batches_done):\n    \"\"\"Saves a generated sample from the test set\"\"\"\n    imgs = next(iter(val_dataloader))\n    #print(imgs)\n    X1 = Variable(imgs[\"monet_jpg\"].type(Tensor))\n    X2 = Variable(imgs[\"photo_jpg\"].type(Tensor))\n    _, Z1 = E1(X1)\n    _, Z2 = E2(X2)\n    fake_X1 = G1(Z2)\n    fake_X2 = G2(Z1)\n    \n    #Generate grids\n    X1 = Tensor.cpu(X2) \n    fake_X1 = Tensor.cpu(fake_X1) \n    #print(X1.shape)\n    grid_x =  reverse_normalize(make_grid(X1, nrow=4).permute(1, 2, 0).detach().cpu().numpy())\n    grid_fake_y =  reverse_normalize(make_grid(fake_X1, nrow=4).permute(1, 2, 0).detach().cpu().numpy() )\n    \n    #Transformation from X -> Y\n\n    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(30, 20))\n    \n    #ax1.imshow(X1.permute(1, 2, 0))\n    ax1.imshow(grid_x)\n    \n    ax1.axis('off')\n    ax1.set_title('X')\n    ax2.imshow(grid_fake_y)\n    ax2.axis('off')\n    ax2.set_title('Fake Y, Monet-esque')\n    plt.show()\n    \n    \n    \n    \n    \ndef compute_kl(mu):\n    mu_2 = torch.pow(mu, 2)\n    loss = torch.mean(mu_2)\n    return loss\n\n\n\n\nL = [] # Losses list\nfor rate in LR:\n    criterion_GAN = torch.nn.MSELoss()\n    criterion_pixel = torch.nn.L1Loss()\n\n    input_shape = (3, 256, 256)\n\n# Dimensionality (channel-wise) of image embedding\n    shared_dim = dim * 2 ** n_downsample\n\n# Initialize generator and discriminator\n    shared_E = ResidualBlock(features=shared_dim)\n    E1 = Encoder(dim=dim, n_downsample=n_downsample, shared_block=shared_E)\n    E2 = Encoder(dim=dim, n_downsample=n_downsample, shared_block=shared_E)\n    shared_G = ResidualBlock(features=shared_dim)\n    G1 = Generator(dim=dim, n_upsample=n_downsample, shared_block=shared_G)\n    G2 = Generator(dim=dim, n_upsample=n_downsample, shared_block=shared_G)\n    D1 = Discriminator(input_shape)\n    D2 = Discriminator(input_shape)\n\n    if cuda:\n        E1 = E1.cuda()\n        E2 = E2.cuda()\n        G1 = G1.cuda()\n        G2 = G2.cuda()\n        D1 = D1.cuda()\n        D2 = D2.cuda()\n        criterion_GAN.cuda()\n        criterion_pixel.cuda()\n\n\n    E1.apply(weights_init_normal)\n    E2.apply(weights_init_normal)\n    G1.apply(weights_init_normal)\n    G2.apply(weights_init_normal)\n    D1.apply(weights_init_normal)\n    D2.apply(weights_init_normal)\n\n# Loss weights\n    lambda_0 = 10  # GAN\n    lambda_1 = 0.1  # KL (encoded images)\n    lambda_2 = 100  # ID pixel-wise\n    lambda_3 = 0.1  # KL (encoded translated images)\n    lambda_4 = 100  # Cycle pixel-wise\n\n# Optimizers\n    optimizer_G = torch.optim.Adam(\n        itertools.chain(E1.parameters(), E2.parameters(), G1.parameters(), G2.parameters()),\n        lr=rate,\n        betas=(b1, b2),\n    )\n    optimizer_D1 = torch.optim.Adam(D1.parameters(), lr=rate, betas=(b1, b2))\n    optimizer_D2 = torch.optim.Adam(D2.parameters(), lr=rate, betas=(b1, b2))\n\n# Learning rate update schedulers\n    lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n        optimizer_G, lr_lambda=LambdaLR(n_epochs, 0, 0).step\n    )\n    lr_scheduler_D1 = torch.optim.lr_scheduler.LambdaLR(\n        optimizer_D1, lr_lambda=LambdaLR(n_epochs, 0, 0).step\n    )\n    lr_scheduler_D2 = torch.optim.lr_scheduler.LambdaLR(\n        optimizer_D2, lr_lambda=LambdaLR(n_epochs, 0, 0).step\n    )\n\n    Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\n# Image transformations\n    transforms_ = [\n        transforms.Resize(int(256 * 1.12), Image.BICUBIC),\n        transforms.RandomCrop((256, 256)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]\n\n\n\n# ----------\n#  Training\n# ----------\n    losses = []\n    prev_time = time.time()\n    for epoch in range(0, n_epochs):\n        torch.manual_seed(epoch)\n        \n        average_loss_G1 = 0\n        av_loss = 0\n        \n        # Training data loader\n        dataloader = DataLoader(\n            ImageDataset(\"%s\" % dataset_name, transforms_=transforms_, unaligned=True),\n            batch_size=batch_size,\n            shuffle=True,\n            num_workers=n_cpu,       #changer à 1 ????\n        )\n        # Test data loader\n        val_dataloader = DataLoader(\n            ImageDataset(\"%s\" % dataset_name, transforms_=transforms_, unaligned=True, mode=\"test\"),\n            batch_size=5,\n            shuffle=True,\n            num_workers=1,\n        )\n        for i, batch in enumerate(dataloader):\n\n        # Set model input\n            X1 = Variable(batch[\"monet_jpg\"].type(Tensor))\n            X2 = Variable(batch[\"photo_jpg\"].type(Tensor))\n\n        # Adversarial ground truths\n            valid = Variable(Tensor(np.ones((X1.size(0), *D1.output_shape))), requires_grad=False)\n            fake = Variable(Tensor(np.zeros((X1.size(0), *D1.output_shape))), requires_grad=False)\n\n        # -------------------------------\n        #  Train Encoders and Generators\n        # -------------------------------\n\n            optimizer_G.zero_grad()\n\n        # Get shared latent representation\n            mu1, Z1 = E1(X1)\n            mu2, Z2 = E2(X2)\n\n        # Reconstruct images\n            recon_X1 = G1(Z1)\n            recon_X2 = G2(Z2)\n\n        # Translate images\n            fake_X1 = G1(Z2)\n            fake_X2 = G2(Z1)\n\n        # Cycle translation\n            mu1_, Z1_ = E1(fake_X1)\n            mu2_, Z2_ = E2(fake_X2)\n            cycle_X1 = G1(Z2_)\n            cycle_X2 = G2(Z1_)\n\n        # Losses\n            loss_GAN_1 = lambda_0 * criterion_GAN(D1(fake_X1), valid)\n            loss_GAN_2 = lambda_0 * criterion_GAN(D2(fake_X2), valid)\n            loss_KL_1 = lambda_1 * compute_kl(mu1)\n            loss_KL_2 = lambda_1 * compute_kl(mu2)\n            loss_ID_1 = lambda_2 * criterion_pixel(recon_X1, X1)\n            loss_ID_2 = lambda_2 * criterion_pixel(recon_X2, X2)\n            loss_KL_1_ = lambda_3 * compute_kl(mu1_)\n            loss_KL_2_ = lambda_3 * compute_kl(mu2_)\n            loss_cyc_1 = lambda_4 * criterion_pixel(cycle_X1, X1)\n            loss_cyc_2 = lambda_4 * criterion_pixel(cycle_X2, X2)\n\n        # Total loss\n            loss_G = (\n                loss_KL_1\n                + loss_KL_2\n                + loss_ID_1\n                + loss_ID_2\n                + loss_GAN_1\n                + loss_GAN_2\n                + loss_KL_1_\n                + loss_KL_2_\n                + loss_cyc_1\n                + loss_cyc_2\n            )\n\n            loss_G.backward()\n            optimizer_G.step()\n        \n            av_loss += loss_GAN_1 /len(dataloader)\n            average_loss_G1 += loss_G /(len(dataloader))\n        \n        # -----------------------\n        #  Train Discriminator 1\n        # -----------------------\n\n            optimizer_D1.zero_grad()\n\n            loss_D1 = criterion_GAN(D1(X1), valid) + criterion_GAN(D1(fake_X1.detach()), fake)\n\n            loss_D1.backward()\n            optimizer_D1.step()\n\n        # -----------------------\n        #  Train Discriminator 2\n        # -----------------------\n\n            optimizer_D2.zero_grad()\n\n            loss_D2 = criterion_GAN(D2(X2), valid) + criterion_GAN(D2(fake_X2.detach()), fake)\n\n            loss_D2.backward()\n            optimizer_D2.step()\n\n        # --------------\n        #  Log Progress\n        # --------------\n####\n        # Determine approximate time left\n            batches_done = epoch * len(dataloader) + i\n            batches_left = n_epochs * len(dataloader) - batches_done\n            time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n            prev_time = time.time()\n\n        # Print log\n            sys.stdout.write(\n                \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f] ETA: %s\"\n                % (epoch, n_epochs, i, len(dataloader), (loss_D1 + loss_D2).item(), loss_G.item(), time_left)\n            )\n            if batches_done % sample_interval == 0:\n                sample_images(batches_done)\n\n\n        losses.append(average_loss_G1.item())\n        print(\"\\n averageloss: \",average_loss_G1.item())\n        print(\"gan1_loss\", av_loss.item())\n        # Update learning rates\n        lr_scheduler_G.step()\n        lr_scheduler_D1.step()\n        lr_scheduler_D2.step()\n    L.append(losses)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\n! rm -rf ../images\n! mkdir ../images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\nfrom PIL import Image\nimport imageio\nfrom tqdm import tqdm\n\n\nclass ImageDataset(Dataset):\n        \"\"\"\n        Custom dataset\n        \"\"\"\n        \n        def __init__(self, img_path, img_size=256, normalize=True):\n            self.img_path = img_path\n            \n            if normalize:\n                self.transform = transforms.Compose([\n                    transforms.Resize(img_size),\n                    transforms.ToTensor(),\n                    transforms.Normalize(mean=[0.5], std=[0.5])\n                ])\n            else:\n                self.transform = transforms.Compose([\n                    transforms.Resize(img_size),\n                    transforms.ToTensor()\n                ])\n            \n            #Dictionary entries\n            self.img_idx = dict()\n            for number_, img_ in enumerate(os.listdir(self.img_path)):\n                self.img_idx[number_] = img_\n                \n        def __len__(self):\n            #Length of dataset --> number of images\n            return len(self.img_idx)\n        \n        def __getitem__(self, idx):\n            img_path = os.path.join(self.img_path, self.img_idx[idx])\n            img = Image.open(img_path)\n            img = self.transform(img)\n            \n            return img\n\n        \nE1.eval()        \nG1.eval()\n\npath_photo = \"../input/gan-getting-started/photo_jpg\"\n#Get data loader for final transformation / submission\ndataset_photo = ImageDataset(path_photo, img_size=256, normalize=True)\nsubmit_dataloader = DataLoader(dataset_photo, batch_size=1, shuffle=False)\n\n\ndataiter = iter(submit_dataloader)\n\n#Previous normalization choosen\nmean_=0.5 \nstd_=0.5\n\n#Loop through each picture\nfor image_idx in tqdm(range(0, len(submit_dataloader))):\n        \n    #Get base picture\n    fixed_X = dataiter.next()\n    \n    #Identify correct device\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    \n    #Create fake pictures (monet-esque)\n    #X1 = Variable(imgs[\"monet_jpg\"].type(Tensor))\n    _,encod_fake = E2(fixed_X.to(device))\n    fake_Y = G1(encod_fake)\n    fake_Y = fake_Y.detach().cpu().numpy()\n    fake_Y = reverse_normalize(fake_Y, mean_, std_)\n    fake_Y = fake_Y[0].transpose(1, 2, 0)\n    fake_Y = np.uint8(fake_Y)\n    fake_Y = Image.fromarray(fake_Y)\n    #print(fake_Y.shape)\n    \n    #Save picture\n    fake_Y.save(\"../images/\" + str(image_idx) + \".jpg\")\n\n#Back to it\n        \nE1.train()        \nG1.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}