{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install image-quality\n#https://colab.research.google.com/github/supertramp2/Colab/blob/main/CycleGAN.ipynb#scrollTo=soWFSTxHGoAU\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.autograd import Variable\nimport torch.utils.data as data\nimport random\nfrom torchvision import transforms\nimport torchvision.models.vgg as vgg\nimport torch.utils.model_zoo as model_zoo\nfrom collections import namedtuple\nimport torch\nimport PIL\nfrom PIL import Image\nimport os , itertools\nimport shutil\nfrom pathlib import Path\nfrom collections import OrderedDict\nimport matplotlib.pyplot as plt\nimport time\nimport imquality.brisque as brisque\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:27:38.107528Z","iopub.execute_input":"2021-08-12T12:27:38.107882Z","iopub.status.idle":"2021-08-12T12:27:44.299128Z","shell.execute_reply.started":"2021-08-12T12:27:38.107849Z","shell.execute_reply":"2021-08-12T12:27:44.298167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#model params\nparams = {\n    'batch_size':1,\n    'input_size':256,\n    'resize_scale':286, \n    'crop_size':256,\n    'fliplr':True,\n    'num_epochs':50,\n    'decay_epoch':50,\n    'ngf':32,   #number of generator filters\n    'ndf':64,   #number of discriminator filters\n    'num_resnet':6, #number of resnet blocks\n    'lrG':0.0002,    #learning rate for generator\n    'lrD':0.0002,    #learning rate for discriminator\n    'beta1':0.5 ,    #beta1 for Adam optimizer\n    'beta2':0.999 ,  #beta2 for Adam optimizer\n    'lambdaA':10 ,   #lambdaA for cycle loss\n    'lambdaB':10  ,  #lambdaB for cycle loss\n}\n\nRANDOM = \"monet_jpg_random\"\nOPPOSITE = \"monet_jpg_opposite\"\nOBJECTS = \"monet-jpg-objects\"\norig_data = \"../input/gan-getting-started\"\nchosen_subfolder = \"monet_jpg_quality\"\nmonet_list_dest = './' + chosen_subfolder","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:20:15.883486Z","iopub.execute_input":"2021-08-12T12:20:15.883832Z","iopub.status.idle":"2021-08-12T12:20:15.889757Z","shell.execute_reply.started":"2021-08-12T12:20:15.883786Z","shell.execute_reply":"2021-08-12T12:20:15.888955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_dest(dest):\n  if os.path.exists(dest) and os.path.isdir(dest):\n    if os.listdir(dest):\n        for f in os.listdir(dest):\n          os.remove(os.path.join(dest, f))\n  else:\n      os.makedirs(dest)","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:20:15.891499Z","iopub.execute_input":"2021-08-12T12:20:15.892119Z","iopub.status.idle":"2021-08-12T12:20:15.901306Z","shell.execute_reply.started":"2021-08-12T12:20:15.892079Z","shell.execute_reply":"2021-08-12T12:20:15.900491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_subfolder_30_random = \"/monet_jpg_random/\"\nmonet_subfolder_30_opposites = \"/monet_jpg_opposites/\"\nmonet_list = orig_data + \"/monet_jpg/\"\n\n\nbest_quality = []\ndef train_data_30_quality():\n  global best_quality\n  monet_list_dir = os.listdir(monet_list)\n  for i,filename in enumerate(tqdm(os.listdir(monet_list))):\n    img = Image.open(monet_list + filename)\n    best_quality.append((filename, brisque.score(img)))\n    \n  create_dest(monet_list_dest)\n  best_quality = sorted(best_quality, key=lambda img: img[1], reverse=True)[:30]\n  for img,score in best_quality:\n      if not Path(monet_list_dest + img).is_file():\n        shutil.copy(monet_list + img, monet_list_dest + '/' + img)\n\ndef train_data_30_random():\n  monet_list_dir = os.listdir(monet_list)\n\n  create_dest(monet_list_dest)\n  count = 0\n  while count < 30:\n    m = random.choice(monet_list_dir)\n    if not Path(monet_list_dest + '/' + m).is_file():\n      shutil.copy(monet_list + m, monet_list_dest + '/' + m)\n      count += 1\n\ndef get_dominant_color(pil_img):\n    img = pil_img.copy()\n    img.convert(\"RGB\")\n    img.resize((1, 1), resample=0)\n    dominant_color = '%02x%02x%02x' % img.getpixel((0, 0))\n    return dominant_color \n\ndef train_data_30_opposites():\n  monet_dom_color = {}\n  \n  create_dest(monet_list_dest)\n  for filename in os.listdir(monet_list):\n    img = Image.open(monet_list + filename)\n    img_dom_color = get_dominant_color(img)\n    if img_dom_color in monet_dom_color:\n      monet_dom_color[img_dom_color].append(filename)\n    else:\n      monet_dom_color[img_dom_color] = [filename]\n  \n  sorted_monet_dom_color = OrderedDict(sorted(monet_dom_color.items()))\n\n  if len(sorted_monet_dom_color) <= 30:\n    count = 0\n    while count < 30:\n      i = 0\n      for color, ms in sorted_monet_dom_color.items():\n        shutil.copy(monet_list + ms[i], monet_list_dest + \"/\" + ms[i])\n        count += 1\n      i += 1\n  else:\n    jump = len(sorted_monet_dom_color) / 30\n    count = 0\n    getMonet = 0\n    index = 0\n    for color, monet in sorted_monet_dom_color.items():\n      if count < 15:\n        if getMonet == index:\n          shutil.copy(monet_list + monet[0], monet_list_dest + \"/\" + monet[0])\n          count += 1\n          getMonet += jump\n        index += 1\n    \n    desc_sorted_monet_dom_color = OrderedDict(sorted(sorted_monet_dom_color.items(), reverse=True))\n    getMonet = 0\n    index = 0\n    for color, monet in desc_sorted_monet_dom_color.items():\n      if count < 30:\n        if getMonet == index:\n          if not Path(monet_list_dest + \"/\" + monet[0]).is_file():\n            shutil.copy(monet_list + monet[0], monet_list_dest + \"/\" + monet[0])\n            count += 1\n            getMonet += jump\n        index += 1\n\nif \"random\" in chosen_subfolder:\n  train_data_30_random()\n  print(1)\nelif \"opposites\" in chosen_subfolder:\n  train_data_30_opposites()\n  print(2)\nelse:\n  print(3)\n  train_data_30_quality()","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:20:15.903049Z","iopub.execute_input":"2021-08-12T12:20:15.903458Z","iopub.status.idle":"2021-08-12T12:23:11.405697Z","shell.execute_reply.started":"2021-08-12T12:20:15.903423Z","shell.execute_reply":"2021-08-12T12:23:11.404042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def to_np(x):\n    return x.data.cpu().numpy()\ndef plot_train_result(real_image, gen_image, recon_image, epoch, save=False,  show=True, fig_size=(15, 15)):\n    fig, axes = plt.subplots(2, 3, figsize=fig_size)\n    imgs = [to_np(real_image[0]), to_np(gen_image[0]), to_np(recon_image[0]),\n            to_np(real_image[1]), to_np(gen_image[1]), to_np(recon_image[1])]\n    for ax, img in zip(axes.flatten(), imgs):\n        ax.axis('off')\n        #ax.set_adjustable('box-forced')\n        # Scale to 0-255\n        img = img.squeeze()\n        img = (((img - img.min()) * 255) / (img.max() - img.min())).transpose(1, 2, 0).astype(np.uint8)\n        ax.imshow(img, cmap=None, aspect='equal')\n    plt.subplots_adjust(wspace=0, hspace=0)\n\n    title = 'Epoch {0}'.format(epoch + 1)\n    fig.text(0.5, 0.04, title, ha='center')\n\n    # save figure\n    if save:\n        save_fn = 'Result_epoch_{:d}'.format(epoch+1) + '.png'\n        plt.savefig(save_fn)\n\n    if show:\n        plt.show()\n    else:\n        plt.close()","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:23:11.407371Z","iopub.execute_input":"2021-08-12T12:23:11.407838Z","iopub.status.idle":"2021-08-12T12:23:11.418725Z","shell.execute_reply.started":"2021-08-12T12:23:11.407792Z","shell.execute_reply":"2021-08-12T12:23:11.417808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ImagePool():\n    def __init__(self, pool_size):\n        self.pool_size = pool_size\n        if self.pool_size > 0:\n            self.num_imgs = 0\n            self.images = []\n\n    def query(self, images):\n        if self.pool_size == 0:\n            return images\n        return_images = []\n        for image in images.data:\n            image = torch.unsqueeze(image, 0)\n            if self.num_imgs < self.pool_size:\n                self.num_imgs = self.num_imgs + 1\n                self.images.append(image)\n                return_images.append(image)\n            else:\n                p = random.uniform(0, 1)\n                if p > 0.5:\n                    random_id = random.randint(0, self.pool_size-1)\n                    tmp = self.images[random_id].clone()\n                    self.images[random_id] = image\n                    return_images.append(tmp)\n                else:\n                    return_images.append(image)\n        return_images = Variable(torch.cat(return_images, 0))\n        return return_images\n        \nclass DatasetFromFolder(data.Dataset):\n    def __init__(self, image_dir, subfolder, transform=None, resize_scale=None, crop_size=None, fliplr=False):\n        super(DatasetFromFolder, self).__init__()\n        self.input_path = os.path.join(image_dir, subfolder)\n        self.image_filenames = [x for x in sorted(os.listdir(self.input_path))]\n        self.transform = transform\n        \n        self.resize_scale = resize_scale\n        self.crop_size = crop_size\n        self.fliplr = fliplr\n\n    def __getitem__(self, index):\n        # Load Image\n        img_fn = os.path.join(self.input_path, self.image_filenames[index])\n        img = Image.open(img_fn).convert('RGB')\n\n        # preprocessing\n        if self.resize_scale:\n            img = img.resize((self.resize_scale, self.resize_scale), Image.BILINEAR)\n\n        if self.crop_size:\n            x = random.randint(0, self.resize_scale - self.crop_size + 1)\n            y = random.randint(0, self.resize_scale - self.crop_size + 1)\n            img = img.crop((x, y, x + self.crop_size, y + self.crop_size))\n        if self.fliplr:\n            if random.random() < 0.5:\n                img = img.transpose(Image.FLIP_LEFT_RIGHT)\n\n        if self.transform is not None:\n            img = self.transform(img)\n\n        return img\n\n    def __len__(self):\n        return len(self.image_filenames)","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:23:11.420279Z","iopub.execute_input":"2021-08-12T12:23:11.420709Z","iopub.status.idle":"2021-08-12T12:23:11.436557Z","shell.execute_reply.started":"2021-08-12T12:23:11.420621Z","shell.execute_reply":"2021-08-12T12:23:11.435612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CycleGAN Architecture","metadata":{}},{"cell_type":"code","source":"#https://github.com/aitorzip/PyTorch-CycleGAN/blob/master/models.py\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_features):\n        super(ResidualBlock, self).__init__()\n\n        conv_block = [  nn.ReflectionPad2d(1),\n                        nn.Conv2d(in_features, in_features, 3),\n                        nn.InstanceNorm2d(in_features),\n                        nn.ReLU(inplace=True),\n                        nn.ReflectionPad2d(1),\n                        nn.Conv2d(in_features, in_features, 3),\n                        nn.InstanceNorm2d(in_features)  ]\n\n        self.conv_block = nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        return x + self.conv_block(x)\n\nclass Generator(nn.Module):\n    def __init__(self, input_nc, output_nc, n_residual_blocks=9):\n        super(Generator, self).__init__()\n\n        # Initial convolution block       \n        model = [   nn.ReflectionPad2d(3),\n                    nn.Conv2d(input_nc, 64, 7),\n                    nn.InstanceNorm2d(64),\n                    nn.ReLU(inplace=True) ]\n\n        # Downsampling\n        in_features = 64\n        out_features = in_features*2\n        for _ in range(2):\n            model += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n                        nn.InstanceNorm2d(out_features),\n                        nn.ReLU(inplace=True) ]\n            in_features = out_features\n            out_features = in_features*2\n\n        # Residual blocks\n        for _ in range(n_residual_blocks):\n            model += [ResidualBlock(in_features)]\n\n        # Upsampling\n        out_features = in_features//2\n        for _ in range(2):\n            model += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),\n                        nn.InstanceNorm2d(out_features),\n                        nn.ReLU(inplace=True) ]\n            in_features = out_features\n            out_features = in_features//2\n\n        # Output layer\n        model += [  nn.ReflectionPad2d(3),\n                    nn.Conv2d(64, output_nc, 7),\n                    nn.Tanh() ]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        return self.model(x)\n\nclass Discriminator(nn.Module):\n    def __init__(self, input_nc):\n        super(Discriminator, self).__init__()\n\n        # A bunch of convolutions one after another\n        model = [   nn.Conv2d(input_nc, 64, 4, stride=2, padding=1),\n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(64, 128, 4, stride=2, padding=1),\n                    nn.InstanceNorm2d(128), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(128, 256, 4, stride=2, padding=1),\n                    nn.InstanceNorm2d(256), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(256, 512, 4, padding=1),\n                    nn.InstanceNorm2d(512), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        # FCN classification layer\n        model += [nn.Conv2d(512, 1, 4, padding=1)]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        x =  self.model(x)\n        # Average pooling and flatten\n        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)","metadata":{"execution":{"iopub.status.busy":"2021-08-12T12:23:11.437792Z","iopub.execute_input":"2021-08-12T12:23:11.438283Z","iopub.status.idle":"2021-08-12T12:23:11.458164Z","shell.execute_reply.started":"2021-08-12T12:23:11.438247Z","shell.execute_reply":"2021-08-12T12:23:11.45733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Dataset","metadata":{}},{"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize(size=params['input_size']),\n    transforms.RandomCrop(224),\n    transforms.ColorJitter(0.5),\n    transforms.RandomRotation(degrees=45),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomVerticalFlip(p=0.5),\n    transforms.RandomGrayscale(p=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) #TODO check if these are the actual values, and if not- change to actual values\n])\n#Subfolders\ntrain_data_A = DatasetFromFolder(orig_data, subfolder='photo_jpg', transform=transform,\n                                resize_scale=params['resize_scale'], crop_size=params['crop_size'], fliplr=params['fliplr'])\ntrain_data_loader_A = torch.utils.data.DataLoader(dataset=train_data_A, batch_size=params['batch_size'], shuffle=True)\n\ntrain_data_B = DatasetFromFolder(image_dir='./', subfolder=chosen_subfolder, transform=transform,\n                                resize_scale=params['resize_scale'], crop_size=params['crop_size'], fliplr=params['fliplr'])\ntrain_data_loader_B = torch.utils.data.DataLoader(dataset=train_data_B, batch_size=params['batch_size'], shuffle=True)\n\n#Kaggle GPU\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:10.62695Z","iopub.execute_input":"2021-08-11T17:54:10.62731Z","iopub.status.idle":"2021-08-11T17:54:10.825002Z","shell.execute_reply.started":"2021-08-11T17:54:10.627277Z","shell.execute_reply":"2021-08-11T17:54:10.824122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#After each epoch output of these input images (tensors) will be displayed\ntest_real_A_data = train_data_A.__getitem__(11).unsqueeze(0) \ntest_real_B_data = train_data_B.__getitem__(11).unsqueeze(0)","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:10.82672Z","iopub.execute_input":"2021-08-11T17:54:10.827055Z","iopub.status.idle":"2021-08-11T17:54:10.904248Z","shell.execute_reply.started":"2021-08-11T17:54:10.82702Z","shell.execute_reply":"2021-08-11T17:54:10.903374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Build The Model","metadata":{}},{"cell_type":"code","source":"#https://gist.github.com/jeasinema/ed9236ce743c8efaf30fa2ff732749f5\n\ndef weights_init_normal(m):\n    if isinstance(m, nn.Conv1d):\n        init.normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.Conv2d):\n        init.xavier_normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.Conv3d):\n        init.xavier_normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.ConvTranspose1d):\n        init.normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.ConvTranspose2d):\n        init.xavier_normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.ConvTranspose3d):\n        init.xavier_normal_(m.weight.data)\n        if m.bias is not None:\n            init.normal_(m.bias.data)\n    elif isinstance(m, nn.BatchNorm1d):\n        init.normal_(m.weight.data, mean=1, std=0.02)\n        init.constant_(m.bias.data, 0)\n    elif isinstance(m, nn.BatchNorm2d):\n        init.normal_(m.weight.data, mean=1, std=0.02)\n        init.constant_(m.bias.data, 0)\n    elif isinstance(m, nn.BatchNorm3d):\n        init.normal_(m.weight.data, mean=1, std=0.02)\n        init.constant_(m.bias.data, 0)\n    elif isinstance(m, nn.Linear):\n        init.xavier_normal_(m.weight.data)\n        init.normal_(m.bias.data)\n    elif isinstance(m, nn.LSTM):\n        for param in m.parameters():\n            if len(param.shape) >= 2:\n                init.orthogonal_(param.data)\n            else:\n                init.normal_(param.data)\n    elif isinstance(m, nn.LSTMCell):\n        for param in m.parameters():\n            if len(param.shape) >= 2:\n                init.orthogonal_(param.data)\n            else:\n                init.normal_(param.data)\n    elif isinstance(m, nn.GRU):\n        for param in m.parameters():\n            if len(param.shape) >= 2:\n                init.orthogonal_(param.data)\n            else:\n                init.normal_(param.data)\n    elif isinstance(m, nn.GRUCell):\n        for param in m.parameters():\n            if len(param.shape) >= 2:\n                init.orthogonal_(param.data)\n            else:\n                init.normal_(param.data)\n","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:10.905653Z","iopub.execute_input":"2021-08-11T17:54:10.906058Z","iopub.status.idle":"2021-08-11T17:54:10.92627Z","shell.execute_reply.started":"2021-08-11T17:54:10.906019Z","shell.execute_reply":"2021-08-11T17:54:10.925421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Build Model \n#G_A - real photo->monet style ; G_B - monet style -> real photo\nG_A = Generator(3,3).cuda() \nG_B = Generator(3,3).cuda()\n\n#two Discriminators\nD_A = Discriminator(3).cuda()\nD_B = Discriminator(3).cuda()\n\nG_A.apply(weights_init_normal)\nG_B.apply(weights_init_normal)\nD_A.apply(weights_init_normal)\nD_B.apply(weights_init_normal)\n\n\nG_optimizer = torch.optim.Adam(itertools.chain(G_A.parameters(), G_B.parameters()), lr=params['lrG'], betas=(params['beta1'], params['beta2']))\nD_A_optimizer = torch.optim.Adam(D_A.parameters(), lr=params['lrD'], betas=(params['beta1'], params['beta2']))\nD_B_optimizer = torch.optim.Adam(D_B.parameters(), lr=params['lrD'], betas=(params['beta1'], params['beta2']))","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:10.928322Z","iopub.execute_input":"2021-08-11T17:54:10.928671Z","iopub.status.idle":"2021-08-11T17:54:15.884771Z","shell.execute_reply.started":"2021-08-11T17:54:10.928647Z","shell.execute_reply":"2021-08-11T17:54:15.883881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Functions","metadata":{}},{"cell_type":"code","source":"MSE_Loss = torch.nn.MSELoss().cuda()\nL1_Loss = torch.nn.L1Loss().cuda()\nLossOutput = namedtuple(\"LossOutput\", [\"relu1_2\", \"relu2_2\", \"relu3_3\", \"relu4_3\"])\n\n# https://discuss.pytorch.org/t/how-to-extract-features-of-an-image-from-a-trained-model/119/3\nclass LossNetwork(torch.nn.Module):\n    def __init__(self, vgg_model):\n        super(LossNetwork, self).__init__()\n        self.vgg_layers = vgg_model.features\n        self.layer_name_mapping = {\n            '3': \"relu1_2\",\n            '8': \"relu2_2\",\n            '15': \"relu3_3\",\n            '22': \"relu4_3\"\n        }\n    \n    def forward(self, x):\n        output = {}\n        for name, module in self.vgg_layers._modules.items():\n            x = module(x)\n            if name in self.layer_name_mapping:\n                output[self.layer_name_mapping[name]] = x\n        return LossOutput(**output)\n\nvgg_model = vgg.vgg16(pretrained=True)\nif torch.cuda.is_available():\n    vgg_model.cuda()\nloss_network = LossNetwork(vgg_model)\nloss_network.eval()\ndel vgg_model\n\ndef gram_matrix(y):\n    (b, ch, h, w) = y.size()\n    features = y.view(b, ch, w * h)\n    features_t = features.transpose(1, 2)\n    gram = features.bmm(features_t) / (ch * h * w)\n    return gram\n\ndef compStyle(a,b):\n    #http://pytorch.org/docs/master/notes/autograd.html#volatile\n    styleB_loss_features = loss_network(Variable(a, volatile=True))\n    gram_style = [Variable(gram_matrix(y).data, requires_grad=False) for y in styleB_loss_features]\n        \n    features_y = loss_network(b)\n        \n    style_loss = 0    \n    for m in range(len(features_y)):\n        gram_s = gram_style[m]\n        gram_y = gram_matrix(features_y[m])\n        style_loss += 1e4 * MSE_Loss(gram_y, gram_s.expand_as(gram_y))\n    return style_loss","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:15.88618Z","iopub.execute_input":"2021-08-11T17:54:15.886566Z","iopub.status.idle":"2021-08-11T17:54:40.421781Z","shell.execute_reply.started":"2021-08-11T17:54:15.886516Z","shell.execute_reply":"2021-08-11T17:54:40.419194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Model","metadata":{}},{"cell_type":"code","source":"D_A_avg_losses = []\nD_B_avg_losses = []\nG_A_avg_losses = []\nG_B_avg_losses = []\ncycle_A_avg_losses = []\ncycle_B_avg_losses = []\nSTYLE_WEIGHT = 1e4\n\nnum_pool = 10\nfake_A_pool = ImagePool(num_pool)\nfake_B_pool = ImagePool(num_pool)\n\nstep = 0\nfor epoch in range(params['num_epochs']):\n    D_A_losses = []\n    D_B_losses = []\n    G_A_losses = []\n    G_B_losses = []\n    cycle_A_losses = []\n    cycle_B_losses = []\n    \n    # Learing rate decay \n    if(epoch + 1) > params['decay_epoch']:\n        D_A_optimizer.param_groups[0]['lr'] -= params['lrD'] / (params['num_epochs'] - params['decay_epoch'])\n        D_B_optimizer.param_groups[0]['lr'] -= params['lrD'] / (params['num_epochs'] - params['decay_epoch'])\n        G_optimizer.param_groups[0]['lr'] -= params['lrG'] / (params['num_epochs'] - params['decay_epoch'])\n        \n\n        \n    # training \n    for i, (real_A, real_B) in enumerate(zip(train_data_loader_A, train_data_loader_B)):\n        \n        # input image data\n        real_A = real_A.to(device)\n        real_B = real_B.to(device)\n        \n        # -------------------------- train generator G_A --------------------------\n        # A --> B\n        fake_B = G_A(real_A)\n        a_idt = G_A(real_A)\n        \n        D_B_fake_decision = D_B(fake_B)\n        G_A_loss = MSE_Loss(D_B_fake_decision, Variable(torch.ones(D_B_fake_decision.size()).cuda()))\n        \n        # forward cycle loss\n        recon_A = G_B(fake_B)\n        cycle_A_loss = L1_Loss(recon_A, real_A) * params['lambdaA']\n        \n        #idtA_loss = L1_Loss(a_idt,real_A) * 10*0.5 \n        \n        styleA_loss = compStyle(real_A,a_idt) \n        \n    \n        #G_B_loss = G_B_loss + (style_loss)/2\n       \n        #ends here\n        \n        # -------------------------- train generator G_B --------------------------\n\n        # B --> A\n        fake_A = G_B(real_B)\n        b_idt = G_B(real_B)\n        \n        D_A_fake_decision = D_A(fake_A)\n        G_B_loss = MSE_Loss(D_A_fake_decision, Variable(torch.ones(D_A_fake_decision.size()).cuda()))\n        \n        # backward cycle loss\n        recon_B = G_A(fake_A)\n        cycle_B_loss = L1_Loss(recon_B, real_B) * params['lambdaB']\n        \n        #idtB_loss = L1_Loss(b_idt,real_B) * 10*0.5 \n    \n        styleB_loss = compStyle(real_B,b_idt) \n\n        style_loss = (styleB_loss + styleA_loss)\n        \n        # Back propagation\n        G_loss = G_A_loss + G_B_loss + cycle_A_loss + cycle_B_loss \n        \n        G_loss = G_loss+style_loss * 2.5\n        \n        \n        G_optimizer.zero_grad()\n        G_loss.backward()\n        G_optimizer.step()\n    \n        \n        # -------------------------- train discriminator D_A --------------------------\n        D_A_real_decision = D_A(real_A)\n        D_A_real_loss = MSE_Loss(D_A_real_decision, Variable(torch.ones(D_A_real_decision.size()).cuda()))\n        \n        fake_A = fake_A_pool.query(fake_A)\n        \n        D_A_fake_decision = D_A(fake_A)\n        D_A_fake_loss = MSE_Loss(D_A_fake_decision, Variable(torch.zeros(D_A_fake_decision.size()).cuda()))\n        \n       # D_A_recon_decision = D_A(recon_A)\n        #D_A_recon_loss = MSE_Loss(D_A_recon_decision, Variable(torch.zeros(D_A_recon_decision.size()).cuda()))\n        \n        # Back propagation\n        D_A_loss = (D_A_real_loss + D_A_fake_loss ) * 0.5\n        D_A_optimizer.zero_grad()\n        D_A_loss.backward()\n        D_A_optimizer.step()\n        \n        \n        # -------------------------- train discriminator D_B --------------------------\n        D_B_real_decision = D_B(real_B)\n        D_B_real_loss = MSE_Loss(D_B_real_decision, Variable(torch.ones(D_B_fake_decision.size()).cuda()))\n        \n        fake_B = fake_B_pool.query(fake_B)\n        \n        D_B_fake_decision = D_B(fake_B)\n        D_B_fake_loss = MSE_Loss(D_B_fake_decision, Variable(torch.zeros(D_B_fake_decision.size()).cuda()))\n        \n        #D_B_recon_decision = D_B(recon_B)\n        #D_B_recon_loss = MSE_Loss(D_B_recon_decision, Variable(torch.zeros(D_B_recon_decision.size()).cuda()))\n        \n        # Back propagation\n        D_B_loss = (D_B_real_loss + D_B_fake_loss ) * 0.5\n        D_B_optimizer.zero_grad()\n        D_B_loss.backward()\n        D_B_optimizer.step()\n        \n        # ------------------------ Print -----------------------------\n        # loss values\n        D_A_losses.append(D_A_loss.item())\n        D_B_losses.append(D_B_loss.item())\n        G_A_losses.append(G_A_loss.item())\n        G_B_losses.append(G_B_loss.item())\n        cycle_A_losses.append(cycle_A_loss.item())\n        cycle_B_losses.append(cycle_B_loss.item())\n\n        if i%100 == 0:\n            print('Epoch [%d/%d], Step [%d/%d], D_A_loss: %.4f, D_B_loss: %.4f, G_A_loss: %.4f, G_B_loss: %.4f'\n                  % (epoch+1, params['num_epochs'], i+1, len(train_data_loader_A), D_A_loss.item(), D_B_loss.item(), G_A_loss.item(), G_B_loss.item()))\n            \n        step += 1\n        \n    D_A_avg_loss = torch.mean(torch.FloatTensor(D_A_losses))\n    D_B_avg_loss = torch.mean(torch.FloatTensor(D_B_losses))\n    G_A_avg_loss = torch.mean(torch.FloatTensor(G_A_losses))\n    G_B_avg_loss = torch.mean(torch.FloatTensor(G_B_losses))\n    cycle_A_avg_loss = torch.mean(torch.FloatTensor(cycle_A_losses))\n    cycle_B_avg_loss = torch.mean(torch.FloatTensor(cycle_B_losses))\n\n    # avg loss values for plot\n    D_A_avg_losses.append(D_A_avg_loss.item())\n    D_B_avg_losses.append(D_B_avg_loss.item())\n    G_A_avg_losses.append(G_A_avg_loss.item())\n    G_B_avg_losses.append(G_B_avg_loss.item())\n    cycle_A_avg_losses.append(cycle_A_avg_loss.item())\n    cycle_B_avg_losses.append(cycle_B_avg_loss.item())\n    \n    # Show result for test image\n    test_real_A = test_real_A_data.cuda()\n    test_fake_B = G_A(test_real_A)\n    test_recon_A = G_B(test_fake_B)\n\n    test_real_B = test_real_B_data.cuda()\n    test_fake_A = G_B(test_real_B)\n    test_recon_B = G_A(test_fake_A)\n\n    plot_train_result([test_real_A, test_real_B], [test_fake_B, test_fake_A], [test_recon_A, test_recon_B],\n                            epoch, save=False)\n","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.423334Z","iopub.status.idle":"2021-08-11T17:54:40.423884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = [e for e in range(1,params['num_epochs']+1)]\n\nplt.figure(0)\nplt.plot(epochs, G_A_avg_losses, label='Gen Photo -> Monet')\nplt.plot(epochs, G_B_avg_losses, label='Gen Monet -> Photo')\nplt.plot(epochs, D_A_avg_losses, label='Disc Monet -> Photo')\nplt.plot(epochs, D_B_avg_losses, label='Disc Photo -> Monet')\n\ntitle = 'Average Loss - {0} Epochs, Monet Pictures by Quality, DA'.format(params['num_epochs'])\n\n# naming the x axis\nplt.xlabel('Epochs')\n# naming the y axis\nplt.ylabel('Average Loss')\n# giving a title to my graph\nplt.title(title)\n\n# show a legend on the plot\nplt.legend(ncol=2)\n\n# function to show the plot\nplt.show()\n\nplt.savefig('./'+title)\n","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.425339Z","iopub.status.idle":"2021-08-11T17:54:40.426075Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(1)\nplt.plot(epochs, cycle_A_avg_losses, label='Cycle A')\nplt.plot(epochs, cycle_B_avg_losses, label='Cycle B')\n\n# naming the x axis\nplt.xlabel('Epochs')\n# naming the y axis\nplt.ylabel('Average Loss')\n# giving a title to my graph\nplt.title(title)\n\n# show a legend on the plot\nplt.legend()\n\n# function to show the plot\nplt.show()\n\nplt.savefig('./'+title)\n","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.427452Z","iopub.status.idle":"2021-08-11T17:54:40.428355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Test Model","metadata":{}},{"cell_type":"code","source":"def reverse_normalize(image, mean_=0.5, std_=0.5):\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)","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.429949Z","iopub.status.idle":"2021-08-11T17:54:40.430729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transform_test = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n])\ntest_data_A = DatasetFromFolder('../input/gan-getting-started/', subfolder='photo_jpg', transform=transform_test)\ntest_data_loader_A = torch.utils.data.DataLoader(dataset=test_data_A, batch_size=params['batch_size'], shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.432294Z","iopub.status.idle":"2021-08-11T17:54:40.433096Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! mkdir ../images\n\nfor i, real_A in enumerate(test_data_loader_A):\n    real_A = real_A.to(device)\n    fake_B = G_A(real_A)\n\n    # Save picture\n    fake_B = fake_B.detach().cpu().numpy()\n    fake_B = reverse_normalize(fake_B, 0.5, 0.5)\n    fake_B = fake_B[0].transpose(1, 2, 0)\n    fake_B = np.uint8(fake_B)\n    fake_B = Image.fromarray(fake_B)\n    fake_B.save(\"../images/\" + str(i) + \".jpg\")\n    \nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2021-08-11T17:54:40.434531Z","iopub.status.idle":"2021-08-11T17:54:40.435227Z"},"trusted":true},"execution_count":null,"outputs":[]}]}