{"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.13"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":1057.598622,"end_time":"2025-07-27T12:17:20.682844","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-07-27T11:59:43.084222","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"0c5633f6","cell_type":"code","source":"import os\nimport random\nimport torch\nfrom torchvision.io import read_image\nfrom torch.utils.data import Dataset\n\nclass CustomUnpairedDataset(Dataset):\n    def __init__(self, domainA_dir, domainB_dir, transform=None):\n        self.A_paths = [os.path.join(domainA_dir, f) for f in os.listdir(domainA_dir) if f.endswith(\".jpg\")]\n        self.B_paths = [os.path.join(domainB_dir, f) for f in os.listdir(domainB_dir) if f.endswith(\".jpg\")]\n        self.transform = transform\n\n    def __getitem__(self, idx):\n        A_img = read_image(self.A_paths[idx % len(self.A_paths)]).float() / 255.0  # scale to [0, 1]\n        B_img = read_image(random.choice(self.B_paths)).float() / 255.0\n\n        if self.transform:\n            A_img = self.transform(A_img)\n            B_img = self.transform(B_img)\n\n        return A_img, B_img\n\n    def __len__(self):\n        return min(len(self.A_paths), len(self.B_paths))\n","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:47.299729Z","iopub.status.busy":"2025-07-27T11:59:47.299527Z","iopub.status.idle":"2025-07-27T11:59:56.244155Z","shell.execute_reply":"2025-07-27T11:59:56.243574Z"},"papermill":{"duration":8.949514,"end_time":"2025-07-27T11:59:56.245541","exception":false,"start_time":"2025-07-27T11:59:47.296027","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"494a04cf","cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n'''\nDiscrinimator architecture for cyclic GAN. \n'''\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3):\n        super(Discriminator, self).__init__()\n\n        def conv_block(in_features, out_features, normalize=True):\n            layers = [nn.Conv2d(in_features, out_features, kernel_size=4, stride=2, padding=1)]\n            if normalize:\n                '''\n                Use instance norm over batch norm for generating images. \n                '''\n                layers.append(nn.InstanceNorm2d(out_features))\n            layers.append(nn.LeakyReLU(0.2, inplace=True))\n            return layers\n\n        # 70x70 PatchGAN\n        self.model = nn.Sequential(\n            *conv_block(in_channels, 64, normalize=False),  # 256 → 128\n            *conv_block(64, 128),                           # 128 → 64\n            *conv_block(128, 256),                          # 64 → 32\n            *conv_block(256, 512),                          # 32 → 16\n            nn.Conv2d(512, 1, kernel_size=4, padding=1)     # 16 → ~14 (patch map)\n        )\n\n    def forward(self, x):\n        return self.model(x)\n","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:56.251596Z","iopub.status.busy":"2025-07-27T11:59:56.250964Z","iopub.status.idle":"2025-07-27T11:59:56.256483Z","shell.execute_reply":"2025-07-27T11:59:56.255948Z"},"papermill":{"duration":0.009231,"end_time":"2025-07-27T11:59:56.257519","exception":false,"start_time":"2025-07-27T11:59:56.248288","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"99c5ad57","cell_type":"code","source":"import torch.nn as nn\n\nclass FeatureNetwork(nn.Module):\n    def __init__(self, in_channels=256, out_channels=256):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, 1)\n        )\n\n    def forward(self, x):\n        return self.net(x)\n","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:56.262632Z","iopub.status.busy":"2025-07-27T11:59:56.262006Z","iopub.status.idle":"2025-07-27T11:59:56.266358Z","shell.execute_reply":"2025-07-27T11:59:56.265829Z"},"papermill":{"duration":0.007848,"end_time":"2025-07-27T11:59:56.267427","exception":false,"start_time":"2025-07-27T11:59:56.259579","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"d350df8b","cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n'''\nA residual block is a building block used in ResNet which helps neural networks learn. \nBasically it learns the change it wants to apply in the input and then apply the change back in. \nThis helps in gradient flow in backpropagation and reuse of low-level features.\n'''\nimport torch\nimport torch.nn as nn\n\nclass ResnetBlock(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(dim, dim, kernel_size=3, bias=False),\n            nn.InstanceNorm2d(dim),\n            nn.ReLU(True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(dim, dim, kernel_size=3, bias=False),\n            nn.InstanceNorm2d(dim)\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\n\nclass ResNetGenerator(nn.Module):\n    def __init__(self, in_channels=3, out_channels=3, ngf=64, n_blocks=9):\n        super().__init__()\n        model = []\n\n        # Initial conv\n        model += [\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(in_channels, ngf, kernel_size=7, bias=False),\n            nn.InstanceNorm2d(ngf),\n            nn.ReLU(True)\n        ]\n\n        # Downsampling\n        in_feat = ngf\n        out_feat = ngf * 2\n        for _ in range(2):\n            model += [\n                nn.Conv2d(in_feat, out_feat, kernel_size=3, stride=2, padding=1, bias=False),\n                nn.InstanceNorm2d(out_feat),\n                nn.ReLU(True)\n            ]\n            in_feat = out_feat\n            out_feat = in_feat * 2\n\n        # Residual blocks\n        for _ in range(n_blocks):\n            model += [ResnetBlock(in_feat)]\n\n        # Upsampling\n        out_feat = in_feat // 2\n        for _ in range(2):\n            model += [\n                nn.ConvTranspose2d(in_feat, out_feat, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),\n                nn.InstanceNorm2d(out_feat),\n                nn.ReLU(True)\n            ]\n            in_feat = out_feat\n            out_feat = in_feat // 2\n\n        # Final conv\n        model += [\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(ngf, out_channels, kernel_size=7),\n            nn.Tanh()\n        ]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x, return_feature=False, feature_layer=3):\n        feat = None\n        for i, layer in enumerate(self.model):\n            x = layer(x)\n            if i == feature_layer:\n                feat = x\n        if return_feature:\n            return x, feat\n        return x\n\n","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:56.272423Z","iopub.status.busy":"2025-07-27T11:59:56.272047Z","iopub.status.idle":"2025-07-27T11:59:56.280540Z","shell.execute_reply":"2025-07-27T11:59:56.280020Z"},"papermill":{"duration":0.012056,"end_time":"2025-07-27T11:59:56.281549","exception":false,"start_time":"2025-07-27T11:59:56.269493","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"9bb8c8de","cell_type":"code","source":"import torch\nimport torch.nn.functional as F\n\nclass PatchNCELoss(nn.Module):\n    def __init__(self, temperature=0.07):\n        super().__init__()\n        self.temperature = temperature\n\n    def forward(self, feat_q, feat_k):\n        B, C, H, W = feat_q.shape\n        feat_q = feat_q.permute(0, 2, 3, 1).reshape(-1, C)\n        feat_k = feat_k.permute(0, 2, 3, 1).reshape(-1, C)\n\n        feat_q = F.normalize(feat_q, dim=1)\n        feat_k = F.normalize(feat_k, dim=1)\n\n        logits = torch.matmul(feat_q, feat_k.T) / self.temperature\n        labels = torch.arange(feat_q.size(0), device=feat_q.device)\n        loss = F.cross_entropy(logits, labels)\n        return loss\n","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:56.286570Z","iopub.status.busy":"2025-07-27T11:59:56.286229Z","iopub.status.idle":"2025-07-27T11:59:56.291023Z","shell.execute_reply":"2025-07-27T11:59:56.290526Z"},"papermill":{"duration":0.008227,"end_time":"2025-07-27T11:59:56.291943","exception":false,"start_time":"2025-07-27T11:59:56.283716","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"fcbedd2c","cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.utils import save_image\nfrom itertools import cycle\nimport os\nimport matplotlib.pyplot as plt\nimport torchvision.transforms as transforms\n\n\n# ========== Config ==========\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nSAVE_EVERY = 5\nBATCH_SIZE = 1\nEPOCHS = 15\nIMAGE_SIZE = 256\nCHECKPOINT_DIR = \"checkpoints_cut\"\n\nos.makedirs(CHECKPOINT_DIR, exist_ok=True)\n\n# ========== Model ==========\ngenerator = ResNetGenerator().to(DEVICE)\ndiscriminator = Discriminator().to(DEVICE)\nfeature_net = FeatureNetwork(in_channels=128).to(DEVICE)\nnce_loss_fn = PatchNCELoss().to(DEVICE)\ncriterion_GAN = nn.MSELoss()\n\n\noptimizer_G = optim.Adam(list(generator.parameters()) + list(feature_net.parameters()), lr=2e-4, betas=(0.5, 0.999))\noptimizer_D = optim.Adam(discriminator.parameters(), lr=1e-4, betas=(0.5, 0.999))\n\n# ========== Data ==========\nfrom torchvision import transforms\n\ntransform = transforms.Compose([\n    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),\n    transforms.Normalize([0.5]*3, [0.5]*3)  # assumes 3-channel RGB\n])\n\ndataset = CustomUnpairedDataset(\"/kaggle/input/gan-getting-started/photo_jpg\", \"/kaggle/input/gan-getting-started/monet_jpg\", transform=transform)\ndataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=4)\ndataloader = cycle(dataloader) \n\n# ========== Forward Feature Extract ==========\ndef extract_features(model, x, layer_index=6):\n    \"\"\"Forward pass and extract feature maps from a specific layer\"\"\"\n    features = None\n    for i, layer in enumerate(model.model):\n        x = layer(x)\n        if i == layer_index:\n            features = x\n    return x, features\n\n# ========== Training ==========\nstart_epoch = 0\n\nfor epoch in range(start_epoch, EPOCHS):\n    generator.train()\n    feature_net.train()\n\n    total_nce_loss = 0\n\n    for step in range(300 // BATCH_SIZE):\n        real_A, real_B = next(dataloader)\n        real_A = real_A.to(DEVICE)\n        real_B = real_B.to(DEVICE)\n\n        # Forward pass through generator and get intermediate features\n        fake_B, feat_q = extract_features(generator, real_A)\n        with torch.no_grad():\n            _, feat_k = extract_features(generator, real_A)\n\n        # Real vs Fake\n        real_validity = discriminator(real_B)\n        fake_validity = discriminator(fake_B.detach())  # don't backprop through G\n        \n        valid = torch.ones_like(real_validity)\n        fake = torch.zeros_like(fake_validity)\n        \n        # PatchNCE loss\n        feat_q = feature_net(feat_q)\n        feat_k = feature_net(feat_k).detach()\n        loss_nce = nce_loss_fn(feat_q, feat_k)\n\n        pred_fake = discriminator(fake_B)\n        loss_G_GAN = criterion_GAN(pred_fake, valid)\n\n        λ_nce = 10.0\n        λ_idt = 5.0\n        idt_B = generator(real_B)[0]  # only take the image output\n        loss_idt = nn.L1Loss()(idt_B, real_B)\n        \n        loss_G = loss_G_GAN + λ_nce * loss_nce + λ_idt*loss_idt\n\n        # Backprop\n        optimizer_G.zero_grad()\n        loss_G.backward()\n        optimizer_G.step()\n\n        total_nce_loss += loss_nce.item()\n        \n        loss_D_real = criterion_GAN(real_validity, valid)\n        loss_D_fake = criterion_GAN(fake_validity, fake)\n        loss_D = 0.5 * (loss_D_real + loss_D_fake)\n        \n        optimizer_D.zero_grad()\n        loss_D.backward()\n        optimizer_D.step()\n\n\n        if step % 100 == 0:\n            print(f\"[Epoch {epoch+1}/{EPOCHS}] [Step {step}] [NCE loss: {loss_nce.item():.4f}] [Generator loss: {loss_G.item():.4f}] [Discriminator loss: {loss_D.item():.4f}]\")\n\n    avg_loss = total_nce_loss / (step + 1)\n    print(f\"==> Epoch {epoch + 1}: Avg NCE Loss = {avg_loss:.4f}\")\n    '''\n    if (epoch + 1) % 3 == 0:\n        generator.eval()\n        with torch.no_grad():\n            real_A_vis, _ = next(dataloader)  # Get a sample batch\n            real_A_vis = real_A_vis.to(DEVICE)\n            \n            fake_B_vis, _ = generator(real_A_vis, return_feature=True)\n            fake_B_vis = fake_B_vis[0].detach().cpu()  # Get first image in batch\n            fake_B_vis = fake_B_vis * 0.5 + 0.5  # De-normalize to [0,1]\n\n            real_A_vis = real_A_vis[0].detach().cpu()\n            real_A_vis = real_A_vis * 0.5 + 0.5\n            real_img = transforms.ToPILImage()(real_A_vis.clamp(0, 1))\n\n        plt.figure(figsize=(10, 5))\n        plt.subplot(1, 2, 1)\n        plt.imshow(real_img)\n        plt.title(\"Input Photo (real_A)\")\n        plt.axis('off')\n\n        # Convert tensor to image and plot\n        img = transforms.ToPILImage()(fake_B_vis.clamp(0, 1))\n        plt.figure(figsize=(5, 5))\n        plt.imshow(img)\n        plt.title(f\"Epoch_{epoch+1}\")\n        plt.axis('off')\n        plt.show()\n'''","metadata":{"execution":{"iopub.execute_input":"2025-07-27T11:59:56.296940Z","iopub.status.busy":"2025-07-27T11:59:56.296577Z","iopub.status.idle":"2025-07-27T12:14:04.073669Z","shell.execute_reply":"2025-07-27T12:14:04.072798Z"},"papermill":{"duration":847.781,"end_time":"2025-07-27T12:14:04.074908","exception":false,"start_time":"2025-07-27T11:59:56.293908","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"418a343d","cell_type":"code","source":"import torchvision.transforms.v2 as transforms\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])","metadata":{"execution":{"iopub.execute_input":"2025-07-27T12:14:04.084261Z","iopub.status.busy":"2025-07-27T12:14:04.084011Z","iopub.status.idle":"2025-07-27T12:14:04.184656Z","shell.execute_reply":"2025-07-27T12:14:04.184136Z"},"papermill":{"duration":0.106463,"end_time":"2025-07-27T12:14:04.185730","exception":false,"start_time":"2025-07-27T12:14:04.079267","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"ab252ed6","cell_type":"code","source":"from PIL import Image\n\nos.makedirs(\"/kaggle/images/\", exist_ok=True)\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    generator.eval()\n    with torch.no_grad(): op = generator(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\")","metadata":{"execution":{"iopub.execute_input":"2025-07-27T12:14:04.194439Z","iopub.status.busy":"2025-07-27T12:14:04.194224Z","iopub.status.idle":"2025-07-27T12:17:16.036097Z","shell.execute_reply":"2025-07-27T12:17:16.035175Z"},"papermill":{"duration":191.848003,"end_time":"2025-07-27T12:17:16.037776","exception":false,"start_time":"2025-07-27T12:14:04.189773","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"c569af07","cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"execution":{"iopub.execute_input":"2025-07-27T12:17:16.048738Z","iopub.status.busy":"2025-07-27T12:17:16.048508Z","iopub.status.idle":"2025-07-27T12:17:18.948150Z","shell.execute_reply":"2025-07-27T12:17:18.947464Z"},"papermill":{"duration":2.906102,"end_time":"2025-07-27T12:17:18.949366","exception":false,"start_time":"2025-07-27T12:17:16.043264","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"d2407065","cell_type":"code","source":"","metadata":{"papermill":{"duration":0.003769,"end_time":"2025-07-27T12:17:18.957408","exception":false,"start_time":"2025-07-27T12:17:18.953639","status":"completed"},"tags":[]},"outputs":[],"execution_count":null}]}