{"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":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30887,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### 1. Dataset and DataLoader","metadata":{}},{"cell_type":"code","source":"import os\nimport random\nimport cv2\nfrom PIL import Image\nfrom torch.utils.data import Dataset,DataLoader\nfrom torchvision.transforms import ToTensor,Compose,Normalize,Resize\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nclass MyDataset(Dataset):\n    def __init__(self, raw_root,true_root, transform=None):\n        self.raw_images_path = [os.path.join(raw_root, i) for i in os.listdir(raw_root)]\n        self.true_images_path = [os.path.join(true_root,i) for i in os.listdir(true_root)]\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.raw_images_path)\n\n    def __getitem__(self, item):\n        raw_image = cv2.imread(self.raw_images_path[item])\n        raw_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)\n        raw_image = Image.fromarray(raw_image)\n\n        true_image = cv2.imread(self.true_images_path[random.randrange(0,len(self.true_images_path))])\n        true_image = cv2.cvtColor(true_image, cv2.COLOR_BGR2RGB)\n        true_image = Image.fromarray(true_image)\n        if (self.transform):\n            raw_image = self.transform(raw_image)\n            true_image = self.transform(true_image)\n\n        return raw_image, true_image\n\nif __name__ == '__main__':\n    data = MyDataset(raw_root='/kaggle/input/gan-getting-started/photo_jpg',\n                     true_root =\"/kaggle/input/gan-getting-started/monet_jpg\",\n                     transform=Compose([\n                         ToTensor(),\n                        Normalize([0.5,0.5,0.5],[0.5,0.5,0.5]),\n                     ]))\n    data_loader = DataLoader(dataset=data,\n                             batch_size=32,\n                             shuffle=True,\n                             num_workers=4)\n\n    for i, j in data_loader:\n        fix, ax = plt.subplots(2, 8, figsize=(15, 6))\n        ax = ax.flatten()\n        i = (i.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        j = (j.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        print(i.shape, j.shape)\n\n        for value in range(0, len(ax), 2):\n            ax[value].imshow(i[value].astype(np.uint8))\n            ax[value + 1].imshow(j[value].astype(np.uint8))\n            ax[value].axis('off')\n            ax[value + 1].axis('off')\n            ax[value].set_title('ori image')\n            ax[value+1].set_title('paint style')\n\n        plt.show()\n        break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-22T07:57:05.918631Z","iopub.execute_input":"2025-02-22T07:57:05.918820Z","iopub.status.idle":"2025-02-22T07:57:15.169856Z","shell.execute_reply.started":"2025-02-22T07:57:05.918802Z","shell.execute_reply":"2025-02-22T07:57:15.168708Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2. Generator","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\nimport matplotlib.pyplot as plt\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_channel, out_channel, stride=1):\n        super().__init__()\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, 3, stride=stride, padding=1, padding_mode='reflect'),\n            nn.InstanceNorm2d(out_channel,affine=True),\n            nn.ReLU())\n\n        self.conv2 = nn.Sequential(\n            nn.Conv2d(out_channel, out_channel, 3, stride=1, padding=1, padding_mode='reflect'),\n            nn.InstanceNorm2d(out_channel,affine=True))\n\n        self.shortcut = nn.Sequential()\n        if stride != 1 or in_channel != out_channel:\n            self.shortcut = nn.Sequential(\n                nn.Conv2d(in_channel, out_channel, 1, stride=stride),\n                nn.InstanceNorm2d(out_channel,affine=True))\n\n    def forward(self, x):\n        shortcut = x\n        noise = torch.normal(0, 0.02, size=x.shape, device=x.device) \n        x = x + noise\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x += shortcut\n        return x\n\nclass Gen(nn.Module):\n    def __init__(self,image_channel,feature_g=64,num_resid_block=9):\n        super().__init__()\n\n        self.down1 = self.down_block_conv(image_channel,feature_g,7,1,3)\n        self.down2 = self.down_block_conv(feature_g,feature_g*2,3,2,1)\n        self.down3 = self.down_block_conv(feature_g*2,feature_g*4,3,2,1)\n\n        self.resid_block = nn.Sequential(\n            *[ResidualBlock(feature_g * 4, feature_g * 4,1) for _ in range(num_resid_block)]\n        )\n\n        self.up1 = self.up_block_conv(feature_g*4,feature_g*2,3,2,1)\n        self.up2 = self.up_block_conv(feature_g*2,feature_g,3,2,1)\n        self.up3 = nn.Sequential(\n            nn.Conv2d(feature_g,image_channel,7,1,padding=3,padding_mode='reflect'),\n            nn.Tanh())\n\n    def forward(self,x):\n        noise = torch.normal(0, 0.02, size=x.shape, device=x.device)\n        x = x + noise\n        x = self.down1(x)\n        x = self.down2(x)\n        x = self.down3(x)\n        x = self.resid_block(x)\n        x = self.up1(x)\n        x = self.up2(x)\n        x = self.up3(x)\n        return x\n\n    def down_block_conv(self, in_channel, out_channel, kernel, stride, padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, kernel, stride, padding, padding_mode='reflect'),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.ReLU()\n            )\n\n    def up_block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_channel,out_channel,kernel,stride,padding,output_padding=1),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.ReLU()\n        )\n\nif __name__ == '__main__':\n    model = Gen(3)\n    print(summary(model,input_size=(32,3,256,256)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-22T08:18:08.351152Z","iopub.execute_input":"2025-02-22T08:18:08.351550Z","iopub.status.idle":"2025-02-22T08:18:08.748776Z","shell.execute_reply.started":"2025-02-22T08:18:08.351521Z","shell.execute_reply":"2025-02-22T08:18:08.747834Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3. Discriminator","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\n\nclass Dis(nn.Module):\n    def __init__(self,image_channel=3,feature_d=64):\n        super().__init__()\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(image_channel,feature_d,4,2,1),\n            nn.LeakyReLU(0.2)\n        )\n        self.conv2 = self.block_conv(feature_d,feature_d*2,4,2,1)\n        self.conv3 = self.block_conv(feature_d*2, feature_d*4, 4, 2, 1)\n        self.conv4 = self.block_conv(feature_d*4, feature_d*8, 4, 1, 1)\n        self.conv5 = nn.Sequential(\n            nn.Conv2d(feature_d*8,1,4,1,1)\n        )\n\n    def block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel,out_channel,kernel_size=kernel,stride=stride,padding=padding),\n            nn.InstanceNorm2d(out_channel),\n            nn.LeakyReLU(0.2)\n        )\n\n    def forward(self,x):\n        input_tensor = self.conv1(x)\n        x = self.conv2(input_tensor)\n        x = self.conv3(x)\n        x = self.conv4(x)\n        x = self.conv5(x)\n        return x\n\nif __name__ == '__main__':\n    model = Dis(image_channel=3,feature_d=64)\n    print(summary(model,input_size=(32,3,256,256)))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-22T08:18:14.489088Z","iopub.execute_input":"2025-02-22T08:18:14.489386Z","iopub.status.idle":"2025-02-22T08:18:14.565605Z","shell.execute_reply.started":"2025-02-22T08:18:14.489364Z","shell.execute_reply":"2025-02-22T08:18:14.564769Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3. Train","metadata":{}},{"cell_type":"code","source":"import os\nimport shutil\nimport torch\nimport torchvision.utils\n# from Dis import Dis\n# from Gen import Gen\nfrom torch.utils.data import DataLoader, Dataset, Subset\n# from dataset import MyDataset\nfrom torchvision.transforms import ToTensor, Resize, Normalize\nfrom torchvision.transforms.v2 import Compose\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.tensorboard import SummaryWriter\nfrom argparse import ArgumentParser\nimport random\n\n\ndef get_arg():\n    parser = ArgumentParser(description='Train')\n    parser.add_argument('--epochs', '-e', type=int, default=10)\n    parser.add_argument('--batch_size', '-b', type=int, default=1)\n    parser.add_argument('--learning_rate', '-lr', type=float, default=0.0002)\n    parser.add_argument('--ratio', '-r', type=float, default=0.9)\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--trained', '-t', type=str, default='trained')\n    parser.add_argument('--log', '-l', type=str, default='tensorboard')\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef train(arg):\n    arg = get_arg()\n\n    if not os.path.isdir(arg.trained):\n        os.mkdir(arg.trained)\n    if os.path.isdir(arg.log):\n        shutil.rmtree(arg.log)\n\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    writer = SummaryWriter(arg.log)\n\n    data = MyDataset(raw_root='/kaggle/input/gan-getting-started/photo_jpg',\n                     true_root=\"/kaggle/input/gan-getting-started/monet_jpg\",\n                     transform=Compose([\n                         ToTensor(),\n                         Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n                     ]))\n\n    random.seed(42)\n    index_data = [i for i in range(0, int(0.1 * len(data)))]\n    random.shuffle(index_data)\n\n    train_size = int(arg.ratio * len(index_data))\n    train_index = index_data[:train_size]\n    val_index = index_data[train_size:]\n    train = Subset(data, train_index)\n    val = Subset(data, val_index)\n    print(len(train))\n    print(len(val))\n    train_loader = DataLoader(train, batch_size=arg.batch_size, shuffle=True)\n    val_loader = DataLoader(val, batch_size=32, shuffle=True)\n\n    gen_a = Gen(3)\n    gen_b = Gen(3)\n    dis_a = Dis(3)\n    dis_b = Dis(3)\n\n    gen_a.to(device)\n    gen_b.to(device)\n    dis_a.to(device)\n    dis_b.to(device)\n\n    def init_weights(m):\n        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n            #nn.init.kaiming_normal_(m.weight)\n            nn.init.normal_(m.weight.data, 0.0, 0.02)\n\n    gen_a.apply(init_weights)\n    gen_b.apply(init_weights)\n    dis_a.apply(init_weights)\n    dis_b.apply(init_weights)\n\n    op_g_a = optim.Adam(gen_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_g_b = optim.Adam(gen_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_d_a = optim.Adam(dis_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_d_b = optim.Adam(dis_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n\n    #bce = nn.BCEWithLogitsLoss()\n    mse = nn.MSELoss()\n    l1 = nn.L1Loss()\n\n    if arg.checkpoint:\n        checkpoint = torch.load(arg.checkpoint,map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),weights_only=True)\n        start_epoch = checkpoint['epoch']\n        dis_a.load_state_dict(checkpoint['dis_a'])\n        dis_b.load_state_dict(checkpoint['dis_b'])\n        gen_a.load_state_dict(checkpoint['gen_a'])\n        gen_b.load_state_dict(checkpoint['gen_b'])\n        op_g_a.load_state_dict(checkpoint['op_g_a'])\n        op_g_b.load_state_dict(checkpoint['op_g_b'])\n        op_d_a.load_state_dict(checkpoint['op_d_a'])\n        op_d_b.load_state_dict(checkpoint['op_d_b'])\n        best_loss = checkpoint['best_loss']\n    else:\n        start_epoch = 0\n        best_loss = float('inf')\n\n    train_iters = len(train) // arg.batch_size\n    val_iters = len(val) // 32\n\n    for epoch in range(start_epoch, arg.epochs):\n        dis_a.train()\n        dis_b.train()\n        gen_a.train()\n        gen_b.train()\n\n        mean_loss_da = 0.0\n        mean_loss_db = 0.0\n        mean_loss_ga = 0.0\n        mean_loss_gb = 0.0\n\n        for raw_image, true_image in tqdm(train_loader):\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            out_g1 = gen_a(raw_image)\n            out_g2 = gen_b(true_image)\n\n            out_d1a = dis_a(true_image)\n            loss_d1a = mse(out_d1a,torch.ones_like(out_d1a)*0.9)\n            out_d2a = dis_a(out_g1.detach())\n            loss_d2a = mse(out_d2a, torch.zeros_like(out_d2a)+0.1)\n            loss_da = (loss_d1a + loss_d2a)/2\n            mean_loss_da += loss_da\n\n            op_d_a.zero_grad()\n            loss_da.backward()\n            op_d_a.step()\n\n            out_d1b = dis_b(raw_image)\n            loss_d1b = mse(out_d1b, torch.ones_like(out_d1b))\n            out_d2b = dis_b(out_g2.detach())\n            loss_d2b = mse(out_d2b, torch.zeros_like(out_d2b))\n            loss_db = (loss_d1b + loss_d2b)/2\n            mean_loss_db += loss_db\n\n            op_d_b.zero_grad()\n            loss_db.backward()\n            op_d_b.step()\n\n            out_da = dis_a(out_g1)\n            loss_ga = mse(out_da, torch.ones_like(out_da))\n\n            cycle_a = gen_a(out_g2.detach())\n            #iden_a = gen_a(true_image)\n            #loss_iden_a = l1(iden_a,true_image)\n            loss_cycle_a = l1(cycle_a,true_image)\n            loss_ga_final = loss_ga + 10 * loss_cycle_a #+ 5 * loss_iden_a\n            mean_loss_ga += loss_ga_final\n\n            op_g_a.zero_grad()\n            loss_ga_final.backward()\n            op_g_a.step()\n\n            out_db = dis_b(out_g2)\n            loss_gb = mse(out_db, torch.ones_like(out_db))\n\n            cycle_b = gen_b(out_g1.detach())\n            #iden_b = gen_b(raw_image)\n            loss_cycle_b = l1(cycle_b,raw_image)\n            #loss_iden_b = l1(iden_b,raw_image)\n            loss_gb_final = loss_gb + 10 * loss_cycle_b #+ 5 * loss_iden_b\n            mean_loss_gb += loss_gb_final\n\n            op_g_b.zero_grad()\n            loss_gb_final.backward()\n            op_g_b.step()\n\n        print('Epoch {}/{} -Da: Loss: {} | -Ga: Loss: {} | -Db: Loss: {} | -Gb: Loss: {}\\n'.format(\n                                                                epoch + 1,arg.epochs,\n                                                                mean_loss_da / train_iters,\n                                                                mean_loss_ga / train_iters,\n                                                                mean_loss_db / train_iters,\n                                                                mean_loss_gb / train_iters))\n\n        gen_a.eval()\n        gen_b.eval()\n        for raw_image, true_image in val_loader:\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            with torch.no_grad():\n                out_ga = gen_a(raw_image)\n                out_gb = gen_b(out_ga)\n                true_image = torchvision.utils.make_grid(true_image[:32], normalize=True)\n                gen_image = torchvision.utils.make_grid(out_ga[:32], normalize=True)\n                raw_image = torchvision.utils.make_grid(raw_image[:32], normalize=True)\n                cycle_image = torchvision.utils.make_grid(out_gb[:32], normalize=True)\n                writer.add_image('true image', true_image, epoch)\n                writer.add_image('gen image', gen_image, epoch)\n                writer.add_image('raw image', raw_image, epoch)\n                writer.add_image('cycle image', cycle_image, epoch)\n                writer.add_scalar('loss da', mean_loss_da / train_iters, epoch)\n                writer.add_scalar('loss ga', mean_loss_ga / train_iters, epoch)\n                writer.add_scalar('loss db', mean_loss_db / train_iters, epoch)\n                writer.add_scalar('loss gb', mean_loss_gb / train_iters, epoch)\n\n        checkpoint = {\n            'epoch': epoch + 1,\n            'gen_a': gen_a.state_dict(),\n            'dis_a': dis_a.state_dict(),\n            'gen_b': gen_b.state_dict(),\n            'dis_b': dis_b.state_dict(),\n            'op_g_a': op_g_a.state_dict(),\n            'op_d_a': op_d_a.state_dict(),\n            'op_g_b': op_g_b.state_dict(),\n            'op_d_b': op_d_b.state_dict(),\n            'best_loss': best_loss\n        }\n        torch.save(checkpoint, '{}/last.pt'.format(arg.trained))\n        if ((mean_loss_ga / train_iters) + (mean_loss_gb / train_iters))/2 < best_loss:\n            best_loss = ((mean_loss_ga / train_iters) + (mean_loss_gb / train_iters))/2\n            torch.save(checkpoint, '{}/best.pt'.format(arg.trained))\n\nif __name__ == '__main__':\n    arg = get_arg()\n    train(arg)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-22T08:37:51.554580Z","iopub.execute_input":"2025-02-22T08:37:51.554936Z","iopub.status.idle":"2025-02-22T08:49:08.991660Z","shell.execute_reply.started":"2025-02-22T08:37:51.554907Z","shell.execute_reply":"2025-02-22T08:49:08.990657Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 4. Test","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport torch\nfrom argparse import ArgumentParser\n#from Gen import Gen\nimport matplotlib.pyplot as plt\nfrom torchvision.transforms import Compose,Resize,Normalize,ToTensor\nfrom torch.utils.data import DataLoader\n\ndef get_arg():\n    parser = ArgumentParser(description='Test Cycle GAN')\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--image','-i',type=str,default=None)\n    parser.add_argument('--threshold','-t',type=float,default=0.0)\n    args, unknown = parser.parse_known_args()\n    return args\n    \nif __name__ == '__main__':\n    arg = get_arg()\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    gen_a = Gen(3)\n    gen_b = Gen(3)\n\n    checkpoint = torch.load('/kaggle/working/trained/last.pt',\n                            map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),\n                            weights_only=True)\n    gen_a.load_state_dict(checkpoint['gen_a'])\n    gen_b.load_state_dict(checkpoint['gen_b'])\n\n    mean = [0.5, 0.5, 0.5]\n    std = [0.5, 0.5, 0.5]\n\n    std = [s if s > 0 else 1e-6 for s in std]\n\n    transform = Compose([\n        ToTensor(),\n        Normalize(mean, std)\n    ])\n    #data = MyDataset(root=day2night_root, transform=transform)\n    data = MyDataset(raw_root='/kaggle/input/gan-getting-started/photo_jpg',\n                     true_root=\"/kaggle/input/gan-getting-started/monet_jpg\",\n                     transform=transform)\n    loader = DataLoader(dataset=data,batch_size=16,shuffle=True)\n    print(len(loader))\n\n    show_iter = 0\n    for in_image, out_image in loader:\n        show_iter += 1\n        fix, ax = plt.subplots(2, 9, figsize=(15, 6))\n        ax = ax.flatten()\n\n        with torch.no_grad():\n            out_a = gen_a(in_image).detach().cpu()\n            out_b = gen_b(out_a).detach().cpu()\n        \n        ori_image = (in_image.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        gen_image_a = (out_a.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        gen_image_b = (out_b.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        print(ori_image.shape,gen_image_a.shape,gen_image_b.shape)\n\n        for value in range(0, len(ax), 3):\n            ax[value].imshow(ori_image[value].astype(np.uint8))\n            ax[value + 1].imshow(gen_image_a[value].astype(np.uint8))\n            ax[value + 2].imshow(gen_image_b[value].astype(np.uint8))\n\n            ax[value].set_title('ori image')\n            ax[value + 1].set_title('paint image')\n            ax[value + 2].set_title('cycle image')\n            ax[value].axis('off')\n            ax[value + 1].axis('off')\n            ax[value + 2].axis('off')\n\n        plt.show()\n        if show_iter == 3:\n            break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-22T09:00:47.720607Z","iopub.execute_input":"2025-02-22T09:00:47.720938Z","iopub.status.idle":"2025-02-22T09:02:19.321626Z","shell.execute_reply.started":"2025-02-22T09:00:47.720915Z","shell.execute_reply":"2025-02-22T09:02:19.320697Z"}},"outputs":[],"execution_count":null}]}