{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30627,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\n\nimport torch \nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.optim as optim\nimport torchvision\nfrom torchvision import transforms\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom tqdm import tqdm\n\nfrom PIL import Image\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-01-21T03:01:59.918189Z","iopub.execute_input":"2024-01-21T03:01:59.918947Z","iopub.status.idle":"2024-01-21T03:01:59.924598Z","shell.execute_reply.started":"2024-01-21T03:01:59.918912Z","shell.execute_reply":"2024-01-21T03:01:59.923624Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Showing images \n\ndef show_images(img, nrow=5, title=\"\"):\n    img = img.detach().cpu() * 0.5 + 0.5\n    img_grid = torchvision.utils.make_grid(img, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(8, 5))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:33.188565Z","iopub.execute_input":"2024-01-20T16:22:33.188934Z","iopub.status.idle":"2024-01-20T16:22:33.195468Z","shell.execute_reply.started":"2024-01-20T16:22:33.18889Z","shell.execute_reply":"2024-01-20T16:22:33.194425Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Creating dataset class\n\nclass MonetPhotoDataset(Dataset):\n    def __init__(self, folder_path_monets, folder_path_photos, transform=None):\n        self.folder_path_monets = folder_path_monets\n        self.folder_path_photos = folder_path_photos\n        self.transform = transform\n        self.monets = [os.path.join(folder_path_monets, f) for f in os.listdir(folder_path_monets) if f.endswith(('.jpg', '.jpeg', '.png'))]\n        self.photos = [os.path.join(folder_path_photos, f) for f in os.listdir(folder_path_photos) if f.endswith(('.jpg', '.jpeg', '.png'))]\n        \n    def __getitem__(self, idx):\n        monet_path = self.monets[idx % len(self.monets)]\n        photo_path = self.photos[idx % len(self.photos)]\n        \n        monet = np.array(Image.open(monet_path).convert('RGB')).astype('float32')\n        photo = np.array(Image.open(photo_path).convert('RGB')).astype('float32')\n        \n                \n        if self.transform:\n            augmentations = transform(image=monet, image0=photo)\n            monet = augmentations[\"image\"]\n            photo = augmentations[\"image0\"]\n            \n        return monet, photo\n    \n    def __len__(self):\n        return max(len(self.monets), len(self.photos))","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:33.457186Z","iopub.execute_input":"2024-01-20T16:22:33.457544Z","iopub.status.idle":"2024-01-20T16:22:33.467628Z","shell.execute_reply.started":"2024-01-20T16:22:33.457513Z","shell.execute_reply":"2024-01-20T16:22:33.46675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Device configuration\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Hyperparameters\n\nnum_epochs = 35\nlr_disc = 2e-4\nlr_gen = 2e-4\nnum_workers = 2\nbatch_size = 8\nimage_size = 128\nchannels_img = 3\n\nLAMBDA_CYCLE = 10.0\nLAMBDA_IDENTITY = 5.0\nSAVE_MODEL = True\n\nCHECKPOINT_DISC_M = \"discM.pth.tar\"\nCHECKPOINT_DISC_P = \"discP.pth.tar\"\nCHECKPOINT_GEN_M = \"genM.pth.tar\"\nCHECKPOINT_GEN_P = \"genP.pth.tar\"\n\n\n# Transformations\n\ntransform = A.Compose(\n    [\n        A.Resize(image_size, image_size),\n        A.HorizontalFlip(),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0),\n        ToTensorV2(),\n    ], additional_targets={\"image0\": \"image\"},\n)\n\ntrain_dataset = MonetPhotoDataset(\"/kaggle/input/gan-getting-started/monet_jpg\", \"/kaggle/input/gan-getting-started/photo_jpg\", transform=transform)\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:33.696742Z","iopub.execute_input":"2024-01-20T16:22:33.697049Z","iopub.status.idle":"2024-01-20T16:22:34.156469Z","shell.execute_reply.started":"2024-01-20T16:22:33.697024Z","shell.execute_reply":"2024-01-20T16:22:34.155403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Total monet images: {len(train_dataset)}\")","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:34.158421Z","iopub.execute_input":"2024-01-20T16:22:34.158736Z","iopub.status.idle":"2024-01-20T16:22:34.16417Z","shell.execute_reply.started":"2024-01-20T16:22:34.15871Z","shell.execute_reply":"2024-01-20T16:22:34.162988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_img, photo_img = next(iter(train_loader))\n\nshow_images(monet_img[:5], nrow=5, title=\"Monet Images\")\nshow_images(photo_img[:5], nrow=5, title=\"Normal Images\")","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:34.179177Z","iopub.execute_input":"2024-01-20T16:22:34.179455Z","iopub.status.idle":"2024-01-20T16:22:35.169444Z","shell.execute_reply.started":"2024-01-20T16:22:34.179431Z","shell.execute_reply":"2024-01-20T16:22:35.168467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Discriminator Model\n\nclass BlockDisc(nn.Module):\n    def __init__(self, in_channels, out_channels, stride=2):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 4, stride, 1, bias=True, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels, features=[64, 128, 256, 512]):\n        super().__init__()\n\n        self.initial = nn.Sequential(\n            nn.Conv2d(in_channels, features[0], kernel_size=4, stride=2, padding=1, padding_mode=\"reflect\"),\n            # paper authors says that we don't use InstanceNorm for the first layer.\n            nn.LeakyReLU(0.2),\n        )\n\n        in_channel = features[0]\n        layers = []\n\n        for feature in features[1:]:\n\n            layers.append(\n                BlockDisc(in_channel, feature, stride=1 if feature == features[-1] else 2),\n            )\n            in_channel = feature\n\n        layers.append(\n            nn.Conv2d(in_channel, 1, kernel_size=4, stride=1, padding=1, padding_mode=\"reflect\")\n        )\n\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:35.171646Z","iopub.execute_input":"2024-01-20T16:22:35.172361Z","iopub.status.idle":"2024-01-20T16:22:35.183967Z","shell.execute_reply.started":"2024-01-20T16:22:35.172324Z","shell.execute_reply":"2024-01-20T16:22:35.182811Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Generator Model\n\nclass BlockGen(nn.Module):\n    def __init__(self, in_channels, out_channels, encoder=True, **kwargs): # **kwargs include kernel_size, stride etc.\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs) if encoder\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(channels, channels, kernel_size=3, padding=1, stride=1, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(channels, channels, kernel_size=3, padding=1, stride=1, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(channels),\n        )\n\n    def forward(self, x):\n        return x + self.block(x) # !!!\n\n\nclass Generator(nn.Module):\n    def __init__(self, in_channels, num_features=64, num_residuals=6):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(in_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\"),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n\n        self.enc_blocks = nn.ModuleList([\n            BlockGen(num_features, num_features*2, encoder=True, kernel_size=3, stride=2, padding=1),\n            BlockGen(num_features*2, num_features*4, encoder=True, kernel_size=3, stride=2, padding=1),\n        ])\n\n        self.residual_blocks = nn.Sequential(\n            *[ResidualBlock(num_features*4) for _ in range(num_residuals)]\n        )\n\n        self.dec_blocks = nn.ModuleList([\n            BlockGen(num_features*4, num_features*2, encoder=False, kernel_size=3, stride=2, padding=1, output_padding=1), # if you don't use output padding test shape is gonna be 253x253 but we want this shape to be 256x256 which same as the input shape.\n            BlockGen(num_features*2, num_features, encoder=False, kernel_size=3, stride=2, padding=1, output_padding=1),\n        ])\n\n        self.last = nn.Sequential(\n            nn.Conv2d(num_features, in_channels, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\"),\n            nn.Tanh(),\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n\n        for layer in self.enc_blocks:\n            x = layer(x)\n\n        x = self.residual_blocks(x)\n\n        for layer in self.dec_blocks:\n            x = layer(x)\n\n        return self.last(x)\n\ndef gen_test():\n    x = torch.randn((10, 3, 128, 128))\n    gen = Generator(in_channels=3, num_residuals=6)\n    img = gen(x)\n    print(img.shape)\n    \ngen_test()","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:35.185482Z","iopub.execute_input":"2024-01-20T16:22:35.185815Z","iopub.status.idle":"2024-01-20T16:22:36.958125Z","shell.execute_reply.started":"2024-01-20T16:22:35.185783Z","shell.execute_reply":"2024-01-20T16:22:36.957139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot for losses\ndef plot(g_loss, d_loss):\n    plt.figure(figsize=(10, 5))\n    plt.title(\"Generator and Discriminator Loss During Training\")\n    plt.plot(g_loss, label=\"Generator\")\n    plt.plot(d_loss, label=\"Discriminator\")\n    plt.xlabel(\"Iterations\")\n    plt.ylabel(\"Loss\")\n    plt.legend()\n    plt.show()\n    \n# Saving the models\ndef save_checkpoint(model, optimizer, file_name = \"my_checkpoint.pth.tar\"):\n    print(\"SAVING CHECKPOINT !\")\n\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n\n    torch.save(checkpoint, file_name)\n","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:36.960287Z","iopub.execute_input":"2024-01-20T16:22:36.961101Z","iopub.status.idle":"2024-01-20T16:22:36.967824Z","shell.execute_reply.started":"2024-01-20T16:22:36.961064Z","shell.execute_reply":"2024-01-20T16:22:36.966783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import itertools\n\nprint(f\"Device: {DEVICE}\")\n\ntorch.backends.cudnn.benchmark = True\n\ndisc_P = Discriminator(in_channels=3).to(DEVICE)\ndisc_M = Discriminator(in_channels=3).to(DEVICE)\ngen_P = Generator(in_channels=3, num_residuals=9).to(DEVICE)\ngen_M = Generator(in_channels=3, num_residuals=9).to(DEVICE)\n\nopt_disc = optim.Adam(\n    list(disc_P.parameters()) + list(disc_M.parameters()) ,\n    lr=lr_disc,\n    betas=(0.5, 0.999)\n)\nopt_gen = optim.Adam(\n    itertools.chain(gen_P.parameters(), gen_M.parameters()),\n    lr=lr_gen,\n    betas=(0.5, 0.999)\n)\n\nL1_loss = nn.L1Loss()\nmse_loss = nn.MSELoss()","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:37.321498Z","iopub.execute_input":"2024-01-20T16:22:37.322404Z","iopub.status.idle":"2024-01-20T16:22:37.752889Z","shell.execute_reply.started":"2024-01-20T16:22:37.32237Z","shell.execute_reply":"2024-01-20T16:22:37.751775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"decay_epoch = 20\n\nlambda_func = lambda epoch: 1 - max(0, epoch-decay_epoch)/(num_epochs-decay_epoch)\n\nlr_scheduler_D = torch.optim.lr_scheduler.LambdaLR(opt_disc, lr_lambda=lambda_func)\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(opt_gen, lr_lambda=lambda_func)","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:22:38.389144Z","iopub.execute_input":"2024-01-20T16:22:38.389483Z","iopub.status.idle":"2024-01-20T16:22:38.395042Z","shell.execute_reply.started":"2024-01-20T16:22:38.389457Z","shell.execute_reply":"2024-01-20T16:22:38.393907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training\n\nd_scaler = torch.cuda.amp.GradScaler()\ng_scaler = torch.cuda.amp.GradScaler()\n\n\nD_LOSS = []\nG_LOSS = []\nfor epoch in range(num_epochs):\n    loop = tqdm(train_loader, leave=True)\n    # --------------------------------------------------------------------------------------------------\n    print(f\"EPOCH {epoch+1}\")\n    for idx, batch in enumerate(loop):\n        m, p = batch\n        m, p = m.to(DEVICE), p.to(DEVICE)\n\n        # >>> DISCRIMINATOR\n\n        with torch.cuda.amp.autocast():\n            # FIRST DISCRIMINATOR\n            fake_m = gen_M(p)\n            d_m_real = disc_M(m)\n            d_m_fake = disc_M(fake_m.detach())\n\n            d_m_real_loss = mse_loss(d_m_real, torch.ones_like(d_m_real))\n            d_m_fake_loss = mse_loss(d_m_fake, torch.zeros_like(d_m_fake))\n            d_m_loss = d_m_real_loss + d_m_fake_loss\n\n            # SECOND DISCRIMINATOR\n            fake_p = gen_P(m)\n            d_p_real = disc_P(p)\n            d_p_fake = disc_P(fake_p.detach())\n\n            d_p_real_loss = mse_loss(d_p_real, torch.ones_like(d_p_real))\n            d_p_fake_loss = mse_loss(d_p_fake, torch.zeros_like(d_p_fake))\n            d_p_loss = d_p_real_loss + d_p_fake_loss\n\n            # PUT IT TOGETHER\n            d_loss = (d_m_loss + d_p_loss) / 2\n\n        opt_disc.zero_grad()\n        d_scaler.scale(d_loss).backward()\n        d_scaler.step(opt_disc)\n        d_scaler.update()\n\n        # >>> GENERATOR\n\n        with torch.cuda.amp.autocast():\n            # Adversarial Loss\n            d_m_fake = disc_M(fake_m)\n            d_p_fake = disc_P(fake_p)\n            g_m_loss = mse_loss(d_m_fake, torch.ones_like(d_m_fake))\n            g_p_loss = mse_loss(d_p_fake, torch.ones_like(d_p_fake))\n\n            # Cycle Loss\n            cycle_m = gen_P(fake_m)\n            cycle_p = gen_M(fake_p)\n\n            cycle_m_loss = L1_loss(m, cycle_m)\n            cycle_p_loss = L1_loss(p, cycle_p)\n\n            # Identity Loss\n            identity_m = gen_M(m)\n            identity_p = gen_P(p)\n\n            identity_m_loss = L1_loss(m, identity_m)\n            identity_p_loss = L1_loss(p, identity_p)\n\n            # PUT ALL TOGETHER\n            g_loss = (\n                    g_m_loss\n                    + g_p_loss\n                    + cycle_m_loss * LAMBDA_CYCLE\n                    + cycle_p_loss * LAMBDA_CYCLE\n                    + identity_m_loss * LAMBDA_IDENTITY\n                    + identity_p_loss * LAMBDA_IDENTITY\n            )\n\n        opt_gen.zero_grad()\n        g_scaler.scale(g_loss).backward()\n        g_scaler.step(opt_gen)\n        g_scaler.update()\n\n        D_LOSS.append(d_loss.item())\n        G_LOSS.append(g_loss.item())\n\n        if idx % 200 == 0:\n            loop.set_postfix(\n                d_loss=d_loss.item(),\n                g_loss=g_loss.item(),\n            )\n\n    with torch.no_grad():\n        show_images(img=p[:5], nrow=5, title=\"Real Photos\")\n        fake_m=fake_m.type(torch.float64)\n        show_images(img=fake_m[:5], nrow=5, title=\"Monets\")\n    \n    lr_scheduler_D.step()\n    lr_scheduler_G.step()\n    \n    # --------------------------------------------------------------------------------------------------\n    if SAVE_MODEL and epoch % 5 == 0:\n        save_checkpoint(disc_M, opt_disc, file_name=CHECKPOINT_DISC_M)\n        save_checkpoint(disc_P, opt_disc, file_name=CHECKPOINT_DISC_P)\n        save_checkpoint(gen_M, opt_gen, file_name=CHECKPOINT_GEN_M)\n        save_checkpoint(gen_P, opt_gen, file_name=CHECKPOINT_GEN_P)","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:23:07.057485Z","iopub.execute_input":"2024-01-20T16:23:07.057845Z","iopub.status.idle":"2024-01-20T16:36:28.361843Z","shell.execute_reply.started":"2024-01-20T16:23:07.057815Z","shell.execute_reply":"2024-01-20T16:36:28.360412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot(G_LOSS, D_LOSS)","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:36:35.436952Z","iopub.execute_input":"2024-01-20T16:36:35.437338Z","iopub.status.idle":"2024-01-20T16:36:35.804759Z","shell.execute_reply.started":"2024-01-20T16:36:35.437307Z","shell.execute_reply":"2024-01-20T16:36:35.803975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_dir = '../images'\nif not os.path.exists(save_dir):\n    os.makedirs(save_dir)","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:36:38.824803Z","iopub.execute_input":"2024-01-20T16:36:38.82581Z","iopub.status.idle":"2024-01-20T16:36:38.831092Z","shell.execute_reply.started":"2024-01-20T16:36:38.825772Z","shell.execute_reply":"2024-01-20T16:36:38.829916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generate_transforms = transforms.Compose([\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\n\nto_image = transforms.ToPILImage()\n\ngen_M.eval()\n\ndata_dir = \"/kaggle/input/gan-getting-started/photo_jpg\"\n\nfiles = [os.path.join(data_dir, name) for name in os.listdir(data_dir)]\n\nTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor\n\nfor i in range(0, len(files), batch_size):\n    # read images\n    imgs = []\n    for j in range(i, min(len(files), i+batch_size)):\n        img = Image.open(files[j])\n        img = generate_transforms(img)\n        imgs.append(img)\n    imgs = torch.stack(imgs, 0).type(Tensor)\n    \n    # generate\n    fake_imgs = gen_M(imgs).detach().cpu()\n    \n    # save\n    for j in range(fake_imgs.size(0)):\n        img = fake_imgs[j].squeeze().permute(1, 2, 0)\n        img_arr = img.numpy()\n        img_arr = (img_arr - np.min(img_arr)) * 255 / (np.max(img_arr) - np.min(img_arr))\n        img_arr = img_arr.astype(np.uint8)\n        \n        img = to_image(img_arr)\n        _, name = os.path.split(files[i+j])\n        img.save(os.path.join(save_dir, name))","metadata":{"execution":{"iopub.status.busy":"2024-01-20T16:43:36.151859Z","iopub.execute_input":"2024-01-20T16:43:36.152774Z","iopub.status.idle":"2024-01-20T16:47:04.945785Z","shell.execute_reply.started":"2024-01-20T16:43:36.152738Z","shell.execute_reply":"2024-01-20T16:47:04.944902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2024-01-20T17:14:48.267604Z","iopub.execute_input":"2024-01-20T17:14:48.268043Z","iopub.status.idle":"2024-01-20T17:14:52.566777Z","shell.execute_reply.started":"2024-01-20T17:14:48.268011Z","shell.execute_reply":"2024-01-20T17:14:52.565841Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}