{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.11"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":627400,"sourceType":"datasetVersion","datasetId":308470}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":25853.762433,"end_time":"2025-05-19T01:02:40.722207","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-05-18T17:51:46.959774","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### 1. Dataset","metadata":{"papermill":{"duration":0.004458,"end_time":"2025-05-18T17:51:51.152417","exception":false,"start_time":"2025-05-18T17:51:51.147959","status":"completed"},"tags":[]}},{"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, raw_transform=None, style_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.raw_transform = raw_transform\n        self.style_transform = style_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[item % len(self.true_images_path)])\n        true_image = cv2.cvtColor(true_image, cv2.COLOR_BGR2RGB)\n        true_image = Image.fromarray(true_image)\n\n        if self.raw_transform:\n            raw_image = self.raw_transform(raw_image)\n        if self.style_transform:\n            true_image = self.style_transform(true_image)\n\n        return raw_image, true_image\n\nif __name__ == '__main__':\n    data = MyDataset(raw_root='/kaggle/input/selfie2anime/trainA',\n                     true_root =\"/kaggle/input/selfie2anime/trainB\",\n                     raw_transform=Compose([\n                         ToTensor(),\n                        Normalize([0.5,0.5,0.5],[0.5,0.5,0.5]),\n                     ]),\n                     style_transform=Compose([\n                         ToTensor(),\n                         Normalize([0.5,0.5,0.5], [0.5,0.5,0.5]),\n                     ])\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), cmap='gray')\n            ax[value + 1].imshow(j[value].astype(np.uint8))\n            ax[value].axis('off')\n            ax[value + 1].axis('off')\n\n        plt.show()\n        break\n\n","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:06.01781Z","iopub.execute_input":"2025-05-26T17:07:06.018141Z","iopub.status.idle":"2025-05-26T17:07:17.291056Z","shell.execute_reply.started":"2025-05-26T17:07:06.018124Z","shell.execute_reply":"2025-05-26T17:07:17.28972Z"},"papermill":{"duration":12.67159,"end_time":"2025-05-18T17:52:03.827794","exception":false,"start_time":"2025-05-18T17:51:51.156204","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2. Multi scale dis","metadata":{"papermill":{"duration":0.01081,"end_time":"2025-05-18T17:52:03.849806","exception":false,"start_time":"2025-05-18T17:52:03.838996","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import torch.nn as nn\n\nclass CAM(nn.Module):\n    def __init__(self, in_channels, out_channels, activation):\n        super().__init__()\n        self.activation = activation\n        self.gap = nn.AdaptiveAvgPool2d(1)\n        self.gap_linear = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)\n        self.gmp = nn.AdaptiveMaxPool2d(1)\n        self.gmp_linear = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)\n        self.out = nn.Conv2d(in_channels * 2, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        gap = self.gap(x)\n        gap_y = self.gap_linear(gap)\n        gap_m = self.gap_linear(x)\n\n        gmp = self.gmp(x)\n        gmp_y = self.gmp_linear(gmp)\n        gmp_m = self.gmp_linear(x)\n\n        cam_logit = torch.cat([gap_y, gmp_y], dim=1)\n        x = torch.cat([gap_m, gmp_m], dim=1)\n        x = self.out(x)\n        return self.activation(x), cam_logit","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:17.292782Z","iopub.execute_input":"2025-05-26T17:07:17.29313Z","iopub.status.idle":"2025-05-26T17:07:17.299292Z","shell.execute_reply.started":"2025-05-26T17:07:17.293109Z","shell.execute_reply":"2025-05-26T17:07:17.298746Z"},"papermill":{"duration":0.020475,"end_time":"2025-05-18T17:52:03.881081","exception":false,"start_time":"2025-05-18T17:52:03.860606","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\nfrom torch.nn.utils import spectral_norm\n\nclass Dis(nn.Module):\n    def __init__(self, image_channel=3, feature_d=64):\n        super().__init__()\n        self.conv1 = nn.Sequential(\n            spectral_norm(nn.Conv2d(image_channel, feature_d, 4, 2, 1, bias=False)),\n            nn.LeakyReLU(0.2,True)\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.cam = CAM(in_channels=feature_d * 8,\n                       out_channels=feature_d * 8,\n                       activation=nn.LeakyReLU(0.2,True))\n        self.dis_out = spectral_norm(nn.Conv2d(feature_d * 8, 1, 4, 1, 1, bias=False))\n        \n        # self.conv5 = self.block_conv(feature_d * 4, feature_d * 8, 4, 2, 1)\n        # self.conv6 = self.block_conv(feature_d * 8, feature_d * 16, 4, 2, 1)\n        # self.conv7 = self.block_conv(feature_d * 16, feature_d * 32, 4, 1, 1)\n        # self.dis_out_2 = nn.Conv2d(feature_d * 32, 1, 4, 1, 1)\n\n    def block_conv(self, in_channel, out_channel, kernel, stride, padding):\n        return nn.Sequential(\n            spectral_norm(nn.Conv2d(in_channel,\n                      out_channel,\n                      kernel_size=kernel,\n                      stride=stride,\n                      padding=padding,\n                      bias=False)),\n            nn.LeakyReLU(0.2,True)\n        )\n\n    def forward(self, x):\n        x1 = self.conv1(x)\n        x2 = self.conv2(x1)\n        x3 = self.conv3(x2)\n        x4 = self.conv4(x3)\n        cam_map, logit = self.cam(x4)\n        dis_out = self.dis_out(cam_map)\n\n        return dis_out, logit\n\n        # x5 = self.conv5(x3)\n        # x6 = self.conv6(x5)\n        # x7 = self.conv7(x6)\n        # dis_out_2 = self.dis_out_2(x7)\n\n        # if return_feature:\n        #     layer_outputs = [x1, x2, x3, x4]\n        #     features = []\n        #     if feature_layers:\n        #         features = [feat for idx, feat in enumerate(layer_outputs, start=1) if idx in feature_layers]\n        #     return dis_out, features\n\nif __name__ == '__main__':\n    model = Dis(image_channel=3, feature_d=64)\n    print(summary(model, input_size=(32, 3, 256, 256)))\n\n","metadata":{"papermill":{"duration":0.923996,"end_time":"2025-05-18T17:52:04.816893","exception":false,"start_time":"2025-05-18T17:52:03.892897","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T17:07:17.299884Z","iopub.execute_input":"2025-05-26T17:07:17.300094Z","iopub.status.idle":"2025-05-26T17:07:18.305555Z","shell.execute_reply.started":"2025-05-26T17:07:17.300074Z","shell.execute_reply":"2025-05-26T17:07:18.304704Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3. Gen with 6 residual block","metadata":{"papermill":{"duration":0.011283,"end_time":"2025-05-18T17:52:04.841159","exception":false,"start_time":"2025-05-18T17:52:04.829876","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"#### Using 6 residual blocks in CycleGAN due to limited time and computational resources.\n#### Although the original images are 256x256, they are loaded at 128x128 resolution to speed up training.\n#### This approach results in approximately 3x faster training.","metadata":{"papermill":{"duration":0.010588,"end_time":"2025-05-18T17:52:04.8627","exception":false,"start_time":"2025-05-18T17:52:04.852112","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\nimport torch.onnx\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_channel, out_channel):\n        super().__init__()\n        self.conv1 = nn.Sequential(\n            spectral_norm(nn.Conv2d(in_channel, out_channel, 3, stride=1, padding=1, padding_mode='reflect',bias=False)),\n            nn.ReLU(True))\n\n        self.conv2 = nn.Sequential(\n            spectral_norm(nn.Conv2d(out_channel, out_channel, 3, stride=1, padding=1, padding_mode='reflect',bias=False)))\n\n    def forward(self, x):\n        x1 = self.conv1(x)\n        x2 = self.conv2(x1)\n        return x+x2\n\nclass Gen(nn.Module):\n    def __init__(self,image_channel,feature_g=64,num_resid_block=6):\n        super().__init__()\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) for _ in range(num_resid_block)]\n        )\n\n        self.cam = CAM(in_channels=feature_g * 4,\n                       out_channels=feature_g * 4,\n                       activation=nn.ReLU(True))\n\n        self.up1 = self.up_block_conv(feature_g*4,feature_g*2,3,1,1)\n        self.up2 = self.up_block_conv(feature_g*2,feature_g,3,1,1)\n        self.up3 = nn.Sequential(\n            spectral_norm(nn.Conv2d(feature_g,image_channel,7,1,3,padding_mode='reflect',bias=False)),\n            nn.Tanh())\n\n    def forward(self,x):\n        d1 = self.down1(x)\n        d2 = self.down2(d1)\n        d3 = self.down3(d2)\n        bottom = self.resid_block(d3)\n        cam_map, logit = self.cam(bottom)\n        u1 = self.up1(cam_map)\n        u2 = self.up2(u1)\n        u3 = self.up3(u2)\n        return u3, logit\n\n    def down_block_conv(self, in_channel, out_channel, kernel, stride, padding):\n        return nn.Sequential(\n            spectral_norm(nn.Conv2d(in_channel,out_channel,kernel,stride,padding,padding_mode='reflect',bias=False)),\n            nn.ReLU(True)\n        )\n\n    def up_block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.Upsample(scale_factor=2, mode=\"nearest\"),\n            spectral_norm(nn.Conv2d(in_channel,out_channel,kernel,stride,padding,padding_mode='reflect',bias=False)),\n            nn.ReLU(True)\n        )\n\nif __name__ == '__main__':\n    model = Gen(3)\n    print(summary(model,input_size=(32,3,256,256)))\n\n    # model.eval()\n    # dummy_input = (torch.rand(1, 3, 256, 256, dtype=torch.float32),)\n    # torch.onnx.export(model, dummy_input, \"model.onnx\", opset_version=11)\n\n","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.306424Z","iopub.execute_input":"2025-05-26T17:07:18.306691Z","iopub.status.idle":"2025-05-26T17:07:18.76984Z","shell.execute_reply.started":"2025-05-26T17:07:18.306653Z","shell.execute_reply":"2025-05-26T17:07:18.768981Z"},"papermill":{"duration":0.273009,"end_time":"2025-05-18T17:52:05.146285","exception":false,"start_time":"2025-05-18T17:52:04.873276","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 4. Func to train gen and dis  ","metadata":{"papermill":{"duration":0.011216,"end_time":"2025-05-18T17:52:05.224375","exception":false,"start_time":"2025-05-18T17:52:05.213159","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\ndef train_discriminator(discriminator=None,\n                        discriminator_fusion=None,\n                        true_data=None,\n                        gen_data=None,\n                        optimizer_d=None,\n                        criterion = nn.MSELoss(),\n                        device = 'cpu'):\n\n    if discriminator_fusion is None:\n        out_real = discriminator(true_data)\n        loss_real = criterion(out_real, torch.ones_like(out_real, device=device))\n        out_fake = discriminator(gen_data.detach())\n        loss_fake = criterion(out_fake, torch.zeros_like(out_fake, device=device))\n        loss_d = (loss_real + loss_fake) / 2\n\n        optimizer_d.zero_grad()\n        loss_d.backward()\n        optimizer_d.step()\n\n        return loss_d.item()\n\n    else:\n        out_real_1, out_real_2 = discriminator_fusion(true_data)\n        out_fake_1, out_fake_2 = discriminator_fusion(gen_data.detach())\n\n        loss_real_1 = criterion(out_real_1, torch.ones_like(out_real_1, device=device))\n        loss_fake_1 = criterion(out_fake_1, torch.zeros_like(out_fake_1, device=device))\n        loss_d1 = (loss_real_1 + loss_fake_1) \n\n        loss_real_2 = criterion(out_real_2, torch.ones_like(out_real_2, device=device))\n        loss_fake_2 = criterion(out_fake_2, torch.zeros_like(out_fake_2, device=device))\n        loss_d2 = (loss_real_2 + loss_fake_2) \n\n        loss_d = (loss_d1 + loss_d2) \n\n        optimizer_d.zero_grad()\n        loss_d.backward()\n        optimizer_d.step()\n\n        return loss_d.item()\n\n\ndef train_generator(\n    discriminator=None,\n    discriminator_global=None,\n    discriminator_fusion=None,\n    true_data=None,\n    gen_data=None,\n    cycle_data=None,\n    iden_data=None,\n    gen_cam=None,\n    iden_cam=None,\n    criterion=nn.MSELoss(),\n    loss_fn=nn.L1Loss(),\n    cam_criterion=nn.BCEWithLogitsLoss(),\n    adv_lambda=None,\n    cycle_lambda=None,\n    iden_lambda=None,\n    cam_lambda=None,\n    perceptual_loss_value=0.0,\n    device='cpu'):\n\n    if discriminator_global:\n        out_adv_patch = discriminator(gen_data)\n        loss_adv_patch = adv_lambda * criterion(out_adv_patch, torch.ones_like(out_adv_patch, device=device))\n\n        out_adv_global = discriminator_global(gen_data)\n        loss_adv_global = adv_lambda * criterion(out_adv_global, torch.ones_like(out_adv_global, device=device))\n\n        loss_adv = (loss_adv_patch + loss_adv_global) / 2\n\n    elif discriminator_fusion:\n        out_adv_1, out_adv_2 = discriminator_fusion(gen_data)\n        loss_adv_1 = adv_lambda * criterion(out_adv_1, torch.ones_like(out_adv_1, device=device))\n        loss_adv_2 = adv_lambda * criterion(out_adv_2, torch.ones_like(out_adv_2, device=device))\n\n        loss_adv = (loss_adv_1 + loss_adv_2)\n\n    else:\n        out_adv = discriminator(gen_data)\n        loss_adv = adv_lambda * criterion(out_adv, torch.ones_like(out_adv, device=device))\n\n    if cycle_data is not None:\n        loss_cycle = cycle_lambda * loss_fn(cycle_data,true_data)\n    else:\n        loss_cycle = torch.tensor(0.0)\n\n    if iden_data is not None:\n        loss_iden = iden_lambda * loss_fn(iden_data,true_data)\n    else:\n        loss_iden = torch.tensor(0.0)\n\n    if gen_cam is not None and iden_cam is not None:\n        loss_cam_real = cam_criterion(gen_cam, torch.ones_like(gen_cam, device=device)) \n        loss_cam_fake = cam_criterion(iden_cam, torch.zeros_like(iden_cam, device=device)) \n        loss_cam = cam_lambda * (loss_cam_real + loss_cam_fake)\n\n    total_loss = loss_adv + loss_cycle + loss_iden + loss_cam + perceptual_loss_value\n\n    return (total_loss,\n            total_loss.item(),\n            loss_adv.item(),\n            loss_cycle.item(),\n            loss_iden.item(),\n            loss_cam.item())","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.771962Z","iopub.execute_input":"2025-05-26T17:07:18.772192Z","iopub.status.idle":"2025-05-26T17:07:18.783311Z","shell.execute_reply.started":"2025-05-26T17:07:18.772174Z","shell.execute_reply":"2025-05-26T17:07:18.782669Z"},"papermill":{"duration":0.025372,"end_time":"2025-05-18T17:52:05.261094","exception":false,"start_time":"2025-05-18T17:52:05.235722","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 5. Buffer","metadata":{"papermill":{"duration":0.011567,"end_time":"2025-05-18T17:52:05.284113","exception":false,"start_time":"2025-05-18T17:52:05.272546","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import random\nimport torch\n\nclass ImageBuffer:\n    def __init__(self, max_size=None):\n        self.max_size = max_size\n        self.buffer = []\n\n    def push_and_pop(self, images):\n        return_images = []\n        for image in images:\n            image = torch.unsqueeze(image.data, 0)\n            if len(self.buffer) < self.max_size:\n                self.buffer.append(image)\n                return_images.append(image)\n            else:\n                if random.uniform(0, 1) > 0.5:\n                    idx = random.randint(0, self.max_size - 1)\n                    return_images.append(self.buffer[idx].clone())\n                    self.buffer[idx] = image\n                else:\n                    return_images.append(image)\n\n        return torch.cat(return_images, 0)\n","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.784Z","iopub.execute_input":"2025-05-26T17:07:18.784161Z","iopub.status.idle":"2025-05-26T17:07:18.801825Z","shell.execute_reply.started":"2025-05-26T17:07:18.784148Z","shell.execute_reply":"2025-05-26T17:07:18.801226Z"},"papermill":{"duration":0.019395,"end_time":"2025-05-18T17:52:05.314972","exception":false,"start_time":"2025-05-18T17:52:05.295577","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 6. Callback for gamma ( if using feature extraction from dis ) and lr ","metadata":{"papermill":{"duration":0.011275,"end_time":"2025-05-18T17:52:05.338129","exception":false,"start_time":"2025-05-18T17:52:05.326854","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import numpy as np\nfrom torch.optim.lr_scheduler import LambdaLR\n\n\ndef get_lambda_scheduler(epoch, total_epochs, decay_start_epoch, start_lambda, end_lambda, mode='linear'):\n    if epoch < decay_start_epoch:\n        return start_lambda\n\n    decay_epochs = total_epochs - decay_start_epoch\n    progress = (epoch - decay_start_epoch) / float(decay_epochs)\n\n    if mode == 'linear':\n        return max(start_lambda - (start_lambda - end_lambda) * progress, end_lambda)\n\n    elif mode == 'linear_reverse':\n        return min(start_lambda + (end_lambda - start_lambda) * progress, end_lambda)\n\n    elif mode == 'exp':\n        decay_rate = (end_lambda / start_lambda) ** (1 / decay_epochs)\n        return max(start_lambda * (decay_rate ** (epoch - decay_start_epoch)), end_lambda)\n\n    elif mode == 'sigmoid_reverse':\n        return start_lambda + (end_lambda - start_lambda) * (1 / (1 + np.exp(-10 * (progress - 0.5))))\n\n    else:\n        raise ValueError(\"Mode not found! Choose from ['linear', 'linear_reverse', 'exp', 'sigmoid_reverse']\")\n\n\ndef get_lr_scheduler(optimizer, total_epochs, decay_start_epoch, min_factor=0.0):\n    def lambda_rule(epoch):\n        if epoch < decay_start_epoch:\n            return 1.0\n        else:\n            decay_ratio = (epoch - decay_start_epoch) / float(total_epochs - decay_start_epoch)\n            return max(min_factor, 1.0 - decay_ratio * (1.0 - min_factor))\n\n    return LambdaLR(optimizer, lr_lambda=lambda_rule)\n","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.80243Z","iopub.execute_input":"2025-05-26T17:07:18.802613Z","iopub.status.idle":"2025-05-26T17:07:18.824629Z","shell.execute_reply.started":"2025-05-26T17:07:18.802599Z","shell.execute_reply":"2025-05-26T17:07:18.824096Z"},"papermill":{"duration":0.020285,"end_time":"2025-05-18T17:52:05.370324","exception":false,"start_time":"2025-05-18T17:52:05.350039","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 7. Preprocessing data","metadata":{"papermill":{"duration":0.011523,"end_time":"2025-05-18T17:52:05.393286","exception":false,"start_time":"2025-05-18T17:52:05.381763","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class RandomRotateResize:\n    def __init__(self, max_angle: float, resize_size: int):\n        self.max_angle = max_angle\n        self.resize_size = resize_size\n        self._flags = [\n            cv2.INTER_NEAREST,\n            cv2.INTER_LINEAR,\n            cv2.INTER_CUBIC,\n            cv2.INTER_AREA,\n            cv2.INTER_LANCZOS4,\n        ]\n\n    def __call__(self, img: Image.Image) -> Image.Image:\n        arr = np.array(img)\n        angle = random.uniform(-self.max_angle, self.max_angle)\n        h, w = arr.shape[:2]\n        M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)\n        flag_rot = random.choice(self._flags)\n        arr = cv2.warpAffine(arr, M, (w, h), flags=flag_rot)\n        h, w = arr.shape[:2]\n        if w < h:\n            new_w = self.resize_size\n            new_h = int(h * self.resize_size / w)\n        else:\n            new_h = self.resize_size\n            new_w = int(w * self.resize_size / h)\n        flag_rs = random.choice(self._flags)\n        arr = cv2.resize(arr, (new_w, new_h), interpolation=flag_rs)\n\n        return Image.fromarray(arr)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T17:07:18.825516Z","iopub.execute_input":"2025-05-26T17:07:18.826189Z","iopub.status.idle":"2025-05-26T17:07:18.858627Z","shell.execute_reply.started":"2025-05-26T17:07:18.82617Z","shell.execute_reply":"2025-05-26T17:07:18.858122Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\nfrom torch.utils.data import Subset\nfrom torchvision import transforms as T\nfrom torchvision.transforms import functional as F\nfrom torchvision.transforms import (Compose,\n                                    Resize,\n                                    RandomCrop,\n                                    CenterCrop,\n                                    RandomHorizontalFlip,\n                                    ColorJitter,\n                                    ToTensor,\n                                    Normalize)\n\ndef seed_all(seed: int = 42):\n    random.seed(seed)\n    torch.manual_seed(seed)\n    torch.backends.cudnn.benchmark = True\n\ndef split_data(data, ratio_train_size, val_size):\n    index_data = list(range(len(data)))\n    random.shuffle(index_data)\n    train_size = int(ratio_train_size * len(index_data))\n    train_index = index_data[:train_size]\n    val_index = index_data[train_size:train_size + val_size]\n\n    train = Subset(data, train_index)\n    val = Subset(data, val_index)\n    print(f'Training set size: {len(train)}')\n    print(f'Validation set size: {len(val)}')\n    return train, val\n\ndef split_data_indices(data_len, ratio_train_size, val_size):\n    indices = list(range(data_len))\n    random.shuffle(indices)\n    train_size = int(ratio_train_size * data_len)\n    train_idx = indices[:train_size]\n    val_idx = indices[train_size:train_size + val_size]\n    print(f'Training set size: {len(train_idx)}')\n    print(f'Validation set size: {len(val_idx)}')\n    return train_idx, val_idx\n\ndef get_transforms(\n    resize_size=158,\n    image_size=128,\n    normalize_mean=(0.5, 0.5, 0.5),\n    normalize_std=(0.5, 0.5, 0.5)\n):\n    \n    def aug_pipeline():\n        return T.Compose([\n            RandomRotateResize(max_angle=20, resize_size=resize_size),\n            T.RandomCrop(image_size),\n            T.RandomHorizontalFlip(p=0.5),\n            T.ColorJitter(\n                brightness=32/255.0,\n                contrast=(0.5, 1.5),\n                saturation=(0.5, 1.5),\n                hue=18/360.0\n            ),\n            T.ToTensor(),\n            T.Normalize(mean=normalize_mean, std=normalize_std),\n        ])\n\n    raw_transform   = aug_pipeline()\n    style_transform = aug_pipeline()\n    val_transform   = T.Compose([\n        T.Resize(resize_size),\n        T.CenterCrop(image_size),\n        T.ToTensor(),\n        T.Normalize(mean=normalize_mean, std=normalize_std),\n    ])\n\n    return raw_transform, style_transform, val_transform","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.859344Z","iopub.execute_input":"2025-05-26T17:07:18.859595Z","iopub.status.idle":"2025-05-26T17:07:18.879602Z","shell.execute_reply.started":"2025-05-26T17:07:18.859573Z","shell.execute_reply":"2025-05-26T17:07:18.878859Z"},"papermill":{"duration":0.022021,"end_time":"2025-05-18T17:52:05.427336","exception":false,"start_time":"2025-05-18T17:52:05.405315","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 8. Train","metadata":{"papermill":{"duration":0.010954,"end_time":"2025-05-18T17:52:05.449514","exception":false,"start_time":"2025-05-18T17:52:05.43856","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import os\nos.environ[\"TF_ENABLE_ONEDNN_OPTS\"] = \"0\"\nimport shutil\nimport torch\nimport itertools\nimport torchvision.utils\n#from Dis import Dis\n#from Dis_SN import Dis\n#from Gen_ori import Gen\n#from Gen_unet import Gen\n#from Gen_unet_modify import Gen\n#from Gen_resu_net import Gen\n#from Gen_resnet_freeze import Gen\n#from gen_resnet_freeze_vanilla import Gen\nfrom torch.utils.data import DataLoader\n#from dataset import MyDataset\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.tensorboard import SummaryWriter\nfrom argparse import ArgumentParser\n#from callback import get_lambda_scheduler,get_lr_scheduler\n#from util import *\n#from loss import PerceptualLoss, feature_matching_style_loss\n#from buffer import ImageBuffer\n#from base_model import train_discriminator, train_generator\n#from mix_real_fake_random import mix_with_random_fake_decay\n#from APA_loss import *\n\n\ndef get_arg():\n    parser = ArgumentParser(description='Train CycleGAN')\n    parser.add_argument('--epochs', '-e', type=int, default=100)\n    parser.add_argument('--batch_size', '-b', type=int, default=4)\n    parser.add_argument('--gen_learning_rate', '-glr', type=float, default=0.0001)\n    parser.add_argument('--dis_learning_rate', '-dlr', type=float, default=0.0003)\n    parser.add_argument('--ratio_train_size', '-rt', type=float, default=0.9)\n    parser.add_argument('--val_size', '-r', type=int, default=32)\n    parser.add_argument('--resize_image', '-rm', type=int, default=158)\n    parser.add_argument('--input_image', '-im', type=int, default=128)\n    parser.add_argument('--buffer_size', '-bs', type=int, default=50)\n    parser.add_argument('--adv_lambda', '-al', type=float, default=1)\n    parser.add_argument('--cycle_lambda', '-cyl', type=int, default=10)\n    parser.add_argument('--identity_lambda', '-il', type=float, default=5)\n    parser.add_argument('--cam_lambda', '-cl', type=int, default=1)\n    parser.add_argument('--perceptual_lambda', '-pl', type=int, default=1)\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    parser.add_argument('--raw_image_path', '-rp', type=str, default='/kaggle/input/selfie2anime/trainA')\n    parser.add_argument('--style_image_path', '-sp', type=str, default='/kaggle/input/selfie2anime/trainB')\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef train(arg):\n\n    seed_all(42)\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    raw_transform, style_transform, val_transform = get_transforms(resize_size=arg.resize_image,\n                                                                   image_size=arg.input_image)\n\n    train_dataset = MyDataset(raw_root=arg.raw_image_path,\n                                   true_root=arg.style_image_path,\n                                   raw_transform=raw_transform,\n                                   style_transform=style_transform)\n\n    val_dataset = MyDataset(raw_root=arg.raw_image_path,\n                                 true_root=arg.style_image_path,\n                                 raw_transform=val_transform,\n                                 style_transform=val_transform)\n\n    train_idx, val_idx = split_data_indices(len(train_dataset), arg.ratio_train_size, arg.val_size)\n    train = Subset(train_dataset, train_idx)\n    val = Subset(val_dataset, val_idx)\n\n    train_loader = DataLoader(train, batch_size=arg.batch_size, shuffle=True, num_workers=0)\n    val_loader = DataLoader(val, batch_size=arg.val_size, shuffle=False, num_workers=0)\n\n    buffer_a = ImageBuffer(arg.buffer_size)\n    buffer_b = ImageBuffer(arg.buffer_size)\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, nn.ConvTranspose2d)):\n            nn.init.normal_(m.weight, mean=0.0, std=0.02)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n            \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 = optim.Adam(itertools.chain(gen_a.parameters(),gen_b.parameters()),lr=arg.gen_learning_rate, betas=(0.5, 0.999))\n    op_da = optim.Adam(dis_a.parameters(), lr=arg.dis_learning_rate, betas=(0.5, 0.999))\n    op_db = optim.Adam(dis_b.parameters(), lr=arg.dis_learning_rate, betas=(0.5, 0.999))\n\n    scheduler_g = get_lr_scheduler(op_g, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n    scheduler_da = get_lr_scheduler(op_da, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n    scheduler_db = get_lr_scheduler(op_db, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n\n    if arg.checkpoint:\n        checkpoint = torch.load(arg.checkpoint,\n                                map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),\n                                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.load_state_dict(checkpoint['op_g'])\n        op_da.load_state_dict(checkpoint['op_da'])\n        op_db.load_state_dict(checkpoint['op_db'])\n        best_loss = checkpoint['best_loss']\n\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) // arg.batch_size\n    mse = nn.MSELoss()\n    bce = nn.BCEWithLogitsLoss()\n    # perceptual_loss = PerceptualLoss(content_weight=1,\n    #                                  use_l1=True,\n    #                                  content_layer='relu3_3',\n    #                                  device=device)\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        mean_loss_adv_a = 0.0\n        mean_loss_adv_b = 0.0\n        mean_loss_cycle_a = 0.0\n        mean_loss_cycle_b = 0.0\n        mean_loss_iden_a = 0.0\n        mean_loss_iden_b = 0.0\n        mean_loss_cam_a = 0.0\n        mean_loss_cam_b = 0.0\n        \n\n        with tqdm(train_loader, desc=f\"Epoch {epoch + 1}/{arg.epochs}\") as pbar:\n            for iteration, (raw_image, true_image) in enumerate(pbar):\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                out_g1_da = buffer_a.push_and_pop(out_g1.detach())\n                out_g2_db = buffer_b.push_and_pop(out_g2.detach())\n\n                loss_da = train_discriminator(discriminator_fusion=dis_a,\n                                              true_data=true_image,\n                                              gen_data=out_g1_da,\n                                              optimizer_d=op_da,\n                                              criterion=mse,\n                                              device=device)\n                mean_loss_da += loss_da\n        \n                loss_db = train_discriminator(discriminator_fusion=dis_b,\n                                              true_data=raw_image,\n                                              gen_data=out_g2_db,\n                                              optimizer_d=op_db,\n                                              criterion=mse,\n                                              device=device)\n                mean_loss_db += loss_db\n                \n\n                out_g1,out_g1_cam = gen_a(raw_image)\n                out_g2,out_g2_cam = gen_b(true_image)\n                cycle_a,_ = gen_a(out_g2)\n                cycle_b,_ = gen_b(out_g1)\n                iden_a,iden_a_cam = gen_a(true_image)\n                iden_b,iden_b_cam = gen_b(raw_image)\n\n                (total_loss_ga,\n                 loss_ga_final,\n                 loss_adv_a,\n                 loss_cycle_a,\n                 loss_iden_a,\n                 loss_cam_a) = train_generator(discriminator_fusion=dis_a,\n                                                true_data=true_image,\n                                                gen_data=out_g1,\n                                                cycle_data=cycle_a,\n                                                iden_data=iden_a,\n                                                gen_cam=out_g1_cam,\n                                                iden_cam=iden_a_cam,\n                                                criterion=mse,\n                                                cam_criterion=bce,\n                                                adv_lambda=arg.adv_lambda,\n                                                cycle_lambda=arg.cycle_lambda,\n                                                iden_lambda=arg.identity_lambda,\n                                                cam_lambda=arg.cam_lambda,\n                                                device=device)\n\n                mean_loss_ga += loss_ga_final\n                mean_loss_adv_a += loss_adv_a\n                mean_loss_cycle_a += loss_cycle_a\n                mean_loss_iden_a += loss_iden_a\n                mean_loss_cam_a += loss_cam_a\n\n                (total_loss_gb,\n                 loss_gb_final,\n                 loss_adv_b,\n                 loss_cycle_b,\n                 loss_iden_b,\n                 loss_cam_b) = train_generator(discriminator_fusion=dis_b,\n                                                true_data=raw_image,\n                                                gen_data=out_g2,\n                                                cycle_data=cycle_b,\n                                                iden_data=iden_b,\n                                                gen_cam=out_g2_cam,\n                                                iden_cam=iden_b_cam,\n                                                criterion=mse,\n                                                cam_criterion=bce,\n                                                adv_lambda=arg.adv_lambda,\n                                                cycle_lambda=arg.cycle_lambda,\n                                                iden_lambda=arg.identity_lambda,\n                                                cam_lambda=arg.cam_lambda,\n                                                device=device)\n\n                op_g.zero_grad()\n                loss_g = total_loss_ga + total_loss_gb\n                loss_g.backward()\n                op_g.step()\n\n                mean_loss_gb += loss_gb_final\n                mean_loss_adv_b += loss_adv_b\n                mean_loss_cycle_b += loss_cycle_b\n                mean_loss_iden_b += loss_iden_b\n                mean_loss_cam_b += loss_cam_b\n                \n\n            scheduler_g.step()\n            scheduler_da.step()\n            scheduler_db.step()\n\n            print(f'Epoch {epoch + 1}/{arg.epochs} | '\n                  f'Da - Loss: {mean_loss_da / train_iters:.4f} | '\n                  f'Ga - Loss: {mean_loss_ga / train_iters:.4f} | '\n                  f'Ga - main: {mean_loss_adv_a / train_iters:.4f} - cycle: {mean_loss_cycle_a / train_iters:.4f} - iden: {mean_loss_iden_a / train_iters:.4f} | '\n                  f'Db - Loss: {mean_loss_db / train_iters:.4f} | '\n                  f'Gb - Loss: {mean_loss_gb / train_iters:.4f} | '\n                  f'Gb - main: {mean_loss_adv_b / train_iters:.4f} - cycle: {mean_loss_cycle_b / train_iters:.4f} - iden: {mean_loss_iden_b / train_iters:.4f} ')\n            print(f'cam ga: {mean_loss_cam_a / train_iters:.4f} - cam gb: {mean_loss_cam_b / train_iters}')\n\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(true_image)\n                out_ga_cycle,_ = gen_a(out_gb)\n                out_gb_cycle,_ = gen_b(out_ga)\n                true_image = torchvision.utils.make_grid(true_image[:arg.val_size].cpu(), normalize=True)\n                gen_image_a = torchvision.utils.make_grid(out_ga[:arg.val_size].cpu(), normalize=True)\n                gen_image_b = torchvision.utils.make_grid(out_gb[:arg.val_size].cpu(), normalize=True)\n                raw_image = torchvision.utils.make_grid(raw_image[:arg.val_size].cpu(), normalize=True)\n                cycle_image_a = torchvision.utils.make_grid(out_ga_cycle[:arg.val_size].cpu(), normalize=True)\n                cycle_image_b = torchvision.utils.make_grid(out_gb_cycle[:arg.val_size].cpu(), normalize=True)\n\n                writer.add_image('true image', true_image, epoch)\n                writer.add_image('gen image a', gen_image_a, epoch)\n                writer.add_image('gen image b', gen_image_b, epoch)\n                writer.add_image('raw image', raw_image, epoch)\n                writer.add_image('cycle image a', cycle_image_a, epoch)\n                writer.add_image('cycle image b', cycle_image_b, epoch)\n\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                writer.add_scalar('loss adv ga', mean_loss_adv_a / train_iters, epoch)\n                writer.add_scalar('loss adv gb', mean_loss_adv_b / train_iters, epoch)\n                writer.add_scalar('loss cycle ga', mean_loss_cycle_a / train_iters, epoch)\n                writer.add_scalar('loss cycle gb', mean_loss_cycle_b / train_iters, epoch)\n                writer.add_scalar('loss iden ga', mean_loss_iden_a / train_iters, epoch)\n                writer.add_scalar('loss iden gb', mean_loss_iden_b / train_iters, epoch)\n                writer.add_scalar('cam ga', mean_loss_cam_a / train_iters, epoch)\n                writer.add_scalar('cam gb', mean_loss_cam_b / 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': op_g.state_dict(),\n            'op_da': op_da.state_dict(),\n            'op_db': op_db.state_dict(),\n            'best_loss': best_loss,\n        }\n        torch.save(checkpoint, '{}/last.pt'.format(arg.trained))\n        if mean_loss_ga/train_iters < best_loss:\n            best_loss = mean_loss_ga/train_iters\n            torch.save(checkpoint, '{}/best.pt'.format(arg.trained))\n\nif __name__ == '__main__':\n    arg = get_arg()\n    train(arg)","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:07:18.880381Z","iopub.execute_input":"2025-05-26T17:07:18.88084Z","iopub.status.idle":"2025-05-26T17:13:01.593547Z","shell.execute_reply.started":"2025-05-26T17:07:18.880824Z","shell.execute_reply":"2025-05-26T17:13:01.592563Z"},"papermill":{"duration":25799.824035,"end_time":"2025-05-19T01:02:05.284749","exception":false,"start_time":"2025-05-18T17:52:05.460714","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 9. Test ( after 20 epoch )","metadata":{"papermill":{"duration":3.220832,"end_time":"2025-05-19T01:02:11.752558","exception":false,"start_time":"2025-05-19T01:02:08.531726","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport torch\nfrom argparse import ArgumentParser\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 CycleGAN')\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--image','-i',type=str,default=None)\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\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    print(f\"Trained epoch: {checkpoint['epoch']}\")\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    image_size = 128\n\n    transform = Compose([\n        Resize(image_size+30),\n        CenterCrop(image_size),\n        ToTensor(),\n        Normalize(mean, std)\n    ])\n\n    data = MyDataset(raw_root='/kaggle/input/selfie2anime/testA',\n                     true_root=\"/kaggle/input/selfie2anime/testB\",\n                     raw_transform=transform,\n                     style_transform=transform)\n    loader = DataLoader(dataset=data, batch_size=16, shuffle=False)\n    \n    print(f'Test data size: {len(data)}')\n    \n    for batch_idx, (in_image, out_image) in enumerate(loader):\n        with torch.no_grad():\n            out_a,_ = gen_a(in_image)\n            out_a = out_a.detach().cpu()\n    \n        ori_image = (in_image.cpu().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    \n        num_images = ori_image.shape[0]\n    \n        for i in range(0, num_images, 5):\n            end = min(i + 5, num_images)\n            n = end - i  \n    \n            fig, ax = plt.subplots(1, n * 2, figsize=(n * 4, 4))\n            ax = np.array(ax).flatten()  \n    \n            for j in range(n):\n                ax[2 * j].imshow(ori_image[i + j].astype(np.uint8))\n                ax[2 * j + 1].imshow(gen_image_a[i + j].astype(np.uint8))\n                ax[2 * j].set_title('Original')\n                ax[2 * j + 1].set_title('Generated')\n                ax[2 * j].axis('off')\n                ax[2 * j + 1].axis('off')\n    \n            plt.tight_layout()\n            plt.show()\n","metadata":{"execution":{"iopub.status.busy":"2025-05-26T17:13:06.789182Z","iopub.execute_input":"2025-05-26T17:13:06.789454Z","iopub.status.idle":"2025-05-26T17:13:40.186687Z","shell.execute_reply.started":"2025-05-26T17:13:06.789434Z","shell.execute_reply":"2025-05-26T17:13:40.185973Z"},"papermill":{"duration":19.261419,"end_time":"2025-05-19T01:02:34.235373","exception":false,"start_time":"2025-05-19T01:02:14.973954","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}