{"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":"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils import data\nimport torchvision\nfrom torchvision import transforms\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nimport glob\nfrom PIL import Image\nimport itertools\n\n\nos.listdir('/kaggle/input/gan-getting-started')\nphotos_path = glob.glob('/kaggle/input/gan-getting-started/photo_jpg/*.jpg')\nmonets_path = glob.glob('/kaggle/input/gan-getting-started/monet_jpg/*.jpg')\nlen(photos_path)\nlen(monets_path)\ntransform = transforms.Compose([\n    transforms.ToTensor(),                        # 将图像转换为张量，并将取值范围归一化到(0, 1)之间\n    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 标准化图像，设置均值和方差为0.5\n])\n\nclass Photo2monet_Dataset(data.Dataset):\n    def __init__(self, imgs_path):\n        self.imgs_path = imgs_path\n\n    def __getitem__(self, index):\n        img_path = self.imgs_path[index]\n        pil_img = Image.open(img_path)\n        pil_img = transform(pil_img)\n        return pil_img\n    \n    def __len__(self):\n        return len(self.imgs_path)\n    \n    \nphotos_dataset = Photo2monet_Dataset(photos_path)\nmonets_dataset = Photo2monet_Dataset(monets_path)\nlen(photos_dataset), len(monets_dataset)              ","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.258375Z","iopub.execute_input":"2023-06-26T07:01:04.258706Z","iopub.status.idle":"2023-06-26T07:01:04.288013Z","shell.execute_reply.started":"2023-06-26T07:01:04.258682Z","shell.execute_reply":"2023-06-26T07:01:04.287152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BTACH_SIZE = 32                                   \n\nphotos_dl = torch.utils.data.DataLoader(\n    photos_dataset,\n    batch_size=BTACH_SIZE,\n    shuffle=True\n)\n\nmonets_dl = torch.utils.data.DataLoader(\n    monets_dataset,\n    batch_size=BTACH_SIZE,\n    shuffle=True\n)\n\nphotos_batch = next(iter(photos_dl))   \nmonets_batch = next(iter(monets_dl))\n\nphotos_batch.shape, monets_batch.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.289423Z","iopub.execute_input":"2023-06-26T07:01:04.289672Z","iopub.status.idle":"2023-06-26T07:01:04.578275Z","shell.execute_reply.started":"2023-06-26T07:01:04.289652Z","shell.execute_reply":"2023-06-26T07:01:04.577210Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Downsample(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(Downsample, self).__init__()\n        self.conv_relu = nn.Sequential(\n                            nn.Conv2d(in_channels, out_channels, \n                                      kernel_size=3, stride=2, padding=1),\n                            nn.LeakyReLU(inplace=True),\n            )\n        self.bn = nn.InstanceNorm2d(out_channels)\n    def forward(self, x, is_bn=True):\n        x = self.conv_relu(x)\n        if is_bn:\n            x = self.bn(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.579499Z","iopub.execute_input":"2023-06-26T07:01:04.579836Z","iopub.status.idle":"2023-06-26T07:01:04.585839Z","shell.execute_reply.started":"2023-06-26T07:01:04.579813Z","shell.execute_reply":"2023-06-26T07:01:04.584884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Upsample(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(Upsample, self).__init__()\n        self.upconv_relu = nn.Sequential(\n                               nn.Upsample(scale_factor=2, mode='bilinear'),\n                               nn.Conv2d(in_channels, out_channels, \n                                                  kernel_size=3,\n                                                  padding=1),\n                               nn.LeakyReLU(inplace=True)\n            )\n        self.bn = nn.InstanceNorm2d(out_channels)\n        \n    def forward(self, x, is_drop=False):\n        x = self.upconv_relu(x)\n        x = self.bn(x)\n        if is_drop:\n            x = F.dropout2d(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.588560Z","iopub.execute_input":"2023-06-26T07:01:04.589151Z","iopub.status.idle":"2023-06-26T07:01:04.600963Z","shell.execute_reply.started":"2023-06-26T07:01:04.589121Z","shell.execute_reply":"2023-06-26T07:01:04.600190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self):\n        super(Generator, self).__init__()\n        self.down1 = Downsample(3, 64)\n        self.down2 = Downsample(64, 128)\n        self.down3 = Downsample(128, 256)\n        self.down4 = Downsample(256, 512)\n        self.down5 = Downsample(512, 512)\n        self.down6 = Downsample(512, 512)\n        \n        self.up1 = Upsample(512, 512)\n        self.up2 = Upsample(1024, 512)\n        self.up3 = Upsample(1024, 256)\n        self.up4 = Upsample(512, 128)\n        self.up5 = Upsample(256, 64)\n\n        self.last = nn.ConvTranspose2d(128, 3, \n                                       kernel_size=3,\n                                       stride=2,\n                                       padding=1,\n                                       output_padding=1)\n        \n\n    def forward(self, x):\n        x1 = self.down1(x, is_bn=False)   # torch.Size([8, 64, 128, 128])\n        x2 = self.down2(x1)               # torch.Size([8, 128, 64, 64])\n        x3 = self.down3(x2)               # torch.Size([8, 256, 32, 32])\n        x4 = self.down4(x3)               # torch.Size([8, 512, 16, 16])\n        x5 = self.down5(x4)               # torch.Size([8, 512, 8, 8])\n        x6 = self.down6(x5)               # torch.Size([8, 512, 4, 4])\n        \n        x6 = self.up1(x6, is_drop=True)   # torch.Size([8, 512, 8, 8])\n        x6 = torch.cat([x5, x6], dim=1)   # torch.Size([8, 1024, 8, 8])\n        \n        x6 = self.up2(x6, is_drop=True)   # torch.Size([8, 512, 16, 16])\n        x6 = torch.cat([x4, x6], dim=1)   # torch.Size([8, 1024, 16, 16])\n        \n        x6 = self.up3(x6, is_drop=True)                         \n        x6 = torch.cat([x3, x6], dim=1)          \n        \n        x6 = self.up4(x6)                       \n        x6 = torch.cat([x2, x6], dim=1)         \n        \n        x6 = self.up5(x6)                        \n        x6 = torch.cat([x1, x6], dim=1)         \n        \n        x6 = torch.tanh(self.last(x6))           \n        return x6","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.602219Z","iopub.execute_input":"2023-06-26T07:01:04.603159Z","iopub.status.idle":"2023-06-26T07:01:04.623589Z","shell.execute_reply.started":"2023-06-26T07:01:04.603125Z","shell.execute_reply":"2023-06-26T07:01:04.622617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nclass Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n        self.down1 = Downsample(3, 64)             # 128\n        self.down2 = Downsample(64, 128)           # 64\n        self.down3 = Downsample(128, 256)           # 32\n        self.last = nn.Conv2d(256, 1, 3)\n\n    def forward(self, img):\n        x = self.down1(img)\n        x = F.dropout2d(self.down2(x))\n        x = F.dropout2d(self.down3(x))\n        x = torch.sigmoid(self.last(x))\n        return x","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.624964Z","iopub.execute_input":"2023-06-26T07:01:04.626139Z","iopub.status.idle":"2023-06-26T07:01:04.636955Z","shell.execute_reply.started":"2023-06-26T07:01:04.626098Z","shell.execute_reply":"2023-06-26T07:01:04.636141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torchvision.utils as vutils\nimport torch.optim as optim\nimport itertools\nfrom torch.optim.lr_scheduler import StepLR\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ngen_AB = Generator().to(device)\ngen_BA = Generator().to(device)\ndis_A = Discriminator().to(device)\ndis_B = Discriminator().to(device)\n\nbceloss_fn = nn.BCELoss()\nl1loss_fn = nn.L1Loss()\n\ngen_optimizer = optim.Adam(itertools.chain(gen_AB.parameters(), gen_BA.parameters()), \n                           lr=2e-4, betas=(0.5, 0.999), weight_decay=0.0001)\ndis_A_optimizer = optim.Adam(dis_A.parameters(), lr=1e-4, betas=(0.5, 0.999), weight_decay=0.0001)\ndis_B_optimizer = optim.Adam(dis_B.parameters(), lr=1e-4, betas=(0.5, 0.999), weight_decay=0.0001)\n\ngen_scheduler = StepLR(gen_optimizer, step_size=5, gamma=0.7)\ndis_A_scheduler = StepLR(dis_A_optimizer, step_size=5, gamma=0.7)\ndis_B_scheduler = StepLR(dis_B_optimizer, step_size=5, gamma=0.7)","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.638188Z","iopub.execute_input":"2023-06-26T07:01:04.638655Z","iopub.status.idle":"2023-06-26T07:01:04.927811Z","shell.execute_reply.started":"2023-06-26T07:01:04.638632Z","shell.execute_reply":"2023-06-26T07:01:04.927139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_input = (photos_batch[:4]).to(device)\ntest_input.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.929288Z","iopub.execute_input":"2023-06-26T07:01:04.929880Z","iopub.status.idle":"2023-06-26T07:01:04.936853Z","shell.execute_reply.started":"2023-06-26T07:01:04.929852Z","shell.execute_reply":"2023-06-26T07:01:04.935581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"D_loss = []\nG_loss = []\nnum_epochs= 30\nfrom tqdm import tqdm\n\n\nfor epoch in range(30):\n    D_epoch_loss= 0\n    G_epoch_loss= 0\n    with tqdm(total=len(photos_dl), desc=f\"Epoch {epoch+1}/{num_epochs}\", unit=\"step\") as pbar:\n        photos_list = list(photos_dl)\n        monets_list = list(monets_dl)\n    \n        for step, real_A in enumerate(monets_list):\n            real_B = monets_list[step] \n            real_A = real_A.to(device)\n            real_B = real_B.to(device)\n\n            with torch.no_grad():\n                \n                same_B = gen_AB(real_B)\n                identity_B_loss = l1loss_fn(same_B, real_B)\n                same_A = gen_BA(real_A)\n                identity_A_loss = l1loss_fn(same_A, real_A)\n\n           \n            gen_optimizer.zero_grad()\n\n            # GAN loss\n            fake_B = gen_AB(real_A)\n            D_pred_fake_B = dis_B(fake_B)\n            gan_loss_AB = bceloss_fn(D_pred_fake_B, torch.ones_like(D_pred_fake_B, device=device))\n\n            fake_A = gen_BA(real_B)\n            D_pred_fake_A = dis_A(fake_A)\n            gan_loss_BA = bceloss_fn(D_pred_fake_A, torch.ones_like(D_pred_fake_A, device=device))\n\n       \n            recovered_A = gen_BA(fake_B)\n            cycle_loss_ABA = l1loss_fn(recovered_A, real_A)\n\n            recovered_B = gen_AB(fake_A)\n            cycle_loss_BAB = l1loss_fn(recovered_B, real_B)\n\n  \n            g_loss = (identity_B_loss + identity_A_loss + gan_loss_AB + gan_loss_BA\n                     + cycle_loss_ABA + cycle_loss_BAB)\n\n            g_loss.backward()\n            gen_optimizer.step()\n\n           \n            dis_A_optimizer.zero_grad()\n            dis_A_real_output = dis_A(real_A)              \n            dis_A_real_loss = bceloss_fn(dis_A_real_output, torch.ones_like(dis_A_real_output, device=device))\n\n            dis_A_fake_output = dis_A(fake_A.detach())     \n            dis_A_fake_loss = bceloss_fn(dis_A_fake_output, torch.zeros_like(dis_A_fake_output, device=device))\n\n            dis_A_loss = (dis_A_real_loss + dis_A_fake_loss)*0.5\n\n            dis_A_loss.backward()\n            dis_A_optimizer.step()\n\n\n            # dis_B 训练\n            dis_B_optimizer.zero_grad()\n            dis_B_real_output = dis_B(real_B)              # 判别器输入真实图片\n            dis_B_real_loss = bceloss_fn(dis_B_real_output, torch.ones_like(dis_B_real_output, device=device))\n\n            dis_B_fake_output = dis_B(fake_B.detach())     # 判别器输入生成图片\n            dis_B_fake_loss = bceloss_fn(dis_B_fake_output, torch.zeros_like(dis_B_fake_output, device=device))\n\n            dis_B_loss = (dis_B_real_loss + dis_B_fake_loss)*0.5\n\n            dis_B_loss.backward()\n            dis_B_optimizer.step()\n            pbar.update(1)\n            pbar.set_postfix(D_loss=D_epoch_loss / (step+1), G_loss=G_epoch_loss / (step+1))\n\n\n            \n            with torch.no_grad():\n                D_epoch_loss += (dis_A_loss + dis_B_loss).item()\n                G_epoch_loss += g_loss.item()\n        dis_A_scheduler.step()\n        dis_B_scheduler.step()\n        \n    with torch.no_grad():        \n        D_epoch_loss /= len(monets_list)\n        G_epoch_loss /= len(monets_list)\n        D_loss.append(D_epoch_loss)\n        G_loss.append(G_epoch_loss)\n\n    \n    if epoch % 5 == 0:\n        print(\"Epoch:\", epoch, \"D_epoch_loss:\", D_epoch_loss, \"G_epoch_loss:\", G_epoch_loss)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-26T07:01:04.939955Z","iopub.execute_input":"2023-06-26T07:01:04.940202Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(gen_AB, 'monet.pt') ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\noutput_dir = \"/kaggle/working/output\"\nos.makedirs(output_dir, exist_ok=True)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ngen_AB = torch.load('monet.pt', map_location=device)\ngen_AB.to(device)\n\n\ndef generate_output_images(input_folder, output_dir):\n    gen_AB.eval()  # 设置生成器模型为评估模式\n    \n    # 获取输入文件夹中的所有图片文件路径\n    input_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if os.path.isfile(os.path.join(input_folder, f))]\n    \n    for i, input_file in enumerate(input_files):\n        input_image = Image.open(input_file).convert(\"RGB\")  # 打开并转换图片为RGB格式\n        input_tensor = torchvision.transforms.ToTensor()(input_image).unsqueeze(0).to(device)  # 转换为张量并添加批次维度\n        \n        with torch.no_grad():\n            output_image = gen_AB(input_tensor)\n        \n        output_path = os.path.join(output_dir, f\"output_{i}.jpg\")\n        torchvision.utils.save_image(output_image, output_path)\n        \n        print(f\"Generated and saved output image {i+1}/{len(input_files)}\")\n\n# 调用函数来生成输出图片\ninput_folder = \"/kaggle/input/gan-getting-started/photo_jpg\"  # 替换成你的输入文件夹路径\ngenerate_output_images(input_folder, output_dir)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}