{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11688381,"sourceType":"datasetVersion","datasetId":7336173}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":40058.097693,"end_time":"2025-05-06T03:00:02.386684","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-05-05T15:52:24.288991","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"19223935","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/gan-getting-started/photo_jpg',\n                     true_root =\"/kaggle/input/arcylic-painting/arcylic_painting\",\n                     raw_transform=Compose([\n                         Resize((256,256)),\n                         ToTensor(),\n                        Normalize([0.5,0.5,0.5],[0.5,0.5,0.5]),\n                     ]),\n                     style_transform=Compose([\n                         Resize((256,256)),\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].set_title('A')\n            ax[value + 1].axis('off')\n            ax[value + 1].set_title('B')\n\n        plt.show()\n        break\n\n","metadata":{"execution":{"iopub.status.busy":"2025-09-11T19:52:15.095025Z","iopub.execute_input":"2025-09-11T19:52:15.095306Z","iopub.status.idle":"2025-09-11T19:52:18.174379Z","shell.execute_reply.started":"2025-09-11T19:52:15.095284Z","shell.execute_reply":"2025-09-11T19:52:18.173469Z"},"papermill":{"duration":11.194079,"end_time":"2025-05-05T15:52:39.565605","exception":false,"start_time":"2025-05-05T15:52:28.371526","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"9dbfde3f","cell_type":"code","source":"import 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, noise_std=0.2):\n        super().__init__()\n        self.noise_std = noise_std\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(image_channel, feature_d, 4, 2, 1),\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.dis_out = nn.Conv2d(feature_d * 8, 1, 4, 1, 1)\n\n    def block_conv(self, in_channel, out_channel, kernel, stride, padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel,\n                      out_channel,\n                      kernel_size=kernel,\n                      stride=stride,\n                      padding=padding,\n                      bias=False),\n            nn.InstanceNorm2d(out_channel),\n            nn.LeakyReLU(0.2,True)\n        )\n\n    def forward(self, x, return_feature=False, feature_layers=None):\n        x1 = self.conv1(x)\n        x2 = self.conv2(x1)\n        x3 = self.conv3(x2)\n        x4 = self.conv4(x3)\n        dis_out = self.dis_out(x4)\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        return dis_out\n\nif __name__ == '__main__':\n    model = Dis(image_channel=3, feature_d=64)\n    print(summary(model, input_size=(32, 3, 256, 256)))","metadata":{"execution":{"iopub.execute_input":"2025-05-05T15:52:39.591633Z","iopub.status.busy":"2025-05-05T15:52:39.591330Z","iopub.status.idle":"2025-05-05T15:52:40.459075Z","shell.execute_reply":"2025-05-05T15:52:40.458212Z"},"papermill":{"duration":0.882624,"end_time":"2025-05-05T15:52:40.460575","exception":false,"start_time":"2025-05-05T15:52:39.577951","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"2d6c5cae","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, 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',bias=False),\n            nn.InstanceNorm2d(out_channel),\n            nn.ReLU(True))\n\n        self.conv2 = nn.Sequential(\n            nn.Conv2d(out_channel, out_channel, 3, stride=1, padding=1, padding_mode='reflect',bias=False),\n            nn.InstanceNorm2d(out_channel))\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))\n\n    def forward(self, x):\n        x1 = self.conv1(x)\n        x2 = self.conv2(x1)\n        return x+x2\n\nclass GaussianNoise(nn.Module):\n    def __init__(self, std=0.2):\n        super(GaussianNoise, self).__init__()\n        self.std = std\n\n    def forward(self, x):\n        if self.training:\n            noise = torch.randn_like(x) * self.std\n            return x + noise\n        return x\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,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,3,padding_mode='reflect'),\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        u1 = self.up1(bottom)\n        u2 = self.up2(u1)\n        u3 = self.up3(u2)\n        return u3\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',bias=False),\n            nn.InstanceNorm2d(out_channel),\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.ConvTranspose2d(in_channel,out_channel,kernel,stride,padding,output_padding=1,bias=False),\n            nn.InstanceNorm2d(out_channel),\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.execute_input":"2025-05-05T15:52:40.488301Z","iopub.status.busy":"2025-05-05T15:52:40.487829Z","iopub.status.idle":"2025-05-05T15:52:40.945835Z","shell.execute_reply":"2025-05-05T15:52:40.944931Z"},"papermill":{"duration":0.472587,"end_time":"2025-05-05T15:52:40.947130","exception":false,"start_time":"2025-05-05T15:52:40.474543","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"9bc27e5f","cell_type":"code","source":"def reconstruction_loss_with_feature(\n    dis_x,\n    real_x,\n    gen_x,\n    gamma=None,\n    loss_type=\"l1\",\n    feature_mode=False,\n    feature_pos=None,\n    weights=None\n):\n    if loss_type == \"l1\":\n        loss_fn = nn.L1Loss()\n    elif loss_type == \"l2\":\n        loss_fn = nn.MSELoss()\n    else:\n        raise ValueError(\"loss_type must be either 'l1' or 'l2'\")\n\n    if feature_mode and gamma == 1.0:\n        with torch.no_grad():\n            _, real_feats = dis_x(real_x, return_feature=True, feature_layers=feature_pos)\n            _, gen_feats  = dis_x(gen_x, return_feature=True, feature_layers=feature_pos)\n\n        feature_loss = 0.0\n        if weights is None:\n            weights = [1.0] * len(real_feats)\n\n        for w, f_real, f_gen in zip(weights, real_feats, gen_feats):\n            feature_loss += w * loss_fn(f_gen, f_real)\n\n        return feature_loss / sum(weights)\n\n    if not feature_mode or gamma == 0.0:\n        recon_loss = loss_fn(gen_x, real_x)\n        return recon_loss\n\n    recon_loss = loss_fn(gen_x, real_x)\n    with torch.no_grad():\n        _, real_feats = dis_x(real_x, return_feature=True, feature_layers=feature_pos)\n        _, gen_feats  = dis_x(gen_x, return_feature=True, feature_layers=feature_pos)\n\n    feature_loss = 0.0\n    if weights is None:\n        weights = [1.0] * len(real_feats)\n\n    for w, f_real, f_gen in zip(weights, real_feats, gen_feats):\n        feature_loss += w * loss_fn(f_gen, f_real)\n    feature_loss /= sum(weights)\n\n    return gamma * feature_loss + (1 - gamma) * recon_loss","metadata":{"execution":{"iopub.execute_input":"2025-05-05T15:52:40.973114Z","iopub.status.busy":"2025-05-05T15:52:40.972615Z","iopub.status.idle":"2025-05-05T15:52:40.979727Z","shell.execute_reply":"2025-05-05T15:52:40.979185Z"},"papermill":{"duration":0.020917,"end_time":"2025-05-05T15:52:40.980714","exception":false,"start_time":"2025-05-05T15:52:40.959797","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"134f5274","cell_type":"code","source":"class PerceptualLoss(nn.Module):\n    def __init__(self, device, content_weight=1, style_weight=0, use_l1=True, content_layer='relu2_2'):\n        super(PerceptualLoss, self).__init__()\n        vgg = vgg16(pretrained=True).features.to(device).eval()\n        self.layers = {\n            '3': 'relu1_2',\n            '8': 'relu2_2',\n            '15': 'relu3_3',\n            '22': 'relu4_3',\n        }\n        self.feature_extractor = vgg\n        for param in self.feature_extractor.parameters():\n            param.requires_grad = False\n\n        if use_l1:\n            self.loss_fn = nn.L1Loss()\n        else:\n            self.loss_fn = nn.MSELoss()\n\n        self.content_weight = content_weight\n        self.style_weight = style_weight\n        self.content_layer = content_layer\n\n        self.mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)\n        self.std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)\n\n    @torch.no_grad()\n    def extract_features(self, x, until_layer=None):\n        features = {}\n        current = x\n        for i, layer in enumerate(self.feature_extractor.children()):\n            current = layer(current)\n            if str(i) in self.layers:\n                features[self.layers[str(i)]] = current\n                if until_layer is not None and self.layers[str(i)] == until_layer:\n                    break\n        return features\n\n    def forward(self, pred, target, style_target=None):\n        x_pred = (pred + 1) / 2\n        x_pred = (x_pred - self.mean) / self.std\n        x_target = (target + 1) / 2\n        x_target = (x_target - self.mean) / self.std\n        if style_target is not None:\n            x_style = (style_target + 1) / 2\n            x_style = (x_style - self.mean) / self.std\n\n        if style_target is None:\n            pred_features = self.extract_features(x_pred, until_layer=self.content_layer)\n            target_features = self.extract_features(x_target, until_layer=self.content_layer)\n        else:\n            pred_features = self.extract_features(x_pred)\n            target_features = self.extract_features(x_target)\n            style_features = self.extract_features(x_style)\n\n        content_loss = 0.0\n        if self.content_weight > 0:\n            content_loss = self.loss_fn(pred_features[self.content_layer], target_features[self.content_layer])\n\n        style_loss = 0.0\n        if style_target is not None and self.style_weight > 0:\n            style_grams = {key: gram_matrix(style_features[key]) for key in style_features.keys()}\n            pred_grams = {key: gram_matrix(pred_features[key]) for key in pred_features.keys()}\n\n            for key in self.layers.values():\n                if key in pred_grams and key in style_grams:\n                    style_loss += self.loss_fn(pred_grams[key], style_grams[key])\n\n        total_loss = self.content_weight * content_loss + self.style_weight * style_loss\n        return total_loss","metadata":{"execution":{"iopub.execute_input":"2025-05-05T15:52:41.005959Z","iopub.status.busy":"2025-05-05T15:52:41.005777Z","iopub.status.idle":"2025-05-05T15:52:41.015618Z","shell.execute_reply":"2025-05-05T15:52:41.015062Z"},"papermill":{"duration":0.024033,"end_time":"2025-05-05T15:52:41.016659","exception":false,"start_time":"2025-05-05T15:52:40.992626","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"2f77c3d8","cell_type":"code","source":"import torch\nimport torch.nn as nn\n#from loss import reconstruction_loss_with_feature\n\n\ndef train_discriminator(discriminator,\n                        true_data,\n                        gen_data,\n                        optimizer_d,\n                        criterion = nn.MSELoss(),\n                        device = 'cpu'):\n\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\ndef train_generator(\n    discriminator,\n    true_data,\n    gen_data,\n    cycle_data,\n    iden_data,\n    optimizer_g,\n    criterion=nn.MSELoss(),\n    adv_lambda=None,\n    cycle_lambda=None,\n    iden_lambda=None,\n    have_cycle=True,\n    have_iden=True,\n    cycle_gamma=None,\n    iden_gamma=None,\n    cycle_feature_pos=None,\n    iden_feature_pos=None,\n    cycle_weight=None,\n    iden_weight=None,\n    perceptual_loss_value=0.0,\n    device='cpu'):\n\n    out_adv = discriminator(gen_data)\n    loss_adv = adv_lambda * criterion(out_adv, torch.ones_like(out_adv, device=device))\n\n    if have_cycle:\n        loss_cycle = cycle_lambda * reconstruction_loss_with_feature(\n            discriminator,\n            true_data,\n            cycle_data,\n            gamma=cycle_gamma,\n            feature_mode=True,\n            feature_pos=cycle_feature_pos,\n            weights=cycle_weight)\n    else:\n        loss_cycle = torch.tensor(0.0)\n\n    if have_iden:\n        loss_iden = iden_lambda * reconstruction_loss_with_feature(\n            discriminator,\n            true_data,\n            iden_data,\n            gamma=iden_gamma,\n            feature_mode=True,\n            feature_pos=iden_feature_pos,\n            weights=iden_weight)\n    else:\n        loss_iden = torch.tensor(0.0)\n\n    total_loss = loss_adv + loss_cycle + loss_iden + perceptual_loss_value\n\n    optimizer_g.zero_grad()\n    total_loss.backward()\n    optimizer_g.step()\n\n    return (total_loss.item(),\n            loss_adv.item(),\n            loss_cycle.item(),\n            loss_iden.item())\n","metadata":{"execution":{"iopub.execute_input":"2025-05-05T15:52:41.042491Z","iopub.status.busy":"2025-05-05T15:52:41.042056Z","iopub.status.idle":"2025-05-05T15:52:41.049374Z","shell.execute_reply":"2025-05-05T15:52:41.048671Z"},"papermill":{"duration":0.021983,"end_time":"2025-05-05T15:52:41.050430","exception":false,"start_time":"2025-05-05T15:52:41.028447","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"8e41ffa3","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.execute_input":"2025-05-05T15:52:41.074706Z","iopub.status.busy":"2025-05-05T15:52:41.074507Z","iopub.status.idle":"2025-05-05T15:52:41.079087Z","shell.execute_reply":"2025-05-05T15:52:41.078581Z"},"papermill":{"duration":0.018056,"end_time":"2025-05-05T15:52:41.080135","exception":false,"start_time":"2025-05-05T15:52:41.062079","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"cd42bccf","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.1):\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.execute_input":"2025-05-05T15:52:41.104649Z","iopub.status.busy":"2025-05-05T15:52:41.104223Z","iopub.status.idle":"2025-05-05T15:52:41.110233Z","shell.execute_reply":"2025-05-05T15:52:41.109551Z"},"papermill":{"duration":0.0195,"end_time":"2025-05-05T15:52:41.111320","exception":false,"start_time":"2025-05-05T15:52:41.091820","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"cd5beb58","cell_type":"code","source":"import random\nfrom torch.utils.data import Subset\nfrom torchvision.transforms import (Compose,\n                                    Resize,\n                                    RandomCrop,\n                                    CenterCrop,\n                                    RandomHorizontalFlip,\n                                    ColorJitter,\n                                    ToTensor,\n                                    Normalize)\n\n\ndef split_data(data, ratio_train_size, val_size, seed=42):\n    random.seed(seed)\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, seed=42):\n    random.seed(seed)\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(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    raw_transform = Compose([\n        Resize(resize_size),\n        RandomCrop(image_size),\n        RandomHorizontalFlip(p=0.5),\n        ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n        ToTensor(),\n        Normalize(normalize_mean, normalize_std)\n    ])\n\n    style_transform = Compose([\n        Resize(resize_size),\n        RandomCrop(image_size),\n        RandomHorizontalFlip(p=0.5),\n        ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n        ToTensor(),\n        Normalize(normalize_mean, normalize_std)\n    ])\n\n    val_transform = Compose([\n        Resize(resize_size),\n        CenterCrop(image_size),\n        ToTensor(),\n        Normalize(normalize_mean, normalize_std)\n    ])\n\n    return raw_transform, style_transform, val_transform\n","metadata":{"execution":{"iopub.execute_input":"2025-05-05T15:52:41.135838Z","iopub.status.busy":"2025-05-05T15:52:41.135660Z","iopub.status.idle":"2025-05-05T15:52:41.142726Z","shell.execute_reply":"2025-05-05T15:52:41.142197Z"},"papermill":{"duration":0.020487,"end_time":"2025-05-05T15:52:41.143713","exception":false,"start_time":"2025-05-05T15:52:41.123226","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"47dee952","cell_type":"code","source":"import os\nos.environ[\"TF_ENABLE_ONEDNN_OPTS\"] = \"0\"\nimport shutil\nimport torch\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=2)\n    parser.add_argument('--learning_rate', '-lr', type=float, default=0.0002)\n    parser.add_argument('--ratio_train_size', '-rt', type=float, default=0.5)\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=2)\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('--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/gan-getting-started/photo_jpg')\n    parser.add_argument('--style_image_path', '-sp', type=str, default='/kaggle/input/arcylic-painting/arcylic_painting')\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef train(arg):\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\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_ga = optim.Adam(gen_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_gb = optim.Adam(gen_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_da = optim.Adam(dis_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_db = optim.Adam(dis_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n\n    scheduler_ga = get_lr_scheduler(op_ga, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n    scheduler_gb = get_lr_scheduler(op_gb, 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_ga.load_state_dict(checkpoint['op_ga'])\n        op_gb.load_state_dict(checkpoint['op_gb'])\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        pa = checkpoint['pa']\n        pb = checkpoint['pb']\n        ema_lambda_ra = checkpoint['ema_lambda_ra']\n        ema_lambda_rb = checkpoint['ema_lambda_rb']\n        dict_feature_a = checkpoint['dict_feature_a']\n        dict_feature_b = checkpoint['dict_feature_b']\n    else:\n        start_epoch = 0\n        best_loss = float('inf')\n        pa = 0.0\n        pb = 0.0\n        ema_lambda_ra = 0.0\n        ema_lambda_rb = 0.0\n        dict_feature_a = {}\n        dict_feature_b = {}\n\n    train_iters = len(train) // arg.batch_size\n    val_iters = len(val) // arg.batch_size\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\n        mean_perceptual_a = 0.0\n        mean_perceptual_b = 0.0\n\n        # gamma = get_lambda_scheduler(epoch=epoch,\n        #                              total_epochs=10,\n        #                              decay_start_epoch=0,\n        #                              start_lambda=0,\n        #                              end_lambda=1,\n        #                              mode='linear_reverse')\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                cycle_a = gen_a(out_g2.detach())\n                cycle_b = gen_b(out_g1.detach())\n                iden_a = gen_a(true_image)\n                iden_b = gen_b(raw_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                #mixed_a = mix_pseudo_real(true_image, out_g1_da,pa)\n                #mixed_b = mix_pseudo_real(raw_image, out_g2_db,pb)\n                #perceptual_a = perceptual_loss(cycle_a, true_image)\n                #perceptual_b = perceptual_loss(cycle_b, raw_image)\n\n                loss_da = train_discriminator(discriminator=dis_a,\n                                              true_data=true_image,\n                                              gen_data=out_g1_da,\n                                              optimizer_d=op_da,\n                                              device=device)\n                mean_loss_da += loss_da\n                # if iteration % 64 == 0:\n                #     pa, ema_lambda_ra = update_p(\n                #         discriminator=dis_a,\n                #         real_images=true_image,\n                #         p=pa,\n                #         ema_lambda_r=ema_lambda_ra)\n\n                loss_db = train_discriminator(discriminator=dis_b,\n                                              true_data=raw_image,\n                                              gen_data=out_g2_db,\n                                              optimizer_d=op_db,\n                                              device=device)\n                mean_loss_db += loss_db\n                # if iteration % 64 == 0:\n                #     pb, ema_lambda_rb = update_p(\n                #         discriminator=dis_b,\n                #         real_images=raw_image,\n                #         p=pb,\n                #         ema_lambda_r=ema_lambda_rb)\n\n                (loss_ga_final,\n                 loss_adv_a,\n                 loss_cycle_a,\n                 loss_iden_a) = train_generator(discriminator=dis_a,\n                                                true_data=true_image,\n                                                gen_data=out_g1,\n                                                cycle_data=cycle_a,\n                                                iden_data=iden_a,\n                                                optimizer_g=op_ga,\n                                                adv_lambda=arg.adv_lambda,\n                                                cycle_lambda=arg.cycle_lambda,\n                                                iden_lambda=arg.identity_lambda,\n                                                cycle_gamma=0,\n                                                iden_gamma=0,\n                                                have_cycle=True,\n                                                have_iden=True,\n                                                cycle_feature_pos=[2,3,4],\n                                                iden_feature_pos=[1],\n                                                perceptual_loss_value=0,\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_perceptual_a += perceptual_a\n\n                (loss_gb_final,\n                 loss_adv_b,\n                 loss_cycle_b,\n                 loss_iden_b) = train_generator(discriminator=dis_b,\n                                                true_data=raw_image,\n                                                gen_data=out_g2,\n                                                cycle_data=cycle_b,\n                                                iden_data=iden_b,\n                                                optimizer_g=op_gb,\n                                                adv_lambda=arg.adv_lambda,\n                                                cycle_lambda=arg.cycle_lambda,\n                                                iden_lambda=arg.identity_lambda,\n                                                cycle_gamma=0,\n                                                iden_gamma=0,\n                                                have_cycle=True,\n                                                have_iden=True,\n                                                cycle_feature_pos=[2,3,4],\n                                                iden_feature_pos=[1],\n                                                perceptual_loss_value=0,\n                                                device=device)\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_perceptual_b += perceptual_b\n\n            scheduler_ga.step()\n            scheduler_gb.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'perceptual ga: {mean_perceptual_a / train_iters:.4f} - perceptual gb: {mean_perceptual_b / train_iters}')\n            print(f'Gamma: {0.0} |  Pa: {pa:.4f} - ema_a: {ema_lambda_ra:.4f} | '\n                  f'Pb: {pb:.5f} - ema_b: {ema_lambda_rb:.5f}')\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('perceptual ga', mean_perceptual_a / train_iters, epoch)\n                writer.add_scalar('perceptual gb', mean_perceptual_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_ga': op_ga.state_dict(),\n            'op_gb': op_gb.state_dict(),\n            'op_da': op_da.state_dict(),\n            'op_db': op_db.state_dict(),\n            'best_loss': best_loss,\n            'pa': pa,\n            'pb': pb,\n            'ema_lambda_ra': ema_lambda_ra,\n            'ema_lambda_rb': ema_lambda_rb,\n            'dict_feature_a': dict_feature_a,\n            'dict_feature_b': dict_feature_b\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.execute_input":"2025-05-05T15:52:41.168429Z","iopub.status.busy":"2025-05-05T15:52:41.168246Z","iopub.status.idle":"2025-05-06T02:58:39.133799Z","shell.execute_reply":"2025-05-06T02:58:39.132945Z"},"papermill":{"duration":39957.979795,"end_time":"2025-05-06T02:58:39.135304","exception":false,"start_time":"2025-05-05T15:52:41.155509","status":"completed"},"tags":[],"_kg_hide-output":true},"outputs":[],"execution_count":null},{"id":"93089a30","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                     raw_transform=transform,\n                     style_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, 6, figsize=(20, 5))\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)\n\n        for value in range(0, len(ax), 2):\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 == 5:\n            break","metadata":{"execution":{"iopub.execute_input":"2025-05-06T02:58:54.725340Z","iopub.status.busy":"2025-05-06T02:58:54.724407Z","iopub.status.idle":"2025-05-06T02:59:51.576167Z","shell.execute_reply":"2025-05-06T02:59:51.575410Z"},"papermill":{"duration":64.648457,"end_time":"2025-05-06T02:59:51.589799","exception":false,"start_time":"2025-05-06T02:58:46.941342","status":"completed"},"tags":[]},"outputs":[],"execution_count":null}]}