{"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 os\n\nimport torch\nimport torch.nn as nn\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\n\nimport numpy as np\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\n\nimport warnings\nwarnings.simplefilter('ignore')","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:50.296731Z","iopub.execute_input":"2022-04-28T10:42:50.297038Z","iopub.status.idle":"2022-04-28T10:42:52.434775Z","shell.execute_reply.started":"2022-04-28T10:42:50.296958Z","shell.execute_reply":"2022-04-28T10:42:52.434035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MonetPhotoDataset(Dataset):\n    def __init__(self, root_monet, root_photo, transform=None):\n        self.transform = transform\n        self.root_monet = root_monet\n        self.root_photo = root_photo\n        \n        self.monet_images = os.listdir(root_monet)\n        self.photo_images = os.listdir(root_photo)\n        self.length_dataset = max(len(self.monet_images), len(self.photo_images))\n\n        self.monet_len = len(self.monet_images)\n        self.photo_len = len(self.photo_images)\n\n    def __len__(self):\n        return self.length_dataset\n\n    def __getitem__(self, idx):\n        monet_img = self.monet_images[idx % self.monet_len]\n        photo_img = self.photo_images[idx % self.photo_len]\n        \n        monet_path = os.path.join(self.root_monet, monet_img)\n        photo_path = os.path.join(self.root_photo, photo_img)\n        \n        monet_img = Image.open(monet_path).convert('RGB')\n        photo_img = Image.open(photo_path).convert('RGB')\n    \n        \n        monet_img = self.transform(monet_img)\n        photo_img = self.transform(photo_img)\n        \n        return monet_img, photo_img","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:52.436231Z","iopub.execute_input":"2022-04-28T10:42:52.436476Z","iopub.status.idle":"2022-04-28T10:42:52.445754Z","shell.execute_reply.started":"2022-04-28T10:42:52.436442Z","shell.execute_reply":"2022-04-28T10:42:52.445078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize(256),\n    transforms.ToTensor(),\n    transforms.Normalize(0.5, 0.5)\n])","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:52.446797Z","iopub.execute_input":"2022-04-28T10:42:52.447277Z","iopub.status.idle":"2022-04-28T10:42:52.455375Z","shell.execute_reply.started":"2022-04-28T10:42:52.44724Z","shell.execute_reply":"2022-04-28T10:42:52.454601Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = MonetPhotoDataset('../input/gan-getting-started/monet_jpg', '../input/gan-getting-started/photo_jpg', transform)","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:52.45721Z","iopub.execute_input":"2022-04-28T10:42:52.457895Z","iopub.status.idle":"2022-04-28T10:42:52.917976Z","shell.execute_reply.started":"2022-04-28T10:42:52.457857Z","shell.execute_reply":"2022-04-28T10:42:52.917268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:52.919373Z","iopub.execute_input":"2022-04-28T10:42:52.919798Z","iopub.status.idle":"2022-04-28T10:42:52.983299Z","shell.execute_reply.started":"2022-04-28T10:42:52.919761Z","shell.execute_reply":"2022-04-28T10:42:52.98236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResBlock(nn.Module):\n    def __init__(self, f):\n        super(ResBlock, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(f, f, 3, 1, 1), \n            nn.InstanceNorm2d(f),\n            nn.ReLU(),\n            \n            nn.Conv2d(f, f, 3, 1, 1),\n        )\n        \n        self.norm = nn.InstanceNorm2d(f)\n        \n    def forward(self, x):\n        return torch.relu(self.conv(x) + x)\n    \nclass Generator(nn.Module):\n    def __init__(self, f=64, res_blocks=6):\n        super(Generator, self).__init__()\n        layers = []\n            \n        layers.extend([\n            nn.Conv2d(3, f, 7, 1, 3),\n            nn.InstanceNorm2d(f), \n            nn.ReLU(True),\n\n            nn.Conv2d(f, f * 2, 3, 2, 1),\n            nn.InstanceNorm2d(f * 2),\n            nn.ReLU(True),\n\n            nn.Conv2d(f * 2, f * 4, 3, 2, 1),\n            nn.InstanceNorm2d(f * 4),\n            nn.ReLU(True)\n        ])\n            \n        for i in range(res_blocks):\n            layers.append(ResBlock(f * 4))\n        \n        layers.extend([\n            nn.ConvTranspose2d(f * 4, f * 2 * 4, 3, 1, 1),\n            nn.PixelShuffle(2),\n            nn.InstanceNorm2d(f * 2),\n            nn.ReLU(True),\n            \n            nn.ConvTranspose2d(f * 2, f * 4, 3, 1, 1),\n            nn.PixelShuffle(2),\n            nn.InstanceNorm2d(f),\n            nn.ReLU(True),\n    \n            nn.ReflectionPad2d(3),\n            nn.Conv2d(f, 3, 7, 1, 0),\n            nn.Tanh()\n            \n        ])\n         \n        self.conv = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:52.986848Z","iopub.execute_input":"2022-04-28T10:42:52.988766Z","iopub.status.idle":"2022-04-28T10:42:53.002213Z","shell.execute_reply.started":"2022-04-28T10:42:52.988734Z","shell.execute_reply":"2022-04-28T10:42:53.001255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, nc=3, ndf=64):\n        super(Discriminator, self).__init__()\n        self.init = nn.Sequential(\n            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n            nn.InstanceNorm2d(ndf * 2),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n            nn.InstanceNorm2d(ndf * 4),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(ndf * 4, ndf * 8, 4, 1, 1),\n            nn.InstanceNorm2d(ndf * 8),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(ndf * 8, 1, 4, 1, 1),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        output = self.init(x)\n        return output","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:53.004426Z","iopub.execute_input":"2022-04-28T10:42:53.004903Z","iopub.status.idle":"2022-04-28T10:42:53.015526Z","shell.execute_reply.started":"2022-04-28T10:42:53.004868Z","shell.execute_reply":"2022-04-28T10:42:53.014623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"disc_M = Discriminator().to(device)\ndisc_P = Discriminator().to(device)\n\ngen_M = Generator().to(device)\ngen_P = Generator().to(device)\n\nopt_disc = torch.optim.Adam(\n    list(disc_M.parameters()) + list(disc_P.parameters()),\n    lr=5e-4,\n    betas=(0.5, 0.999)\n)\n\nopt_gen = torch.optim.Adam(\n    list(gen_M.parameters()) + list(gen_P.parameters()),\n    lr=5e-4,\n    betas=(0.5, 0.999)\n)\n\nl1 = nn.L1Loss()\nmse = nn.MSELoss()","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:53.016757Z","iopub.execute_input":"2022-04-28T10:42:53.017063Z","iopub.status.idle":"2022-04-28T10:42:56.001876Z","shell.execute_reply.started":"2022-04-28T10:42:53.017029Z","shell.execute_reply":"2022-04-28T10:42:56.001135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loader = DataLoader(dataset, batch_size=1, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:42:56.004247Z","iopub.execute_input":"2022-04-28T10:42:56.004726Z","iopub.status.idle":"2022-04-28T10:42:56.010089Z","shell.execute_reply.started":"2022-04-28T10:42:56.004686Z","shell.execute_reply":"2022-04-28T10:42:56.009065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs = 6\n\nfor epoch in range(num_epochs):\n    loop = tqdm(loader, leave=True)\n    for idx, (monet, photo) in enumerate(loop):\n        monet = monet.to(device)\n        photo = photo.to(device)\n\n        fake_photo = gen_P(monet)\n        D_P_real = disc_P(photo)\n        D_P_fake = disc_P(fake_photo.detach())\n        D_P_real_loss = mse(D_P_real, torch.ones(len(D_P_real)).to(device))\n        D_P_fake_loss = mse(D_P_fake, torch.zeros(len(D_P_fake)).to(device))\n        D_P_loss = D_P_real_loss + D_P_fake_loss\n\n        fake_monet = gen_M(photo)\n        D_M_real = disc_M(monet)\n        D_M_fake = disc_M(fake_monet.detach())\n        D_M_real_loss = mse(D_M_real, torch.ones(len(D_M_real)).to(device))\n        D_M_fake_loss = mse(D_M_fake, torch.zeros(len(D_M_fake)).to(device))\n        D_M_loss = D_M_real_loss + D_M_fake_loss\n\n        D_loss = (D_P_loss + D_M_loss) / 2\n        opt_disc.zero_grad()\n        D_loss.backward()\n        opt_disc.step()\n\n        D_P_fake = disc_P(fake_photo)\n        D_M_fake = disc_M(fake_monet)\n        loss_G_P = mse(D_P_fake, torch.ones(len(D_P_fake)).to(device))\n        loss_G_M = mse(D_M_fake, torch.ones(len(D_M_fake)).to(device))\n\n        cycle_photo = gen_P(fake_monet)\n        cycle_monet = gen_M(fake_photo)\n        cycle_photo_loss = l1(photo, cycle_photo)\n        cycle_monet_loss = l1(monet, cycle_monet)\n\n        identity_photo = gen_P(photo)\n        identity_monet = gen_M(monet)\n        identity_photo_loss = l1(photo, identity_photo)\n        identity_monet_loss = l1(monet, identity_monet)\n\n        G_loss = (\n            loss_G_P \n            + loss_G_M \n            + cycle_photo_loss * 10\n            + cycle_monet_loss * 10\n            + identity_photo_loss\n            + identity_monet_loss\n        )\n\n        opt_gen.zero_grad()\n        G_loss.backward()\n        opt_gen.step()\n        \n        if idx % 200 == 0:\n            print(\n                'G_P_loss:', loss_G_P.cpu().detach().item(), \n                'G_M_loss:', loss_G_M.cpu().detach().item(), \n                'D_M_loss:', D_M_loss.cpu().detach().item(), \n                'D_P_loss:', D_P_loss.cpu().detach().item()\n            )","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2022-04-28T10:49:39.528254Z","iopub.execute_input":"2022-04-28T10:49:39.528807Z","iopub.status.idle":"2022-04-28T10:52:46.555043Z","shell.execute_reply.started":"2022-04-28T10:49:39.528769Z","shell.execute_reply":"2022-04-28T10:52:46.553843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for idx, (monet, photo) in enumerate(loader):\n    monet = monet.to(device)\n    photo = photo.to(device)\n    fake_photo = gen_P(monet)\n    fake_monet = gen_M(photo).cpu().detach()[0]\n    \n    plt.subplot(121)\n    plt.imshow(photo.cpu().detach()[0].squeeze().permute(1, 2, 0) * 0.5 + 0.5)\n    plt.subplot(122)\n    plt.imshow(fake_monet.squeeze().permute(1, 2, 0) * 0.5 + 0.5)\n        \n    plt.show()\n        \n    if idx == 4:\n        break","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:43:43.829902Z","iopub.status.idle":"2022-04-28T10:43:43.830413Z","shell.execute_reply.started":"2022-04-28T10:43:43.830159Z","shell.execute_reply":"2022-04-28T10:43:43.830184Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import PIL\n! mkdir ./images","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:43:43.832119Z","iopub.status.idle":"2022-04-28T10:43:43.832527Z","shell.execute_reply.started":"2022-04-28T10:43:43.832302Z","shell.execute_reply":"2022-04-28T10:43:43.832339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import PIL\ni = 1\nfor (monet, photo) in tqdm(dataset):\n    \n    photo = photo.to(device)\n    prediction = gen_M(photo[None, :]).cpu().detach()[0]\n    prediction = (127.5 + 127.5 * prediction).permute(1, 2, 0).numpy().astype(np.uint8)\n\n    im = PIL.Image.fromarray(prediction)\n    im.save(\"./images/\" + str(i) + \".jpg\")\n    i += 1\n    if i == 7030:\n        break","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:43:43.83399Z","iopub.status.idle":"2022-04-28T10:43:43.834407Z","shell.execute_reply.started":"2022-04-28T10:43:43.834175Z","shell.execute_reply":"2022-04-28T10:43:43.834196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive('/kaggle/working/images/', 'zip', '/kaggle/working/images/')","metadata":{"execution":{"iopub.status.busy":"2022-04-28T10:43:43.835559Z","iopub.status.idle":"2022-04-28T10:43:43.836189Z","shell.execute_reply.started":"2022-04-28T10:43:43.835943Z","shell.execute_reply":"2022-04-28T10:43:43.835972Z"},"trusted":true},"execution_count":null,"outputs":[]}]}