{"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":"import glob\nimport os\nimport shutil\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn, Tensor\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\nfrom torchvision.io import read_image\nfrom torchvision.utils import make_grid, save_image","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-23T06:11:23.733520Z","iopub.execute_input":"2023-03-23T06:11:23.733850Z","iopub.status.idle":"2023-03-23T06:11:37.668775Z","shell.execute_reply.started":"2023-03-23T06:11:23.733795Z","shell.execute_reply":"2023-03-23T06:11:37.667674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_img(img_tensor, nrow=5, title=\"\"):\n    img_tensor = img_tensor.detach().cpu()*0.5 + 0.5\n    img_grid = make_grid(img_tensor, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(18, 8))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.673219Z","iopub.execute_input":"2023-03-23T06:11:37.674240Z","iopub.status.idle":"2023-03-23T06:11:37.682622Z","shell.execute_reply.started":"2023-03-23T06:11:37.674200Z","shell.execute_reply":"2023-03-23T06:11:37.681574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomRotation(object):\n    def __call__(self, x):\n        angle = np.random.choice([0., 90, 180, 270])\n        return transforms.functional.rotate(x, angle)\n\nclass CustomTransform(object):\n    def __init__(self, load_dim=286, target_dim=256):\n        self.transform = transforms.Compose([\n            transforms.Resize((load_dim, load_dim)),\n            transforms.RandomCrop((target_dim, target_dim)),\n            CustomRotation(),\n            transforms.RandomHorizontalFlip(0.5),\n            transforms.RandomVerticalFlip(0.5),\n        ])\n        \n    def __call__(self, img, stage=\"fit\"):\n        if stage == \"fit\":\n            img = self.transform(img)\n        return img/127.5 - 1","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.683972Z","iopub.execute_input":"2023-03-23T06:11:37.684226Z","iopub.status.idle":"2023-03-23T06:11:37.716102Z","shell.execute_reply.started":"2023-03-23T06:11:37.684200Z","shell.execute_reply":"2023-03-23T06:11:37.715041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, monet_filenames, photo_filenames, transform, stage):\n        self.photo_filenames = photo_filenames\n        self.transform = transform\n        self.stage = stage\n        if stage == \"fit\":\n            self.monet_filenames = monet_filenames\n            self.check = len(monet_filenames)>len(photo_filenames)\n            self.shuffle()\n            \n    def shuffle(self):\n        if self.check:\n            self.indices = torch.randperm(len(self.monet_filenames))[:len(self.photo_filenames)]\n            \n        else:\n            self.indices = torch.randperm(len(self.photo_filenames))[:len(self.monet_filenames)]\n        \n    def __len__(self):\n        if self.stage == \"fit\":\n            return min(\n                len(self.monet_filenames), len(self.photo_filenames),\n            )\n    \n        elif self.stage == \"predict\":\n            return len(self.photo_filenames)\n    \n    def __getitem__(self, idx):\n        if self.stage == \"fit\":\n            monet_idx = self.indices[idx] if self.check else idx\n            photo_idx = idx if self.check else self.indices[idx]\n            monet_name = self.monet_filenames[monet_idx]\n            photo_name = self.photo_filenames[photo_idx]\n            monet = read_image(monet_name)\n            photo = read_image(photo_name)\n            if idx == len(self)-1:\n                self.shuffle()\n            return self.transform(monet), self.transform(photo)\n    \n        elif self.stage == \"predict\":\n            photo_name = self.photo_filenames[idx]\n            photo = read_image(photo_name)\n            return self.transform(photo, stage=self.stage)","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.720711Z","iopub.execute_input":"2023-03-23T06:11:37.721504Z","iopub.status.idle":"2023-03-23T06:11:37.733247Z","shell.execute_reply.started":"2023-03-23T06:11:37.721473Z","shell.execute_reply":"2023-03-23T06:11:37.732053Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MONET_DIR = \"/kaggle/input/gan-getting-started/monet_jpg/*.jpg\"\nPHOTO_DIR = \"/kaggle/input/gan-getting-started/photo_jpg/*.jpg\"\nBATCH_SIZE = 1\ntransform = CustomTransform()\nBATCH_SIZE = 4\ntransform = CustomTransform()","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.734609Z","iopub.execute_input":"2023-03-23T06:11:37.735016Z","iopub.status.idle":"2023-03-23T06:11:37.747846Z","shell.execute_reply.started":"2023-03-23T06:11:37.734976Z","shell.execute_reply":"2023-03-23T06:11:37.746852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule(pl.LightningDataModule):\n    def __init__(\n        self, \n        monet_dir=MONET_DIR,\n        photo_dir=PHOTO_DIR,\n        batch_size=BATCH_SIZE,\n        transform=transform,\n    ):\n        super().__init__()\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        self.batch_size = batch_size\n        self.transform = transform\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.train = CustomDataset(\n                monet_filenames=self.monet_filenames,\n                photo_filenames=self.photo_filenames,\n                transform=self.transform,\n                stage=stage,\n            )\n        \n        elif stage == \"predict\":\n            self.predict = CustomDataset(\n                monet_filenames=self.monet_filenames,\n                photo_filenames=self.photo_filenames,\n                transform=self.transform,\n                stage=stage,\n            )\n            \n    def train_dataloader(self):\n        return DataLoader(\n            self.train, \n            batch_size=self.batch_size,\n            shuffle=True,\n        )\n    \n    def predict_dataloader(self):\n        return DataLoader(\n            self.predict, \n            batch_size=self.batch_size,\n            shuffle=False,\n        )","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.749404Z","iopub.execute_input":"2023-03-23T06:11:37.749766Z","iopub.status.idle":"2023-03-23T06:11:37.760125Z","shell.execute_reply.started":"2023-03-23T06:11:37.749721Z","shell.execute_reply":"2023-03-23T06:11:37.759148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm_sample = CustomDataModule(batch_size=5)\ndm_sample.setup(\"fit\")\ndm_sample.setup(\"predict\")\n\ntrain_loader = dm_sample.train_dataloader()\nmonet_samples, _ = next(iter(train_loader))\n\npredict_loader = dm_sample.predict_dataloader()\nphoto_samples = next(iter(predict_loader))\n\nshow_img(monet_samples, title=\"Samples of Monet Paintings\")\nshow_img(photo_samples, title=\"Samples of Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:37.761902Z","iopub.execute_input":"2023-03-23T06:11:37.762268Z","iopub.status.idle":"2023-03-23T06:11:39.340459Z","shell.execute_reply.started":"2023-03-23T06:11:37.762232Z","shell.execute_reply":"2023-03-23T06:11:39.338467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        norm=True,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,\n                      stride=stride, padding=padding, bias=False)\n        )\n        if norm:\n            self.block.append(nn.InstanceNorm2d(out_channels, affine=True))\n        self.block.append(nn.LeakyReLU(0.3))\n        \n    def forward(self, x:Tensor)->Tensor:\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.342029Z","iopub.execute_input":"2023-03-23T06:11:39.343143Z","iopub.status.idle":"2023-03-23T06:11:39.352718Z","shell.execute_reply.started":"2023-03-23T06:11:39.343100Z","shell.execute_reply":"2023-03-23T06:11:39.350067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Upsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        dropout=False,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size,\n                               stride=stride, padding=padding, bias=False),\n            nn.InstanceNorm2d(out_channels, affine=True),\n        )\n        if dropout:\n            self.block.append(nn.Dropout(0.5))\n        self.block.append(nn.ReLU())\n        \n    def forward(self, x:Tensor)->Tensor:\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.354418Z","iopub.execute_input":"2023-03-23T06:11:39.355667Z","iopub.status.idle":"2023-03-23T06:11:39.366893Z","shell.execute_reply.started":"2023-03-23T06:11:39.355626Z","shell.execute_reply":"2023-03-23T06:11:39.366070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResNetBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):\n        super().__init__()\n        norm_layer = nn.InstanceNorm2d\n\n        self.conv1 = nn.Conv2d(\n            in_channels,\n            out_channels,\n            kernel_size=kernel_size,\n            stride=1,\n            padding=1,\n            bias=False,\n        )\n        self.bn1 = norm_layer(out_channels, affine=True)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = nn.Conv2d(\n            out_channels,\n            out_channels,\n            kernel_size=kernel_size,\n            stride=1,\n            padding=1,\n            bias=False,\n        )\n        self.bn2 = norm_layer(out_channels, affine=True)\n\n    def forward(self, x: Tensor) -> Tensor:\n        sc = x\n\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n\n        x = self.conv2(x)\n        x = self.bn2(x)\n        x = self.relu(x + sc)\n\n        return x\n\n\nclass Generator(nn.Module):\n    def __init__(self, in_channels, out_channels, dim, depth=8):\n        super().__init__()\n        self.downsampling_blocks = nn.Sequential(\n            Downsampling(in_channels, dim, norm=False),  # 64x128x128\n            Downsampling(dim, dim * 2),  # 128x64x64\n            Downsampling(dim * 2, dim * 4),  # 256x32x32\n        )\n\n        # 256*32*32 * depth times resnet block\n        self.resnet = nn.Sequential(\n            *[ResNetBlock(dim * 4, dim * 4) for _ in range(depth)]\n        )\n\n        self.upsampling_blocks = nn.Sequential(\n            Upsampling(dim * 4, dim * 2),  # 128*64*64\n            Upsampling(dim * 2, dim),  # 64*128*128\n        )\n\n        # 3*256*256 (RGB space)\n        self.to_rgb_layer = nn.Sequential(\n            nn.ConvTranspose2d(\n                dim,\n                out_channels,\n                kernel_size=4,\n                stride=2,\n                padding=1,\n            ),\n            nn.Tanh(),\n        )\n\n    def forward(self, x: Tensor) -> Tensor:\n        x = self.downsampling_blocks(x)\n        x = self.resnet(x)\n        x = self.upsampling_blocks(x)\n        return self.to_rgb_layer(x)\n","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.371160Z","iopub.execute_input":"2023-03-23T06:11:39.371534Z","iopub.status.idle":"2023-03-23T06:11:39.385867Z","shell.execute_reply.started":"2023-03-23T06:11:39.371502Z","shell.execute_reply":"2023-03-23T06:11:39.384746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_channels, hid_channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False),  # 64x128x128\n            Downsampling(hid_channels, hid_channels * 2),  # 128x64x64\n            Downsampling(hid_channels * 2, hid_channels * 4),  # 256x32x32\n            Downsampling(\n                hid_channels * 4, hid_channels * 8, stride=1\n            ),  # 512x31x31\n            nn.Conv2d(hid_channels * 8, 1, kernel_size=4, padding=1),  # 1x30x30\n        )\n\n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.387593Z","iopub.execute_input":"2023-03-23T06:11:39.387965Z","iopub.status.idle":"2023-03-23T06:11:39.400405Z","shell.execute_reply.started":"2023-03-23T06:11:39.387927Z","shell.execute_reply":"2023-03-23T06:11:39.399321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IN_CHANNELS = 3\nOUT_CHANNELS = 3\nHID_CHANNELS = 64\nLR = 2e-4\nBETAS = (0.5, 0.999)\nLAMBDA = 10\nNUM_EPOCHS = 180\nDISPLAY_EPOCHS = 30","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.402222Z","iopub.execute_input":"2023-03-23T06:11:39.402745Z","iopub.status.idle":"2023-03-23T06:11:39.414983Z","shell.execute_reply.started":"2023-03-23T06:11:39.402706Z","shell.execute_reply":"2023-03-23T06:11:39.413823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(pl.LightningModule):\n    def __init__(\n        self,\n        in_channels=IN_CHANNELS,\n        out_channels=OUT_CHANNELS,\n        dim=HID_CHANNELS,\n        lr=LR,\n        betas=BETAS,\n        lambda_w=LAMBDA,\n        display_epochs=DISPLAY_EPOCHS,\n        photo_samples=photo_samples,\n    ):\n        super().__init__()\n        self.lr = lr\n        self.betas = betas\n        self.lambda_w = lambda_w\n        self.display_epochs = display_epochs\n        self.photo_samples = photo_samples\n        self.loss_history = []\n        self.epoch_count = 0\n        self.mae_loss = nn.L1Loss()\n        self.cnt_train_step = 0\n        self.reconstr_w = 10\n        self.id_w = 2\n\n        self.gen_style_to_base = Generator(in_channels, out_channels, dim).apply(\n            self.weights_init\n        )\n        self.gen_base_to_style = Generator(in_channels, out_channels, dim).apply(\n            self.weights_init\n        )\n        self.disc_style = Discriminator(in_channels, dim).apply(self.weights_init)\n        self.disc_base = Discriminator(in_channels, dim).apply(self.weights_init)\n\n    def forward(self, z):\n        return self.gen_base_to_style(z)\n\n    def weights_init(self, m):\n        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n            nn.init.normal_(m.weight, 0.0, 0.02)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0.0)\n\n    def configure_optimizers(self):\n        params = {\n            \"lr\": self.lr,\n            \"betas\": self.betas,\n        }\n        opt_gen_PM = torch.optim.Adam(self.gen_style_to_base.parameters(), **params)\n        opt_gen_MP = torch.optim.Adam(self.gen_base_to_style.parameters(), **params)\n\n        opt_disc_M = torch.optim.Adam(self.disc_style.parameters(), **params)\n        opt_disc_P = torch.optim.Adam(self.disc_base.parameters(), **params)\n\n        return [opt_gen_PM, opt_gen_MP, opt_disc_M, opt_disc_P], []\n\n    def training_step(self, batch, batch_idx, optimizer_idx):\n        style_image, base_image = batch\n        b = style_image.size()[0]\n        valid = torch.ones(b, 1, 30, 30).cuda()\n        fake = torch.zeros(b, 1, 30, 30).cuda()\n\n        if optimizer_idx == 0 or optimizer_idx == 1:\n            # Validity\n            # MSELoss\n            val_base = F.mse_loss(\n                self.disc_base(self.gen_style_to_base(style_image)), valid\n            )\n            val_style = F.mse_loss(\n                self.disc_base(self.gen_base_to_style(base_image)), valid\n            )\n            val_loss = (val_base + val_style) / 2\n\n            # Reconstruction\n            reconstr_base = self.mae_loss(\n                self.gen_style_to_base(self.gen_base_to_style(base_image)), base_image\n            )\n            reconstr_style = self.mae_loss(\n                self.gen_base_to_style(self.gen_style_to_base(style_image)), style_image\n            )\n            reconstr_loss = (reconstr_base + reconstr_style) / 2\n\n            # Identity\n            id_base = self.mae_loss(self.gen_style_to_base(base_image), base_image)\n            id_style = self.mae_loss(self.gen_base_to_style(style_image), style_image)\n            id_loss = (id_base + id_style) / 2\n\n            # Loss Weight\n            G_loss = val_loss + self.reconstr_w * reconstr_loss + self.id_w * id_loss\n\n            loss_dict = {\n                \"loss\": G_loss,\n                \"validity\": val_loss,\n                \"reconstr\": reconstr_loss,\n                \"identity\": id_loss,\n            }\n            self.log_dict(loss_dict, prog_bar=True)\n            return loss_dict\n        # Train Discriminator\n        elif optimizer_idx == 2 or optimizer_idx == 3:\n            # MSELoss\n            disc_base_gen_loss = F.mse_loss(\n                self.disc_base(self.gen_style_to_base(style_image)), fake\n            )\n            disc_style_gen_loss = F.mse_loss(\n                self.disc_style(self.gen_base_to_style(base_image)), fake\n            )\n            disc_base_valid_loss = F.mse_loss(self.disc_base(base_image), valid)\n            disc_style_valid_loss = F.mse_loss(self.disc_style(style_image), valid)\n\n            D_gen_loss = (disc_base_gen_loss + disc_style_gen_loss) / 2\n\n            # Loss Weight\n            D_loss = (D_gen_loss + disc_base_valid_loss + disc_style_valid_loss) / 3\n\n            # Count up\n            self.cnt_train_step += 1\n\n            loss_dict = {\"loss\": D_loss}\n            self.log_dict(loss_dict, prog_bar=True)\n            return loss_dict\n\n    def training_epoch_end(self, outputs):\n        self.epoch_count += 1\n\n        if self.epoch_count % 10 == 0:\n            with torch.inference_mode():\n                gen_monets = self(self.photo_samples.to(self.device)).cpu()\n            show_img(\n                torch.cat([self.photo_samples, gen_monets]),\n                nrow=5,\n                title=f\"Epoch {self.epoch_count}: Photo-to-Monet Translation\",\n            )\n\n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n\n    def loss_curves(self):\n        labels = [\"gen_loss_PM\", \"gen_loss_MP\", \"disc_loss_M\", \"disc_loss_P\"]\n        titles = [\"Generator Loss Curves\", \"Discriminator Loss Curves\"]\n        num_epochs = len(self.loss_history)\n        plt.figure(figsize=(18, 4.5))\n        for j in range(4):\n            if j % 2 == 0:\n                plt.subplot(1, 2, (j // 2) + 1)\n                plt.title(titles[j // 2])\n                plt.ylabel(\"Loss\")\n                plt.xlabel(\"Epoch\")\n            plt.plot(\n                np.arange(1, num_epochs + 1),\n                [losses[j] for losses in self.loss_history],\n                label=labels[j],\n            )\n            plt.legend(loc=\"upper right\")\n","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.417567Z","iopub.execute_input":"2023-03-23T06:11:39.418399Z","iopub.status.idle":"2023-03-23T06:11:39.444594Z","shell.execute_reply.started":"2023-03-23T06:11:39.418346Z","shell.execute_reply":"2023-03-23T06:11:39.443585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = pl.Trainer(\n    accelerator=\"gpu\",\n    devices=1,\n    logger=False,\n    enable_checkpointing=False,\n    max_epochs=NUM_EPOCHS,\n)\n\ndm = CustomDataModule()\nmodel = CycleGAN()\ntrainer.fit(model, datamodule=dm)\n","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:11:39.446559Z","iopub.execute_input":"2023-03-23T06:11:39.447375Z","iopub.status.idle":"2023-03-23T06:53:31.010935Z","shell.execute_reply.started":"2023-03-23T06:11:39.447335Z","shell.execute_reply":"2023-03-23T06:53:31.009763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)\n\nos.makedirs(\"../images\", exist_ok=True)\n\nidx = 0\nfor tensor in predictions:\n    for monet in tensor:\n        save_image((monet.squeeze()*0.5+0.5), fp=f\"../images/{idx}.jpg\")\n        idx += 1\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2023-03-23T06:53:31.012869Z","iopub.execute_input":"2023-03-23T06:53:31.013277Z","iopub.status.idle":"2023-03-23T06:53:33.613486Z","shell.execute_reply.started":"2023-03-23T06:53:31.013214Z","shell.execute_reply":"2023-03-23T06:53:33.610884Z"},"trusted":true},"execution_count":null,"outputs":[]}]}