{"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\nimport torch\nimport torch.nn as nn\nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:49:32.648121Z","iopub.execute_input":"2022-01-13T18:49:32.648454Z","iopub.status.idle":"2022-01-13T18:49:32.653147Z","shell.execute_reply.started":"2022-01-13T18:49:32.648426Z","shell.execute_reply":"2022-01-13T18:49:32.652251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Cycle GAN\n![Cycle-GAN](https://miro.medium.com/max/1838/0*S5gn5i6UhfyoRr9S.png)","metadata":{"execution":{"iopub.status.busy":"2021-08-16T13:44:55.51138Z","iopub.execute_input":"2021-08-16T13:44:55.511731Z","iopub.status.idle":"2021-08-16T13:44:56.261058Z","shell.execute_reply.started":"2021-08-16T13:44:55.511703Z","shell.execute_reply":"2021-08-16T13:44:56.259547Z"}}},{"cell_type":"markdown","source":"## Step 1. Define Generator","metadata":{"execution":{"iopub.status.busy":"2021-08-16T13:46:38.779785Z","iopub.execute_input":"2021-08-16T13:46:38.780264Z","iopub.status.idle":"2021-08-16T13:46:38.784886Z","shell.execute_reply.started":"2021-08-16T13:46:38.780222Z","shell.execute_reply":"2021-08-16T13:46:38.784065Z"}}},{"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, in_channels):\n        super(ResidualBlock, self).__init__()\n        self.block = nn.Sequential(\n            nn.ReflectionPad2d(1), # padding, keep the image size constant after next conv2d\n            nn.Conv2d(in_channels, in_channels, 3),\n            nn.InstanceNorm2d(in_channels),\n            nn.ReLU(inplace=True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_channels, in_channels, 3),\n            nn.InstanceNorm2d(in_channels)\n        )\n    \n    def forward(self, x):\n        return x + self.block(x)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:49:36.902913Z","iopub.execute_input":"2022-01-13T18:49:36.903289Z","iopub.status.idle":"2022-01-13T18:49:36.909678Z","shell.execute_reply.started":"2022-01-13T18:49:36.903256Z","shell.execute_reply":"2022-01-13T18:49:36.908726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class GeneratorResNet(nn.Module):\n    def __init__(self, in_channels, num_residual_blocks=9):\n        super(GeneratorResNet, self).__init__()\n        \n        # Inital Convolution  3*256*256 -> 64*256*256\n        out_channels=64\n        self.conv = nn.Sequential(\n            nn.ReflectionPad2d(in_channels), # padding, keep the image size constant after next conv2d\n            nn.Conv2d(in_channels, out_channels, 2*in_channels+1),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        )\n        \n        channels = out_channels\n        \n        # Downsampling   64*256*256 -> 128*128*128 -> 256*64*64\n        self.down = []\n        for _ in range(2):\n            out_channels = channels * 2\n            self.down += [\n                nn.Conv2d(channels, out_channels, 3, stride=2, padding=1),\n                nn.InstanceNorm2d(out_channels),\n                nn.ReLU(inplace=True),\n            ]\n            channels = out_channels\n        self.down = nn.Sequential(*self.down)\n        \n        # Transformation (ResNet)  256*64*64\n        self.trans = [ResidualBlock(channels) for _ in range(num_residual_blocks)]\n        self.trans = nn.Sequential(*self.trans)\n        \n        # Upsampling  256*64*64 -> 128*128*128 -> 64*256*256\n        self.up = []\n        for _ in range(2):\n            out_channels = channels // 2\n            self.up += [\n                nn.Upsample(scale_factor=2), # bilinear interpolation\n                nn.Conv2d(channels, out_channels, 3, stride=1, padding=1),\n                nn.InstanceNorm2d(out_channels),\n                nn.ReLU(inplace=True),\n            ]\n            channels = out_channels\n        self.up = nn.Sequential(*self.up)\n        \n        # Out layer  64*256*256 -> 3*256*256\n        self.out = nn.Sequential(\n            nn.ReflectionPad2d(in_channels),\n            nn.Conv2d(channels, in_channels, 2*in_channels+1),\n            nn.Tanh()\n        )\n    \n    def forward(self, x):\n        x = self.conv(x)\n        x = self.down(x)\n        x = self.trans(x)\n        x = self.up(x)\n        x = self.out(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:49:43.017236Z","iopub.execute_input":"2022-01-13T18:49:43.017555Z","iopub.status.idle":"2022-01-13T18:49:43.029755Z","shell.execute_reply.started":"2022-01-13T18:49:43.017525Z","shell.execute_reply":"2022-01-13T18:49:43.028932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 2. Define Discriminator","metadata":{"execution":{"iopub.status.busy":"2021-08-17T02:12:55.676304Z","iopub.execute_input":"2021-08-17T02:12:55.676929Z","iopub.status.idle":"2021-08-17T02:12:55.694692Z","shell.execute_reply.started":"2021-08-17T02:12:55.676889Z","shell.execute_reply":"2021-08-17T02:12:55.693976Z"}}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_channels):\n        super(Discriminator, self).__init__()\n        \n        self.model = nn.Sequential(\n            # why normalize=False?\n            *self.block(in_channels, 64, normalize=False), # 3*256*256 -> 64*128*128 \n            *self.block(64, 128),  # 64*128*128 -> 128*64*64\n            *self.block(128, 256), # 128*64*64 -> 256*32*32\n            *self.block(256, 512), # 256*32*32 -> 512*16*16\n            \n            # Why padding first then convolution?\n            nn.ZeroPad2d((1,0,1,0)), # padding left and top   512*16*16 -> 512*17*17\n            nn.Conv2d(512, 1, 4, padding=1) # 512*17*17 -> 1*16*16\n        )\n        \n        self.scale_factor = 16\n    \n    @staticmethod\n    def block(in_channels, out_channels, normalize=True):\n        layers = [nn.Conv2d(in_channels, out_channels, 4, stride=2, padding=1)]\n        if normalize:\n            layers.append(nn.InstanceNorm2d(out_channels))\n        layers.append(nn.LeakyReLU(0.2, inplace=True))\n        \n        return layers\n        \n    def forward(self, x):\n        return self.model(x)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:49:52.187102Z","iopub.execute_input":"2022-01-13T18:49:52.187430Z","iopub.status.idle":"2022-01-13T18:49:52.195286Z","shell.execute_reply.started":"2022-01-13T18:49:52.187400Z","shell.execute_reply":"2022-01-13T18:49:52.194152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 3. Define Loss","metadata":{"execution":{"iopub.status.busy":"2021-08-17T03:15:37.924281Z","iopub.execute_input":"2021-08-17T03:15:37.92481Z","iopub.status.idle":"2021-08-17T03:15:38.24994Z","shell.execute_reply.started":"2021-08-17T03:15:37.924769Z","shell.execute_reply":"2021-08-17T03:15:38.249301Z"}}},{"cell_type":"code","source":"criterion_GAN = nn.MSELoss()\ncriterion_cycle = nn.L1Loss()\ncriterion_identity = nn.L1Loss()","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:50:08.679375Z","iopub.execute_input":"2022-01-13T18:50:08.679684Z","iopub.status.idle":"2022-01-13T18:50:08.685563Z","shell.execute_reply.started":"2022-01-13T18:50:08.679655Z","shell.execute_reply":"2022-01-13T18:50:08.684688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 4. Initalize G and D","metadata":{}},{"cell_type":"code","source":"G_AB = GeneratorResNet(3, num_residual_blocks=9)\nD_B = Discriminator(3)\n\nG_BA = GeneratorResNet(3, num_residual_blocks=9)\nD_A = Discriminator(3)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:50:16.320260Z","iopub.execute_input":"2022-01-13T18:50:16.320566Z","iopub.status.idle":"2022-01-13T18:50:16.572821Z","shell.execute_reply.started":"2022-01-13T18:50:16.320539Z","shell.execute_reply":"2022-01-13T18:50:16.572002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cuda = torch.cuda.is_available()\nprint(f'cuda: {cuda}')\nif cuda:\n    G_AB = G_AB.cuda()\n    D_B = D_B.cuda()\n    G_BA = G_BA.cuda()\n    D_A = D_A.cuda()\n    \n    criterion_GAN = criterion_GAN.cuda()\n    criterion_cycle = criterion_cycle.cuda()\n    criterion_identity = criterion_identity.cuda()","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:50:19.717087Z","iopub.execute_input":"2022-01-13T18:50:19.717402Z","iopub.status.idle":"2022-01-13T18:50:24.072718Z","shell.execute_reply.started":"2022-01-13T18:50:19.717374Z","shell.execute_reply":"2022-01-13T18:50:24.071747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 5. Configure Optimizers","metadata":{"execution":{"iopub.status.busy":"2021-08-17T03:20:13.704545Z","iopub.execute_input":"2021-08-17T03:20:13.704889Z","iopub.status.idle":"2021-08-17T03:20:13.710944Z","shell.execute_reply.started":"2021-08-17T03:20:13.704859Z","shell.execute_reply":"2021-08-17T03:20:13.709978Z"}}},{"cell_type":"code","source":"import itertools\nlr = 0.0005\nb1 = 0.5\nb2 = 0.996\n\noptimizer_G = torch.optim.Adam(\n    itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=lr, betas=(b1, b2)\n)\n\noptimizer_D_A = torch.optim.Adam(\n    D_A.parameters(), lr=lr, betas=(b1, b2)\n)\n\noptimizer_D_B = torch.optim.Adam(\n    D_B.parameters(), lr=lr, betas=(b1, b2)\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:50:31.145624Z","iopub.execute_input":"2022-01-13T18:50:31.145947Z","iopub.status.idle":"2022-01-13T18:50:31.153445Z","shell.execute_reply.started":"2022-01-13T18:50:31.145917Z","shell.execute_reply":"2022-01-13T18:50:31.152621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 6. Learning Rate Scheduler Setting","metadata":{"execution":{"iopub.status.busy":"2021-08-17T03:26:53.38953Z","iopub.execute_input":"2021-08-17T03:26:53.389895Z","iopub.status.idle":"2021-08-17T03:26:53.409687Z","shell.execute_reply.started":"2021-08-17T03:26:53.389861Z","shell.execute_reply":"2021-08-17T03:26:53.408743Z"}}},{"cell_type":"code","source":"n_epoches = 120\ndecay_epoch = 20\n\nlambda_func = lambda epoch: 1 - max(0, epoch-decay_epoch)/(n_epoches-decay_epoch)\n\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=lambda_func)\nlr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda=lambda_func)\nlr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda=lambda_func)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:03.247353Z","iopub.execute_input":"2022-01-13T18:53:03.247691Z","iopub.status.idle":"2022-01-13T18:53:03.253243Z","shell.execute_reply.started":"2022-01-13T18:53:03.247660Z","shell.execute_reply":"2022-01-13T18:53:03.252265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 7. DataLoader","metadata":{"execution":{"iopub.status.busy":"2021-08-17T03:57:20.181606Z","iopub.execute_input":"2021-08-17T03:57:20.182081Z","iopub.status.idle":"2021-08-17T03:57:20.185286Z","shell.execute_reply.started":"2021-08-17T03:57:20.18205Z","shell.execute_reply":"2021-08-17T03:57:20.184648Z"}}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\n\nclass ImageDataset(Dataset):\n    def __init__(self, data_dir, mode='train', transforms=None):\n        A_dir = os.path.join(data_dir, 'monet_jpg')\n        B_dir = os.path.join(data_dir, 'photo_jpg')\n        \n        if mode == 'train':\n            self.files_A = [os.path.join(A_dir, name) for name in sorted(os.listdir(A_dir))[:250]]\n            self.files_B = [os.path.join(B_dir, name) for name in sorted(os.listdir(B_dir))[:250]]\n        elif mode == 'test':\n            self.files_A = [os.path.join(A_dir, name) for name in sorted(os.listdir(A_dir))[250:]]\n            self.files_B = [os.path.join(B_dir, name) for name in sorted(os.listdir(B_dir))[250:301]]\n        \n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.files_A)\n    \n    def __getitem__(self, index):\n        file_A = self.files_A[index]\n        file_B = self.files_B[index]\n        \n        img_A = Image.open(file_A)\n        img_B = Image.open(file_B)\n        \n        if self.transforms is not None:\n            img_A = self.transforms(img_A)\n            img_B = self.transforms(img_B)\n        \n        return img_A, img_B","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:07.281858Z","iopub.execute_input":"2022-01-13T18:53:07.282244Z","iopub.status.idle":"2022-01-13T18:53:07.291718Z","shell.execute_reply.started":"2022-01-13T18:53:07.282210Z","shell.execute_reply":"2022-01-13T18:53:07.290866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torchvision.transforms as transforms\n\ndata_dir = '/kaggle/input/gan-getting-started'\n\ntransforms_ = transforms.Compose([\n   # transforms.Resize(int(256*1.12), Image.BICUBIC),\n    #transforms.RandomCrop(256, 256),\n    transforms.RandomHorizontalFlip(),\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\nbatch_size = 5\n\ntrainloader = DataLoader(\n    ImageDataset(data_dir, mode='train', transforms=transforms_),\n    batch_size = batch_size,\n    shuffle = True,\n    num_workers = 3\n)\n\ntestloader = DataLoader(\n    ImageDataset(data_dir, mode='test', transforms=transforms_),\n    batch_size = batch_size,\n    shuffle = False,\n    num_workers = 3\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:09.293171Z","iopub.execute_input":"2022-01-13T18:53:09.293631Z","iopub.status.idle":"2022-01-13T18:53:09.323434Z","shell.execute_reply.started":"2022-01-13T18:53:09.293584Z","shell.execute_reply":"2022-01-13T18:53:09.322351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 8. Sample images to show","metadata":{"execution":{"iopub.status.busy":"2021-08-17T05:43:38.634944Z","iopub.execute_input":"2021-08-17T05:43:38.6353Z","iopub.status.idle":"2021-08-17T05:43:38.643016Z","shell.execute_reply.started":"2021-08-17T05:43:38.635271Z","shell.execute_reply":"2021-08-17T05:43:38.641759Z"}}},{"cell_type":"code","source":"from torchvision.utils import make_grid\n\nTensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n\ndef sample_images(real_A, real_B, figside=1.5):\n    assert real_A.size() == real_B.size(), 'The image size for two domains must be the same'\n    \n    G_AB.eval()\n    G_BA.eval()\n    \n    real_A = real_A.type(Tensor)\n    fake_B = G_AB(real_A).detach()\n    real_B = real_B.type(Tensor)\n    fake_A = G_BA(real_B).detach()\n    \n    nrows = real_A.size(0)\n    real_A = make_grid(real_A, nrow=nrows, normalize=True)\n    fake_B = make_grid(fake_B, nrow=nrows, normalize=True)\n    real_B = make_grid(real_B, nrow=nrows, normalize=True)\n    fake_A = make_grid(fake_A, nrow=nrows, normalize=True)\n    \n    image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1).cpu().permute(1, 2, 0)\n    \n    plt.figure(figsize=(figside*nrows, figside*4))\n    plt.imshow(image_grid)\n    plt.axis('off')\n    plt.show()\n    ","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:11.318818Z","iopub.execute_input":"2022-01-13T18:53:11.319172Z","iopub.status.idle":"2022-01-13T18:53:11.327655Z","shell.execute_reply.started":"2022-01-13T18:53:11.319142Z","shell.execute_reply":"2022-01-13T18:53:11.326484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"real_A, real_B = next(iter(testloader))\nsample_images(real_A, real_B)","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:15.979628Z","iopub.execute_input":"2022-01-13T18:53:15.979945Z","iopub.status.idle":"2022-01-13T18:53:16.921814Z","shell.execute_reply.started":"2022-01-13T18:53:15.979914Z","shell.execute_reply":"2022-01-13T18:53:16.921031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 9. Training","metadata":{}},{"cell_type":"code","source":"for epoch in range(n_epoches):\n    for i, (real_A, real_B) in enumerate(trainloader):\n        real_A, real_B = real_A.type(Tensor), real_B.type(Tensor)\n        \n        # groud truth\n        out_shape = [real_A.size(0), 1, real_A.size(2)//D_A.scale_factor, real_A.size(3)//D_A.scale_factor]\n        valid = torch.ones(out_shape).type(Tensor)\n        fake = torch.zeros(out_shape).type(Tensor)\n        \n        \"\"\"Train Generators\"\"\"\n        # set to training mode in the begining, beacause sample_images will set it to eval mode\n        G_AB.train()\n        G_BA.train()\n        \n        optimizer_G.zero_grad()\n        \n        fake_B = G_AB(real_A)\n        fake_A = G_BA(real_B)\n        \n        # identity loss\n        loss_id_A = criterion_identity(fake_B, real_A)\n        loss_id_B = criterion_identity(fake_A, real_B)\n        loss_identity = (loss_id_A + loss_id_B) / 2\n        \n        # GAN loss, train G to make D think it's true\n        loss_GAN_AB = criterion_GAN(D_B(fake_B), valid) \n        loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)\n        loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2\n        \n        # cycle loss\n        recov_A = G_BA(fake_B)\n        recov_B = G_AB(fake_A)\n        loss_cycle_A = criterion_cycle(recov_A, real_A)\n        loss_cycle_B = criterion_cycle(recov_B, real_B)\n        loss_cycle = (loss_cycle_A + loss_cycle_B) / 2\n        \n        # G totol loss\n        loss_G = 5.0*loss_identity + loss_GAN + 10.0*loss_cycle\n        \n        loss_G.backward()\n        optimizer_G.step()\n        \n        \"\"\"Train Discriminator A\"\"\"\n        optimizer_D_A.zero_grad()\n        \n        loss_real = criterion_GAN(D_A(real_A), valid)\n        loss_fake = criterion_GAN(D_A(fake_A.detach()), fake)\n        loss_D_A = (loss_real + loss_fake) / 2\n        \n        loss_D_A.backward()\n        optimizer_D_A.step()\n        \n        \"\"\"Train Discriminator B\"\"\"\n        optimizer_D_B.zero_grad()\n        \n        loss_real = criterion_GAN(D_B(real_B), valid)\n        loss_fake = criterion_GAN(D_B(fake_B.detach()), fake)\n        loss_D_B = (loss_real + loss_fake) / 2\n        \n        loss_D_B.backward()\n        optimizer_D_B.step()\n    \n    lr_scheduler_G.step()\n    lr_scheduler_D_A.step()\n    lr_scheduler_D_B.step()\n    \n    # test\n    if (epoch+1) % 10 == 0:\n        test_real_A, test_real_B = next(iter(testloader))\n        sample_images(test_real_A, test_real_B)\n\n        loss_D = (loss_D_A + loss_D_B) / 2\n        print(f'[Epoch {epoch+1}/{n_epoches}]')\n        print(f'[G loss: {loss_G.item()} | identity: {loss_identity.item()} GAN: {loss_GAN.item()} cycle: {loss_cycle.item()}]')\n        print(f'[D loss: {loss_D.item()} | D_A: {loss_D_A.item()} D_B: {loss_D_B.item()}]')    ","metadata":{"execution":{"iopub.status.busy":"2022-01-13T18:53:25.179563Z","iopub.execute_input":"2022-01-13T18:53:25.179897Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Step 10. Generate Images","metadata":{}},{"cell_type":"code","source":"photo_dir = os.path.join(data_dir, 'photo_jpg')\nfiles = [os.path.join(photo_dir, name) for name in os.listdir(photo_dir)]\nlen(files)","metadata":{"execution":{"iopub.status.busy":"2021-08-17T16:13:55.679376Z","iopub.execute_input":"2021-08-17T16:13:55.679792Z","iopub.status.idle":"2021-08-17T16:13:55.702252Z","shell.execute_reply.started":"2021-08-17T16:13:55.679748Z","shell.execute_reply":"2021-08-17T16:13:55.701317Z"},"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":"2021-08-17T16:13:55.703418Z","iopub.execute_input":"2021-08-17T16:13:55.703778Z","iopub.status.idle":"2021-08-17T16:13:55.708333Z","shell.execute_reply.started":"2021-08-17T16:13:55.703743Z","shell.execute_reply":"2021-08-17T16:13:55.707205Z"},"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\nG_BA.eval()\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 = G_BA(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":"2021-08-17T16:13:55.709618Z","iopub.execute_input":"2021-08-17T16:13:55.709978Z","iopub.status.idle":"2021-08-17T16:18:14.669934Z","shell.execute_reply.started":"2021-08-17T16:13:55.709942Z","shell.execute_reply":"2021-08-17T16:18:14.66908Z"},"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":"2021-08-17T16:18:14.671164Z","iopub.execute_input":"2021-08-17T16:18:14.671525Z","iopub.status.idle":"2021-08-17T16:18:17.508148Z","shell.execute_reply.started":"2021-08-17T16:18:14.671459Z","shell.execute_reply":"2021-08-17T16:18:17.507156Z"},"trusted":true},"execution_count":null,"outputs":[]}]}