{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30559,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Monet paintings GAN transformer\n\n#### The idea for this project (taken from kaggle contests) was to utilize photos taken in 256x256 RGB format and turn them into *Monet-styled* paintings.\n#### We will be utilizing a WGAN with Gradient Penalty!","metadata":{"execution":{"iopub.status.busy":"2024-01-12T20:16:02.694311Z","iopub.execute_input":"2024-01-12T20:16:02.694633Z","iopub.status.idle":"2024-01-12T20:16:02.699249Z","shell.execute_reply.started":"2024-01-12T20:16:02.694608Z","shell.execute_reply":"2024-01-12T20:16:02.698256Z"}}},{"cell_type":"markdown","source":"#### A big part of the data preparation in this notebook was coppied from <a href=\"https://www.kaggle.com/code/enzou3/simple-gans-in-monets-paintings/notebook\"> this notebook </a>, so big shoutout to ENZOU3!","metadata":{}},{"cell_type":"markdown","source":"## Basic imports","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\nfrom tqdm.notebook import tqdm\nimport torch.optim as optim\nimport torch.nn as nn\nfrom PIL import Image\nimport torchvision\nimport shutil\nimport torch\nimport os\nimport random\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torchvision.utils import make_grid","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-01-23T15:25:31.132410Z","iopub.execute_input":"2024-01-23T15:25:31.133178Z","iopub.status.idle":"2024-01-23T15:25:32.951052Z","shell.execute_reply.started":"2024-01-23T15:25:31.133144Z","shell.execute_reply":"2024-01-23T15:25:32.950298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset constructor\n\n#### More specs on how this constructor works are included below but in a nutshell it will just take in the data in a folder and allow it to be used by the Dataloader in torch.utils","metadata":{}},{"cell_type":"code","source":"class PaintingsDataset(Dataset):\n    '''\n    Basic dataset class, will be used to manipulate the monet dataset \n    \n    __init__ args:          folder_path       :       path to folder with dataset images\n                            transform         :       transform module from torchvision for data augmentation\n                            __len__           :       number of images in dataset\n                            __get_item__      :       returns image at index idx tranformed\n                            __get_crude_item_ :       returns image at index idx \n                            \n                \n    '''\n    def __init__(self, folder_path, transform=None):\n        self.folder_path = folder_path\n        self.transform = transform\n        self.images = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(('.jpg', '.jpeg', '.png'))]\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        image_path = self.images[idx]\n        image = Image.open(image_path).convert('RGB')\n        if self.transform:\n            image = self.transform(image)\n        return image\n    \n    def __get_crude_item__(self, idx):\n        image_path = self.images[idx]\n        image = Image.open(image_path).convert('RGB')\n        return image\n        ","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:32.952920Z","iopub.execute_input":"2024-01-23T15:25:32.953857Z","iopub.status.idle":"2024-01-23T15:25:32.962023Z","shell.execute_reply.started":"2024-01-23T15:25:32.953820Z","shell.execute_reply":"2024-01-23T15:25:32.961135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Device adptation, hyperparameters used in training/creating both networks and data augmentation\n\n#### Basically, we will be using <code>torch.cuda.is_available()</code> to check if cuda is available and change model to cuda based on the result\n#### Additionally, the hyperparameters and filters in data augmentation (in order to reduce overfitting) were fine-tuned during training and also searched from other public kaggle notebooks for reference (more on that later, in References)","metadata":{}},{"cell_type":"code","source":"# Device and tensors configuration\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor\n\n# Hyperparameters\nbatch_size = 6\nlr = 2e-4\nn_epoches = 100\ndisplay_epoch = 10\n\n\"\"\"\nTranform module:\n\nonly applying GaussianBlur to dataset, more transformations seemed like a over-kill\n\n\"\"\"\ntransform = transforms.Compose([\n    transforms.Resize((256, 256)),\n\n    # Data augmentation \n    #transforms.ElasticTransform(alpha=150.0),\n    transforms.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 3.)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.15, 0.15, 0.15], std=[0.5, 0.5, 0.5])\n])","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:32.963114Z","iopub.execute_input":"2024-01-23T15:25:32.963409Z","iopub.status.idle":"2024-01-23T15:25:33.003008Z","shell.execute_reply.started":"2024-01-23T15:25:32.963385Z","shell.execute_reply":"2024-01-23T15:25:33.002210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initialize datasets\nmonet_dataset = PaintingsDataset('/kaggle/input/gan-getting-started/monet_jpg', transform)\nphoto_dataset = PaintingsDataset('/kaggle/input/gan-getting-started/photo_jpg', transform)\n\n\n# Initialize data loaders\nmonet_loader = DataLoader(monet_dataset, batch_size=batch_size, shuffle=True)\nphoto_loader = DataLoader(photo_dataset, batch_size=batch_size, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:33.004850Z","iopub.execute_input":"2024-01-23T15:25:33.005135Z","iopub.status.idle":"2024-01-23T15:25:33.535463Z","shell.execute_reply.started":"2024-01-23T15:25:33.005110Z","shell.execute_reply":"2024-01-23T15:25:33.534516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Make sure I read the data correctly:","metadata":{}},{"cell_type":"code","source":"monet_loader.dataset.__len__()","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:33.536980Z","iopub.execute_input":"2024-01-23T15:25:33.537286Z","iopub.status.idle":"2024-01-23T15:25:33.544947Z","shell.execute_reply.started":"2024-01-23T15:25:33.537246Z","shell.execute_reply":"2024-01-23T15:25:33.543937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Taking a look at the images:","metadata":{}},{"cell_type":"code","source":"\"\"\"\nShow a random image n from dataset with it's original format and transformed format\n\nargs:\n        n : example image index\n        dataset : dataset of interest\n\nreturns:\n        plot's with original and transformed image side-by-side\n\"\"\"\n\ndef show_images(n, dataset):\n    ex_im_transf = dataset.__getitem__(n)\n    ex_im_crude  = dataset.__get_crude_item__(n)\n\n    im_transf = np.squeeze(ex_im_transf) \n    im_crude = np.squeeze(ex_im_crude)\n\n    f, axarr = plt.subplots(1,2)\n\n    for i in range(2):\n        axarr[i].axes.get_xaxis().set_visible(False)\n        axarr[i].axes.get_yaxis().set_visible(False)\n\n    axarr[1].imshow(im_transf.permute(1, 2, 0))\n    axarr[0].imshow(im_crude)\n    axarr[1].set_title(\"Transformed image\")\n    axarr[0].set_title(\"Original image\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:33.700194Z","iopub.execute_input":"2024-01-23T15:25:33.701144Z","iopub.status.idle":"2024-01-23T15:25:33.708598Z","shell.execute_reply.started":"2024-01-23T15:25:33.701109Z","shell.execute_reply":"2024-01-23T15:25:33.707462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualize data augmentation in both datasets\nrandn = random.randint(1,298)\nshow_images(n=randn, dataset=monet_dataset)\nshow_images(n=randn, dataset=photo_dataset)","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:33.954178Z","iopub.execute_input":"2024-01-23T15:25:33.955068Z","iopub.status.idle":"2024-01-23T15:25:34.624737Z","shell.execute_reply.started":"2024-01-23T15:25:33.955031Z","shell.execute_reply":"2024-01-23T15:25:34.623749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Models used for the generator and discriminator\n\n#### Here, I've implemented a classic Critic and Generator which I could late on use in the way of WGAN paper with the Gradient Penalty loss.\n#### This model and its training was inspired by <a href=\"https://www.kaggle.com/code/enzou3/simple-gans-in-monets-paintings/notebook\"> this notebook </a> and <a href=\" https://www.youtube.com/watch?v=pG0QZ7OddX4&t=1354s \"> this video tutorial on WGAN</a> Which was made as a part of the amazing series done by the Alladin Persson youtube channel, be sure to check it out!","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n\n        channels_img = 3\n        features_d = 64\n        \n        self.model = nn.Sequential(\n            # Input : N x 3 x 256 x 256\n            nn.Conv2d(channels_img, features_d, kernel_size=4, stride=2, padding=1),\n            nn.LeakyReLU(0.2),\n\n            # Downsampling blocks\n            self._block(features_d, features_d * 2, 4, 2, 1),\n            self._block(features_d * 2, features_d * 4, 4, 2, 1),\n            self._block(features_d * 4, features_d * 8, 4, 2, 1),\n\n            # Upscaling blocks\n            self._conv2d_transpose(features_d * 8, features_d * 4, 3, 2, 1, 1),\n            self._conv2d_transpose(features_d * 4, features_d * 2, 3, 2, 1, 1),\n            self._conv2d_transpose(features_d * 2, features_d, 3, 2, 1, 1),\n\n            # Final convolutional layer\n            nn.Conv2d(features_d, 1, kernel_size=3, stride=2, padding=1)\n        )\n\n    def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),\n            nn.InstanceNorm2d(out_channels, affine=True),\n            nn.LeakyReLU(0.2),\n        )\n\n    def _conv2d_transpose(self, in_channels, out_channels, kernel_size, stride, padding, output_padding):\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=False),\n            nn.InstanceNorm2d(out_channels, affine=True),\n            nn.LeakyReLU(0.2),\n        )\n\n    def forward(self, x):\n        return self.model(x)\n    \n    \nclass Generator(nn.Module):\n    def __init__(self, num_residual_blocks=9):\n        super(Generator, self).__init__()\n        \n        channels_img = 3\n        features_g = 64\n        \n        # Initial convolution layer\n        self.initial = nn.Sequential(\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(channels_img, features_g, kernel_size=7, stride=1, padding=0),\n            nn.InstanceNorm2d(features_g),\n            nn.ReLU(inplace=True)\n        )\n\n        # Downsample layers\n        self.downsample_blocks = nn.Sequential(\n            self._downsample_block(features_g, features_g * 2),\n            self._downsample_block(features_g * 2, features_g * 4)\n        )\n\n        # Residual blocks\n        self.residual_blocks = nn.Sequential(\n            *[ResidualBlock(features_g * 4) for _ in range(num_residual_blocks)]\n        )\n\n        # Upsample layers\n        self.upsample_blocks = nn.Sequential(\n            self._upsample_block(features_g * 4, features_g * 2),\n            self._upsample_block(features_g * 2, features_g)\n        )\n\n        # Output layer\n        self.output = nn.Sequential(\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(features_g, channels_img, kernel_size=7, stride=1, padding=0),\n            nn.Tanh()\n        )\n\n    def _downsample_block(self, in_channels, out_channels):\n        return nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def _upsample_block(self, in_channels, out_channels):\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n        x = self.downsample_blocks(x)\n        x = self.residual_blocks(x)\n        x = self.upsample_blocks(x)\n        return self.output(x)\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super(ResidualBlock, self).__init__()\n        self.block = nn.Sequential(\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(channels, channels, 3),\n            nn.InstanceNorm2d(channels),\n            nn.ReLU(inplace=True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(channels, channels, 3),\n            nn.InstanceNorm2d(channels)\n        )\n\n    def forward(self, x):\n        return x + self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:34.626181Z","iopub.execute_input":"2024-01-23T15:25:34.626494Z","iopub.status.idle":"2024-01-23T15:25:34.649969Z","shell.execute_reply.started":"2024-01-23T15:25:34.626469Z","shell.execute_reply":"2024-01-23T15:25:34.649041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Instantiating generator and discriminator object and declaring loss functions and optimezers used\n#### As shown in the quoted video, I've used the RMSprop optimizer, since it works better on WGAN's, providing stable training.","metadata":{}},{"cell_type":"code","source":"G = Generator()\nD = Discriminator()\n\nmodels = [G, D]\nfor model in models:\n    if torch.cuda.is_available():\n        model.cuda()\n\noptimizer_G = torch.optim.RMSprop(G.parameters(), lr=lr)\noptimizer_D = torch.optim.RMSprop(D.parameters(), lr=lr)","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:36.029853Z","iopub.execute_input":"2024-01-23T15:25:36.030735Z","iopub.status.idle":"2024-01-23T15:25:36.325838Z","shell.execute_reply.started":"2024-01-23T15:25:36.030700Z","shell.execute_reply":"2024-01-23T15:25:36.324849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Auxiliary function to display artificial monet images before and after generator training\n#### This function proved very effective at demonstrating the progress of the generator with traning.","metadata":{}},{"cell_type":"code","source":"\"\"\"\nSample images from batch size n with current generator model parameters\n\nargs:\n        images with format (n x 3 x size(image) x size(image))\n\nreturns:\n        plot's with original and transformed image side-by-side\n\"\"\"\n\ndef sample_images(real_X):    \n    # turn off some learning parameters in order to evaluate model:\n    G.eval()\n    \n    # transform to tensor and plot images\n    real_pic = real_X.type(Tensor)\n    artificial_pic = G(real_X).detach()\n\n    ncols = real_X.size(0)\n    real_pic_grid = make_grid(real_pic, nrow=ncols, normalize=True)\n    artificial_pic_grid = make_grid(artificial_pic, nrow=ncols, normalize=True)\n\n    fig, axs = plt.subplots(2, 1, figsize=(20, 4))  \n\n    axs[0].imshow(real_pic_grid.permute(1, 2, 0).cpu())\n    axs[0].set_title(\"Augmented real photos from photo dataset\")\n    axs[0].axis('off')\n\n    axs[1].imshow(artificial_pic_grid.permute(1, 2, 0).cpu())\n    axs[1].set_title(\"Generated monet paintings from photos\")\n    axs[1].axis('off')\n\n\n    plt.show()\n\nreal_X = next(iter(photo_loader))\nreal_Y = next(iter(photo_loader))\nsample_images(real_X.to(device))","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:25:36.732324Z","iopub.execute_input":"2024-01-23T15:25:36.733287Z","iopub.status.idle":"2024-01-23T15:25:38.034922Z","shell.execute_reply.started":"2024-01-23T15:25:36.733240Z","shell.execute_reply":"2024-01-23T15:25:38.034001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Network training\n#### Another very significant part of the project!\n#### In a nutshell, the generator losses were divided in 2 parts: <code>identity loss</code> and <code>gan loss</code> and both had to impact the backward propagation in different proportions. Therefore both were added and used in the optimizer.\n#### The dicriminator loss was also divided in two parts:  <code>real painting missclassification loss</code> and  <code>artificial painting missclassification loss</code>. Basically the loss when misclassifying a fake painting as a true painting and misclassifying a true painting as a fake paiting.","metadata":{}},{"cell_type":"code","source":"CRITIC_ITERATIONS = 5\nWEIGHT_CLIP = 0.01\nlambda_identity = 5\nidentity_loss_function = nn.L1Loss()\n\nfor epoch in tqdm(range(n_epoches)):\n    for i, (real_photo, real_monet) in enumerate(zip(photo_loader, monet_loader)):\n\n        # Prepare the real and generated data\n        real_photo = real_photo.type(Tensor)\n        real_monet = real_monet.type(Tensor)\n\n        # Train the Critic - We wish to maximize the difference between the fake and the real images - Max(E[critic(real)] - E[critic(fake)])\n        for _ in range(CRITIC_ITERATIONS):\n            optimizer_D.zero_grad()\n\n            # Generate a batch of images\n            artificial_monet = G(real_photo)\n\n            # Critic scores for real and fake images\n            critic_real = D(real_monet).reshape(-1)\n            critic_fake = D(artificial_monet.detach()).reshape(-1)\n\n            # Calculate the loss for the critic\n            loss_critic = -(torch.mean(critic_real) - torch.mean(critic_fake))\n            loss_critic.backward(retain_graph=True)\n            optimizer_D.step()\n\n            # Weight clipping - to satisfy lipshitz constraint\n            for p in D.parameters():\n                p.data.clamp_(-WEIGHT_CLIP, WEIGHT_CLIP)\n\n        # Train the Generator - We wish to minimize \n        optimizer_G.zero_grad()\n\n        # Since we just updated D, perform another forward pass of all-fake batch through D\n        gen_fake = D(artificial_monet).reshape(-1)\n\n        # Calculate G's loss\n        loss_G_Wasserstein = -torch.mean(gen_fake)\n\n        # Calculate identity loss\n        identity_loss = identity_loss_function(artificial_monet, real_photo)\n\n        # Total generator loss\n        loss_G = loss_G_Wasserstein + lambda_identity * identity_loss\n        \n        loss_G.backward()\n        optimizer_G.step()\n        \n    if (epoch + 1) % display_epoch == 0:\n        test_real_photo, test_real_monet = next(iter(zip(photo_loader, monet_loader)))\n        sample_images(test_real_photo.type(Tensor))\n\n        print(f'Epoch {epoch + 1}/{n_epoches}')\n        print(f'Generator loss: {loss_G.item():.4f}')\n        print(f'Discriminator (Critic) loss: {loss_critic.item():.4f}')\n    ","metadata":{"execution":{"iopub.status.busy":"2024-01-23T15:26:46.850012Z","iopub.execute_input":"2024-01-23T15:26:46.850866Z","iopub.status.idle":"2024-01-23T16:33:39.495929Z","shell.execute_reply.started":"2024-01-23T15:26:46.850831Z","shell.execute_reply":"2024-01-23T16:33:39.495002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Plot losses","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create directory to save new images\n#### Pretty straightfoward, just save the images batch by batch","metadata":{}},{"cell_type":"markdown","source":"### Submit!","metadata":{}},{"cell_type":"code","source":"# create path to save created monet paintings\ntransformed_save_dir = '../images'\nif not os.path.exists(transformed_save_dir):\n    os.makedirs(transformed_save_dir)\n    \n# freeze some network params to generate imgs\nG.eval()\n\n\nfor i, real_photos in enumerate(photo_loader):\n    real_photos = real_photos.to(device)\n    with torch.no_grad():\n        monet_style_imgs = G(real_photos)\n\n    # Save each transformed image\n    for j, img in enumerate(monet_style_imgs):\n        save_path = os.path.join(transformed_save_dir, f'monet_painting_{i * batch_size + j}.png')\n        torchvision.utils.save_image(img, save_path)\n\nprint(f\"Transformed images are saved in {transformed_save_dir}\")\n\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2024-01-23T16:41:15.128166Z","iopub.execute_input":"2024-01-23T16:41:15.129035Z","iopub.status.idle":"2024-01-23T16:47:02.858121Z","shell.execute_reply.started":"2024-01-23T16:41:15.129000Z","shell.execute_reply":"2024-01-23T16:47:02.857195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## References\n#### Most of my work was inspired <a href=\"https://www.kaggle.com/code/enzou3/simple-gans-in-monets-paintings/notebook\"> this notebook </a>, and the \"Alladin Persson\" Youtube channel, huge shoutout!\n#### The monet and photos dataset provided belongs to <a href=\"https://www.kaggle.com\">Kaggle</a>.\n#### Thanks for checking out my project! :)","metadata":{}}]}