{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torchvision.transforms.v2 as transforms\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nfrom typing import Tuple\n\nimport itertools\nimport os\nimport shutil","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T08:29:42.747017Z","iopub.execute_input":"2025-07-19T08:29:42.747234Z","iopub.status.idle":"2025-07-19T08:29:50.375803Z","shell.execute_reply.started":"2025-07-19T08:29:42.747218Z","shell.execute_reply":"2025-07-19T08:29:50.375194Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Custom dataset class and transformation implementations ###\n\nclass ImageStyleDataset(Dataset):\n    def __init__(self, style_dir, image_dir, transform=None):\n        super().__init__()\n        self.style_dir = style_dir\n        self.image_dir = image_dir\n        self.style_paths = os.listdir(self.style_dir)\n        self.image_paths = os.listdir(self.image_dir)\n        self.transform = transform\n\n    def __len__(self):\n        return min(len(self.style_paths), len(self.image_paths))\n    \n    def __getitem__(self, index):\n        style_path = os.path.join(self.style_dir, self.style_paths[index])\n        image_path = os.path.join(self.image_dir, self.image_paths[index])\n        style = Image.open(style_path)\n        image = Image.open(image_path)\n        if self.transform:\n            style = self.transform(style)\n            image = self.transform(image)\n        return style, image\n\n\ndef init_cyclegan_transform(image_size, image_size_resized):\n    transform = transforms.Compose([\n        transforms.Resize(image_size_resized),\n        transforms.RandomCrop(image_size),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToImage(),\n        transforms.ToDtype(torch.float32, scale=True),\n        transforms.Normalize([0.5], [0.5])\n    ])\n    return transform","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T08:29:50.465094Z","iopub.execute_input":"2025-07-19T08:29:50.465907Z","iopub.status.idle":"2025-07-19T08:29:50.481268Z","shell.execute_reply.started":"2025-07-19T08:29:50.465874Z","shell.execute_reply":"2025-07-19T08:29:50.480623Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Construction of generators and discriminators ###\n\nCYCLEGAN_GEN_CONFIG = {\n    \"DownConvBlock\": [\n        (3, 64, 0.0, False),\n        (64, 128, 0.0, True),\n        (128, 256, 0.0, True),\n        (256, 512, 0.5, True),\n        (512, 512, 0.5, True),\n        (512, 512, 0.5, True),\n        (512, 512, 0.5, True),\n        (512, 512, 0.5, False)\n    ],\n    \"UpConvBlock\": [\n        (512, 512, 0.5),\n        (1024, 512, 0.5),\n        (1024, 512, 0.5),\n        (1024, 512, 0.5),\n        (1024, 256, 0.0),\n        (512, 128, 0.0),\n        (256, 64, 0.0)\n    ],\n    \"Upsample\": 2,\n    \"ZeroPad2d\": (1, 0, 1, 0),\n    \"final_conv\": (128, 3, 4, 1, 1)\n}\n\nCYCLEGAN_DSC_CONFIG = {\n    \"disc_conv_block\": [\n        (3, 64, False),\n        (64, 128, True),\n        (128, 256, True),\n        (256, 512, True)\n    ],\n    \"ZeroPad2d\": (1, 0, 1, 0),\n    \"final_conv\": (512, 1, 4, 1, 1)\n}\n\n\nclass DownConvBlock(nn.Module):\n    def __init__(self, input_size, output_size, dropout=0.0, norm=True):\n        super().__init__()\n        self.layers = nn.ModuleList([\n            nn.Conv2d(input_size, output_size, 4, 2, 1)\n        ])\n        if norm: self.layers.append(nn.InstanceNorm2d(output_size))\n        self.layers.append(nn.LeakyReLU(0.2))\n        if dropout: self.layers.append(nn.Dropout(dropout))\n\n    def forward(self, x):\n        res = nn.Sequential(*self.layers)(x)\n        return res\n\n\nclass UpConvBlock(nn.Module):\n    def __init__(self, input_size, output_size, dropout=0.0):\n        super().__init__()\n        self.layers = nn.ModuleList([\n            nn.ConvTranspose2d(input_size, output_size, 4, 2, 1),\n            nn.InstanceNorm2d(output_size),\n            nn.ReLU()\n        ])\n        if dropout: self.layers.append(nn.Dropout(dropout))\n\n    def forward(self, x, encoding_input):\n        x = nn.Sequential(*self.layers)(x)\n        res = torch.cat((x, encoding_input), dim=1)\n        return res\n\n\nclass CycleGANGenerator(nn.Module):\n    def __init__(self, config: dict):\n        super(CycleGANGenerator, self).__init__()\n        \n        self.downconv_layers = nn.ModuleList([])\n        for layer in config.get(\"DownConvBlock\", []):\n            self.downconv_layers.append(DownConvBlock(*layer))\n\n        self.upconv_layers = nn.ModuleList([])\n        for layer in config.get(\"UpConvBlock\", []):\n            self.upconv_layers.append(UpConvBlock(*layer))\n\n        self.final_layers = nn.ModuleList([\n            nn.Upsample(scale_factor=config[\"Upsample\"]),\n            nn.ZeroPad2d(padding=config[\"ZeroPad2d\"]),\n            nn.Conv2d(*config[\"final_conv\"]),\n            nn.Tanh()\n        ])\n\n\n    def forward(self, x):\n        encs = [x]\n        \n        m = len(self.downconv_layers)\n        n = len(self.upconv_layers)\n\n        if self.downconv_layers:\n            for i in range(m):\n                encs.append(self.downconv_layers[i](encs[i]))\n\n        res = encs[-1]\n        if self.upconv_layers and len(encs) >= 2:\n            for i in range(n):\n                res = self.upconv_layers[i](res, encs[n-i])\n\n        res = nn.Sequential(*self.final_layers)(res)\n        return res\n\n\nclass CycleGANDiscriminator(nn.Module):\n    def __init__(self, config: dict):\n        super(CycleGANDiscriminator, self).__init__()\n\n        def disc_conv_block(input_size, output_size, norm=True):\n            mods = nn.ModuleList([\n                nn.Conv2d(input_size, output_size, 4, 2, 1)\n            ])\n            if norm: mods.append(nn.InstanceNorm2d(output_size))\n            mods.append(nn.LeakyReLU(0.2, inplace=True))\n            return mods\n        \n        self.layers = nn.ModuleList([])\n        for layer in config.get(\"disc_conv_block\", []):\n            mods = disc_conv_block(*layer)\n            for mod in mods:\n                self.layers.append(mod)\n\n        self.layers.append(nn.ZeroPad2d(padding=config[\"ZeroPad2d\"]))\n        self.layers.append(nn.Conv2d(*config[\"final_conv\"]))\n\n    def forward(self, x):\n        res = nn.Sequential(*self.layers)(x)\n        return res","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T08:29:50.560946Z","iopub.execute_input":"2025-07-19T08:29:50.561120Z","iopub.status.idle":"2025-07-19T08:29:50.566358Z","shell.execute_reply.started":"2025-07-19T08:29:50.561105Z","shell.execute_reply":"2025-07-19T08:29:50.565511Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Set of helper functions to aid with model training ###\n\ndef set_requires_grad(model, requires_grad):\n    for p in model.parameters(): p.requires_grad_(requires_grad)\n\n\ndef init_weights(net, gain=0.02):\n    def init_func(m):\n        classname = m.__class__.__name__\n        if hasattr(m, \"weight\") and (classname.find(\"Conv\") != -1 or classname.find(\"Linear\") != -1):\n            torch.nn.init.normal_(m.weight.data, 0.0, gain)\n            if hasattr(m, \"bias\") and m.bias is not None:\n                torch.nn.init.constant_(m.bias.data, 0.0)\n        elif classname.find(\"BatchNorm2d\") != -1:\n            torch.nn.init.normal_(m.weight.data, 1.0, gain)\n            torch.nn.init.constant_(m.bias.data, 0.0)\n    net.apply(init_func)\n\n\ndef calc_gen_loss(gen_values, loss_fn, device=\"cpu\"):\n    labels = torch.FloatTensor(gen_values.shape).fill_(1.0).to(device)\n    labels.requires_grad_(False)\n    return loss_fn(gen_values, labels)\n\n\ndef calc_dsc_loss(gen_values, tru_values, loss_fn, device=\"cpu\"):\n    gt_labels = torch.FloatTensor(tru_values.shape).fill_(1.0).to(device)\n    fk_labels = torch.FloatTensor(gen_values.shape).fill_(0.0).to(device)\n    gt_labels.requires_grad_(False)\n    fk_labels.requires_grad_(False)\n    gt_loss = loss_fn(tru_values, gt_labels)\n    fk_loss = loss_fn(gen_values, fk_labels)\n    return (gt_loss + fk_loss) / 2\n\n\ndef calc_cyc_loss(cyc_values, tru_values, loss_fn, factor):\n    return factor * loss_fn(cyc_values, tru_values)\n\n\ndef calc_idt_loss(sme_values, tru_values, loss_fn, factor):\n    return factor * loss_fn(sme_values, tru_values)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-19T08:29:50.605927Z","iopub.execute_input":"2025-07-19T08:29:50.606104Z","iopub.status.idle":"2025-07-19T08:29:50.618862Z","shell.execute_reply.started":"2025-07-19T08:29:50.606091Z","shell.execute_reply":"2025-07-19T08:29:50.618259Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Model training function ###\n\ndef train_cyclegan(\n    epochs: int,\n    train_dataloader: torch.utils.data.DataLoader,\n    generators: Tuple[CycleGANGenerator, CycleGANGenerator],\n    discriminators: Tuple[CycleGANDiscriminator, CycleGANDiscriminator],\n    optimizers: Tuple[torch.optim.Optimizer, torch.optim.Optimizer],\n    loss_fn: torch.nn.Module,\n    mae_fn: torch.nn.Module,\n    mae_fn_lambda: float,\n    device\n):\n    gen_g, gen_f = generators\n    dsc_x, dsc_y = discriminators\n    opt_gen, opt_dsc = optimizers\n    \n    gen_g.train()\n    gen_f.train()\n    dsc_x.train()\n    dsc_y.train()\n\n    for epoch in range(epochs):\n        for i, (tru_style, tru_image) in enumerate(train_dataloader):\n            tru_style = tru_style.to(device)\n            tru_image = tru_image.to(device)\n    \n            set_requires_grad(dsc_x, False)\n            set_requires_grad(dsc_y, False)\n    \n            opt_gen.zero_grad()\n    \n            gen_image = gen_g(tru_style)\n            gen_style = gen_f(tru_image)\n    \n            cyc_style = gen_f(gen_image)\n            cyc_image = gen_g(gen_style)\n    \n            sme_style = gen_f(tru_style)\n            sme_image = gen_g(tru_image)\n    \n            dsc_gen_style = dsc_x(gen_style)\n            dsc_gen_image = dsc_y(gen_image)\n    \n            total_gen_loss = (\n                calc_cyc_loss(cyc_style, tru_style, mae_fn, mae_fn_lambda) +\n                calc_cyc_loss(cyc_image, tru_image, mae_fn, mae_fn_lambda) +\n                calc_idt_loss(sme_style, tru_style, mae_fn, 0.5*mae_fn_lambda) +\n                calc_idt_loss(sme_image, tru_image, mae_fn, 0.5*mae_fn_lambda) +\n                calc_gen_loss(dsc_gen_style, loss_fn, device) +\n                calc_gen_loss(dsc_gen_image, loss_fn, device)\n            )\n\n            total_gen_loss.backward()\n            opt_gen.step()\n    \n            gen_style = gen_style.detach().clone().requires_grad_(True)\n            gen_image = gen_image.detach().clone().requires_grad_(True)\n    \n            set_requires_grad(dsc_x, True)\n            set_requires_grad(dsc_y, True)\n    \n            opt_dsc.zero_grad()\n    \n            dsc_tru_style = dsc_x(tru_style)\n            dsc_tru_image = dsc_y(tru_image)\n            dsc_gen_style = dsc_x(gen_style)\n            dsc_gen_image = dsc_y(gen_image)\n    \n            total_dsc_loss = (\n                calc_dsc_loss(dsc_gen_style, dsc_tru_style, loss_fn, device) +\n                calc_dsc_loss(dsc_gen_image, dsc_tru_image, loss_fn, device)\n            )\n    \n            total_dsc_loss.backward()\n            opt_dsc.step()\n            \n    \n            if i % 30 == 0:\n                print(f\"[INFO] Epoch: {(epoch+1):03d} | Gen Loss: {total_gen_loss:.4f} | Dsc Loss: {total_dsc_loss:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Execution of training phase ###\n\nBATCH_SIZE = 1\nEPOCH_NUMS = 50\nIMAGE_SIZE = 256\nIMAGE_SIZE_RESIZED = 286\nLAMBDA = 10\nLEARN_RATE = 2e-4\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndataset = ImageStyleDataset(\n    style_dir=\"/kaggle/input/gan-getting-started/monet_jpg\",\n    image_dir=\"/kaggle/input/gan-getting-started/photo_jpg\",\n    transform=init_cyclegan_transform(IMAGE_SIZE, IMAGE_SIZE_RESIZED)\n)\n\ndloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n\ngen_g = CycleGANGenerator(config=CYCLEGAN_GEN_CONFIG).to(device)\ngen_f = CycleGANGenerator(config=CYCLEGAN_GEN_CONFIG).to(device)\ndsc_x = CycleGANDiscriminator(config=CYCLEGAN_DSC_CONFIG).to(device)\ndsc_y = CycleGANDiscriminator(config=CYCLEGAN_DSC_CONFIG).to(device)\n\ninit_weights(gen_g)\ninit_weights(gen_f)\ninit_weights(dsc_x)\ninit_weights(dsc_y)\n\nmae_fn = torch.nn.L1Loss().to(device)\nloss_fn = torch.nn.BCEWithLogitsLoss().to(device)\nopt_gen = torch.optim.Adam(itertools.chain(gen_g.parameters(), gen_f.parameters()), lr=LEARN_RATE, betas=(0.5, 0.999))\nopt_dsc = torch.optim.Adam(itertools.chain(dsc_x.parameters(), dsc_y.parameters()), lr=LEARN_RATE, betas=(0.5, 0.999))\n\ntrain_cyclegan(\n    epochs=EPOCH_NUMS,\n    train_dataloader=dloader,\n    generators=(gen_g, gen_f),\n    discriminators=(dsc_x, dsc_y),\n    optimizers=(opt_gen, opt_dsc),\n    loss_fn=loss_fn,\n    mae_fn=mae_fn,\n    mae_fn_lambda=LAMBDA,\n    device=device\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Testing Phase ###\n\nphoto_dir = \"/kaggle/input/gan-getting-started/photo_jpg\"\nphoto_paths = os.listdir(photo_dir)\n\ntfm = transforms.Compose([\n    transforms.ToImage(),\n    transforms.ToDtype(torch.float32, scale=True),\n    transforms.Normalize([0.5], [0.5])\n])\n\nos.makedirs(\"/kaggle/images/\")\nfor i, path in enumerate(photo_paths):\n    photo_path = os.path.join(photo_dir, path)\n    photo_img = Image.open(photo_path)\n    photo_tsr = tfm(photo_img).unsqueeze(0).to(device)\n\n    gen_f.eval()\n    with torch.no_grad(): op = gen_f(photo_tsr).squeeze(0).cpu().detach()\n\n    op_unnormalized = op * 0.5 + 0.5\n    op_image = transforms.ToPILImage()(op_unnormalized).convert(\"RGB\")\n    op_image.save(\"/kaggle/images/\" + str(i+1) + \".jpg\")\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}