{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","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":30674,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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-04-05T10:20:35.299105Z","iopub.execute_input":"2024-04-05T10:20:35.299583Z","iopub.status.idle":"2024-04-05T10:20:41.748767Z","shell.execute_reply.started":"2024-04-05T10:20:35.299537Z","shell.execute_reply":"2024-04-05T10:20:41.747736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:41.750451Z","iopub.execute_input":"2024-04-05T10:20:41.751041Z","iopub.status.idle":"2024-04-05T10:20:41.759922Z","shell.execute_reply.started":"2024-04-05T10:20:41.751014Z","shell.execute_reply":"2024-04-05T10:20:41.758734Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:41.760961Z","iopub.execute_input":"2024-04-05T10:20:41.761250Z","iopub.status.idle":"2024-04-05T10:20:41.791489Z","shell.execute_reply.started":"2024-04-05T10:20:41.761226Z","shell.execute_reply":"2024-04-05T10:20:41.790377Z"},"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-04-05T10:20:41.793887Z","iopub.execute_input":"2024-04-05T10:20:41.794250Z","iopub.status.idle":"2024-04-05T10:20:42.192488Z","shell.execute_reply.started":"2024-04-05T10:20:41.794218Z","shell.execute_reply":"2024-04-05T10:20:42.191659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_loader.dataset.__len__()","metadata":{"execution":{"iopub.status.busy":"2024-04-05T10:20:42.193599Z","iopub.execute_input":"2024-04-05T10:20:42.193886Z","iopub.status.idle":"2024-04-05T10:20:42.200928Z","shell.execute_reply.started":"2024-04-05T10:20:42.193861Z","shell.execute_reply":"2024-04-05T10:20:42.199992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:42.202119Z","iopub.execute_input":"2024-04-05T10:20:42.202457Z","iopub.status.idle":"2024-04-05T10:20:42.210324Z","shell.execute_reply.started":"2024-04-05T10:20:42.202433Z","shell.execute_reply":"2024-04-05T10:20:42.209456Z"},"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-04-05T10:20:42.211616Z","iopub.execute_input":"2024-04-05T10:20:42.211888Z","iopub.status.idle":"2024-04-05T10:20:42.790142Z","shell.execute_reply.started":"2024-04-05T10:20:42.211865Z","shell.execute_reply":"2024-04-05T10:20:42.789158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:42.793638Z","iopub.execute_input":"2024-04-05T10:20:42.793940Z","iopub.status.idle":"2024-04-05T10:20:42.818830Z","shell.execute_reply.started":"2024-04-05T10:20:42.793914Z","shell.execute_reply":"2024-04-05T10:20:42.817643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:42.820485Z","iopub.execute_input":"2024-04-05T10:20:42.821651Z","iopub.status.idle":"2024-04-05T10:20:43.162388Z","shell.execute_reply.started":"2024-04-05T10:20:42.821618Z","shell.execute_reply":"2024-04-05T10:20:43.161602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:43.164969Z","iopub.execute_input":"2024-04-05T10:20:43.165290Z","iopub.status.idle":"2024-04-05T10:20:44.407637Z","shell.execute_reply.started":"2024-04-05T10:20:43.165259Z","shell.execute_reply":"2024-04-05T10:20:44.406753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T10:20:44.408695Z","iopub.execute_input":"2024-04-05T10:20:44.408987Z","iopub.status.idle":"2024-04-05T11:27:41.973114Z","shell.execute_reply.started":"2024-04-05T10:20:44.408961Z","shell.execute_reply":"2024-04-05T11:27:41.972223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-04-05T11:27:41.974205Z","iopub.execute_input":"2024-04-05T11:27:41.974488Z"},"trusted":true},"execution_count":null,"outputs":[]}]}