{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":1217826,"sourceType":"datasetVersion","datasetId":298806},{"sourceId":1504266,"sourceType":"datasetVersion","datasetId":885385},{"sourceId":2520560,"sourceType":"datasetVersion","datasetId":1526975},{"sourceId":3956508,"sourceType":"datasetVersion","datasetId":850761},{"sourceId":7457578,"sourceType":"datasetVersion","datasetId":1698586},{"sourceId":10700105,"sourceType":"datasetVersion","datasetId":6630834},{"sourceId":253477,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":216764,"modelId":238476}],"dockerImageVersionId":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -U albumentations","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:32.949095Z","iopub.execute_input":"2025-02-05T18:48:32.949456Z","iopub.status.idle":"2025-02-05T18:48:44.45593Z","shell.execute_reply.started":"2025-02-05T18:48:32.949417Z","shell.execute_reply":"2025-02-05T18:48:44.45506Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\n\nimport torch.nn as nn\n\nfrom PIL import Image\nimport os\nimport numpy as np\nimport torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport sys\nfrom tqdm import tqdm\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset\nimport numpy as np\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:44.457959Z","iopub.execute_input":"2025-02-05T18:48:44.458269Z","iopub.status.idle":"2025-02-05T18:48:50.229619Z","shell.execute_reply.started":"2025-02-05T18:48:44.45824Z","shell.execute_reply":"2025-02-05T18:48:50.228629Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"base_path = '/kaggle/working/' \nmonet_path = os.path.join(base_path, 'monet') \nactual_path = os.path.join(base_path, 'actual')\nos.makedirs(monet_path, exist_ok=True)\nos.makedirs(actual_path, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.231001Z","iopub.execute_input":"2025-02-05T18:48:50.231551Z","iopub.status.idle":"2025-02-05T18:48:50.238576Z","shell.execute_reply.started":"2025-02-05T18:48:50.231509Z","shell.execute_reply":"2025-02-05T18:48:50.237632Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define the new dataset paths\n\nHUMAN_PATH = \"/kaggle/input/human-faces/Humans\"\nPORTRAIT_PATHS = [\n    \"/kaggle/input/portrait-paintings/Images\",\n    \"/kaggle/input/art-portraits/Portraits/Portraits\",\n    \"/kaggle/input/art-portraits/Portraits_update/Portraits\"\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:50.239797Z","iopub.execute_input":"2025-02-05T18:48:50.2402Z","iopub.status.idle":"2025-02-05T18:48:50.250515Z","shell.execute_reply.started":"2025-02-05T18:48:50.240161Z","shell.execute_reply":"2025-02-05T18:48:50.249644Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Configuration**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename=\"/kaggle/working/.pth.tar\"):\n    print(\"=> Saving checkpoint\")\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(\"=> Loading checkpoint\")\n    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n    # If we don't do this then it will just have learning rate of old checkpoint\n    # and it will lead to many hours of debugging \\:\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\ndef seed_everything(seed=42):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n#TRAIN_DIR = \"data/train\"\n#VAL_DIR = \"data/val\"\nBATCH_SIZE = 20\nLEARNING_RATE = 1e-5\nLAMBDA_IDENTITY = 0.0\nLAMBDA_CYCLE = 10\nNUM_WORKERS = 4\nNUM_EPOCHS = 19\nLOAD_MODEL = False\nSAVE_MODEL = True\nCHECKPOINT_GEN_A = \"/kaggle/working/genH.pth.tar\"\nCHECKPOINT_GEN_M = \"/kaggle/working/genP.pth.tar\"\nCHECKPOINT_CRITIC_A = \"/kaggle/working/criticH.pth.tar\"\nCHECKPOINT_CRITIC_M = \"/kaggle/working/criticP.pth.tar\"\n\ntransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n    ],\n    additional_targets={\"image0\": \"image\"},\n    is_check_shapes=False\n)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.251713Z","iopub.execute_input":"2025-02-05T18:48:50.25202Z","iopub.status.idle":"2025-02-05T18:48:50.296432Z","shell.execute_reply.started":"2025-02-05T18:48:50.251986Z","shell.execute_reply":"2025-02-05T18:48:50.29556Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Generator**","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity(),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\nclass Generator(nn.Module):\n    def __init__(self, img_channels, num_features=64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                img_channels,\n                num_features,\n                kernel_size=7,\n                stride=1,\n                padding=3,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n\n        self.down_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features, num_features * 2, kernel_size=3, stride=2, padding=1\n                ),\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 4,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n            ]\n        )\n\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n        )\n\n        self.up_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features * 4,\n                    num_features * 2,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 1,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n            ]\n        )\n\n        self.last = nn.Conv2d(\n            num_features * 1,\n            img_channels,\n            kernel_size=7,\n            stride=1,\n            padding=3,\n            padding_mode=\"reflect\",\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.297927Z","iopub.execute_input":"2025-02-05T18:48:50.298381Z","iopub.status.idle":"2025-02-05T18:48:50.31493Z","shell.execute_reply.started":"2025-02-05T18:48:50.298343Z","shell.execute_reply":"2025-02-05T18:48:50.314022Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Discriminator**","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                out_channels,\n                4,\n                stride,\n                1,\n                bias=True,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                features[0],\n                kernel_size=4,\n                stride=2,\n                padding=1,\n                padding_mode=\"reflect\",\n            ),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(\n                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)\n            )\n            in_channels = feature\n        layers.append(\n            nn.Conv2d(\n                in_channels,\n                1,\n                kernel_size=4,\n                stride=1,\n                padding=1,\n                padding_mode=\"reflect\",\n            )\n        )\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.318066Z","iopub.execute_input":"2025-02-05T18:48:50.318337Z","iopub.status.idle":"2025-02-05T18:48:50.328951Z","shell.execute_reply.started":"2025-02-05T18:48:50.318313Z","shell.execute_reply":"2025-02-05T18:48:50.328193Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Buliding**","metadata":{}},{"cell_type":"code","source":"class HumanToPortrait(Dataset):\n    def __init__(self, human_path, portrait_paths, transform=None):\n        self.human_path = human_path\n        self.portrait_paths = portrait_paths\n        self.transform = transform\n\n        self.human_images = os.listdir(human_path)\n        self.portrait_images = []\n        for path in portrait_paths:\n            self.portrait_images.extend(os.listdir(path))\n\n        self.length_dataset = max(len(self.human_images), len(self.portrait_images))\n        self.human_len = len(self.human_images)\n        self.portrait_len = len(self.portrait_images)\n\n    def __len__(self):\n        return self.length_dataset\n\n    def __getitem__(self, index):\n        human_img = self.human_images[index % self.human_len]\n        portrait_img = self.portrait_images[index % self.portrait_len]\n\n        human_img = np.array(Image.open(os.path.join(self.human_path, human_img)).convert(\"RGB\"))\n        portrait_img_path = self.portrait_images[index % self.portrait_len]\n        for path in self.portrait_paths:\n            if portrait_img_path in os.listdir(path):\n                portrait_img = np.array(Image.open(os.path.join(path, portrait_img_path)).convert(\"RGB\"))\n                break\n\n        if self.transform:\n            augmentations = self.transform(image=human_img, image0=portrait_img)\n            human_img = augmentations[\"image\"]\n            portrait_img = augmentations[\"image0\"]\n\n        return human_img, portrait_img","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.33008Z","iopub.execute_input":"2025-02-05T18:48:50.330415Z","iopub.status.idle":"2025-02-05T18:48:50.343543Z","shell.execute_reply.started":"2025-02-05T18:48:50.33038Z","shell.execute_reply":"2025-02-05T18:48:50.342719Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Loading Data**","metadata":{}},{"cell_type":"code","source":"# Create the dataset object\ndataset = HumanToPortrait(\n    human_path=HUMAN_PATH,\n    portrait_paths=PORTRAIT_PATHS,\n    transform=transforms\n)\n\n# Create the DataLoader\nloader = DataLoader(\n    dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=NUM_WORKERS,\n    pin_memory=True\n)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:50.344516Z","iopub.execute_input":"2025-02-05T18:48:50.344803Z","iopub.status.idle":"2025-02-05T18:48:51.338791Z","shell.execute_reply.started":"2025-02-05T18:48:50.344778Z","shell.execute_reply":"2025-02-05T18:48:51.337843Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Training Function**","metadata":{}},{"cell_type":"code","source":"def train_fn(\n    disc_H, disc_P, gen_P, gen_H, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler, epoch\n):\n    H_reals = 0\n    H_fakes = 0\n    P_reals = 0\n    P_fakes = 0\n    loop = tqdm(loader, leave=True)\n\n    for idx, (human, portrait) in enumerate(loop):\n\n        human = human.to(DEVICE)\n        portrait = portrait.to(DEVICE)\n\n        # Train Discriminators H and P\n        with torch.cuda.amp.autocast():\n            fake_portrait = gen_P(human)\n            D_P_real = disc_P(portrait)\n            D_P_fake = disc_P(fake_portrait.detach())\n            P_reals += D_P_real.mean().item()\n            P_fakes += D_P_fake.mean().item()\n            D_P_real_loss = mse(D_P_real, torch.ones_like(D_P_real))\n            D_P_fake_loss = mse(D_P_fake, torch.zeros_like(D_P_fake))\n            D_P_loss = D_P_real_loss + D_P_fake_loss\n\n            fake_human = gen_H(portrait)\n            D_H_real = disc_H(human)\n            D_H_fake = disc_H(fake_human.detach())\n            D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))\n            D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))\n            D_H_loss = D_H_real_loss + D_H_fake_loss\n\n            # put it together\n            D_loss = (D_P_loss + D_H_loss) / 2\n\n        opt_disc.zero_grad()\n        d_scaler.scale(D_loss).backward()\n        d_scaler.step(opt_disc)\n        d_scaler.update()\n\n        # Train Generators H and P\n        with torch.cuda.amp.autocast():\n            # adversarial loss for both generators\n            D_P_fake = disc_P(fake_portrait)\n            D_H_fake = disc_H(fake_human)\n            loss_G_P = mse(D_P_fake, torch.ones_like(D_P_fake))\n            loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))\n\n            # cycle loss\n            cycle_human = gen_H(fake_portrait)\n            cycle_portrait = gen_P(fake_human)\n            cycle_human_loss = l1(human, cycle_human)\n            cycle_portrait_loss = l1(portrait, cycle_portrait)\n\n            # identity loss (remove these for efficiency if you set lambda_identity=0)\n            identity_human = gen_H(human)\n            identity_portrait = gen_P(portrait)\n            identity_human_loss = l1(human, identity_human)\n            identity_portrait_loss = l1(portrait, identity_portrait)\n\n            # add all together\n            G_loss = (\n                loss_G_H\n                + loss_G_P\n                + cycle_human_loss * LAMBDA_CYCLE\n                + cycle_portrait_loss * LAMBDA_CYCLE\n                + identity_human_loss * LAMBDA_IDENTITY\n                + identity_portrait_loss * LAMBDA_IDENTITY\n            )\n\n        opt_gen.zero_grad()\n        g_scaler.scale(G_loss).backward()\n        g_scaler.step(opt_gen)\n        g_scaler.update()\n\n        if idx % 500 == 0:\n            save_image(fake_portrait * 0.5 + 0.5, f\"/kaggle/working/portrait/FAKE_PORTRAIT_index_{idx}_epoch_{epoch}.png\")\n            save_image(fake_human * 0.5 + 0.5, f\"/kaggle/working/human/FAKE_HUMAN_index_{idx}_epoch_{epoch}.png\")\n\n        loop.set_postfix(P_real=P_reals / (idx + 1), P_fake=P_fakes / (idx + 1))","metadata":{"execution":{"iopub.status.busy":"2025-02-05T18:48:51.339941Z","iopub.execute_input":"2025-02-05T18:48:51.340254Z","iopub.status.idle":"2025-02-05T18:48:51.352697Z","shell.execute_reply.started":"2025-02-05T18:48:51.340229Z","shell.execute_reply":"2025-02-05T18:48:51.351754Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef display_image(image_tensor):\n    image = image_tensor.cpu().detach().numpy().transpose(1, 2, 0)\n    image = (image * 0.5) + 0.5  # Unnormalize\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n\ndef main():\n    disc_H = Discriminator(in_channels=3).to(DEVICE)\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # Wrap models with DataParallel\n    disc_H = torch.nn.DataParallel(disc_H)\n    disc_P = torch.nn.DataParallel(disc_P)\n    gen_P = torch.nn.DataParallel(gen_P)\n    gen_H = torch.nn.DataParallel(gen_H)\n\n    opt_disc = optim.Adam(\n        list(disc_H.parameters()) + list(disc_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_P.parameters()) + list(gen_H.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n    dataset = HumanToPortrait(\n        human_path=HUMAN_PATH,\n        portrait_paths=PORTRAIT_PATHS,\n        transform=transforms,\n    )\n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n    )\n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n\n    # Create directories if they don't exist\n    os.makedirs(\"/kaggle/working/portrait\", exist_ok=True)\n    os.makedirs(\"/kaggle/working/human\", exist_ok=True)\n\n    for epoch in range(NUM_EPOCHS):\n        train_fn(\n            disc_H,\n            disc_P,\n            gen_P,\n            gen_H,\n            loader,\n            opt_disc,\n            opt_gen,\n            L1,\n            mse,\n            d_scaler,\n            g_scaler,\n            epoch\n        )\n        save_checkpoint(gen_H, opt_gen, filename=f\"/kaggle/working/gen_H_epoch_{epoch}.pth\")\n        save_checkpoint(gen_P, opt_gen, filename=f\"/kaggle/working/gen_P_epoch_{epoch}.pth\")\n        save_checkpoint(disc_H, opt_disc, filename=f\"/kaggle/working/disc_H_epoch_{epoch}.pth\")\n        save_checkpoint(disc_P, opt_disc, filename=f\"/kaggle/working/disc_P_epoch_{epoch}.pth\")\n\n        # Generate and display an image\n        with torch.no_grad():\n            random_human_image = dataset[0][0].unsqueeze(0).to(DEVICE)  # Take the first image from the dataset\n            fake_portrait = gen_P(random_human_image)\n            display_image(fake_portrait[0])\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:51.353835Z","iopub.execute_input":"2025-02-05T18:48:51.354163Z","iopub.status.idle":"2025-02-05T18:48:58.585877Z","shell.execute_reply.started":"2025-02-05T18:48:51.354112Z","shell.execute_reply":"2025-02-05T18:48:58.58443Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\nimport torch.optim as optim\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nimport os\n\ndef display_image(image_tensor):\n    image = image_tensor.cpu().detach().numpy().transpose(1, 2, 0)\n    image = (image * 0.5) + 0.5  # Unnormalize\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(f\"=> Loading checkpoint from {checkpoint_file}\")\n    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n    # If you want to change the learning rate\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\ndef main():\n    print(\"Initializing models and optimizers...\")\n    disc_H = Discriminator(in_channels=3).to(DEVICE)\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # Wrap models with DataParallel\n    disc_H = torch.nn.DataParallel(disc_H)\n    disc_P = torch.nn.DataParallel(disc_P)\n    gen_P = torch.nn.DataParallel(gen_P)\n    gen_H = torch.nn.DataParallel(gen_H)\n\n    opt_disc = optim.Adam(\n        list(disc_H.parameters()) + list(disc_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_P.parameters()) + list(gen_H.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n    print(\"Loading dataset...\")\n    dataset = HumanToPortrait(\n        human_path=HUMAN_PATH,\n        portrait_paths=PORTRAIT_PATHS,\n        transform=transforms,\n    )\n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n    )\n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n\n    # Load checkpoints if available\n    if os.path.exists(\"gen_H_epoch_16.pth\"):\n        load_checkpoint(\"gen_H_epoch_16.pth\", gen_H, opt_gen, LEARNING_RATE)\n    if os.path.exists(\"gen_P_epoch_16.pth\"):\n        load_checkpoint(\"gen_P_epoch_16.pth\", gen_P, opt_gen, LEARNING_RATE)\n    if os.path.exists(\"disc_H_epoch_16.pth\"):\n        load_checkpoint(\"disc_H_epoch_16.pth\", disc_H, opt_disc, LEARNING_RATE)\n    if os.path.exists(\"disc_P_epoch_16.pth\"):\n        load_checkpoint(\"disc_P_epoch_16.pth\", disc_P, opt_disc, LEARNING_RATE)\n\n    # Create directories if they don't exist\n    os.makedirs(\"/kaggle/working/portrait\", exist_ok=True)\n    os.makedirs(\"/kaggle/working/human\", exist_ok=True)\n\n    print(\"Starting training loop...\")\n    for epoch in range(17, NUM_EPOCHS):  # Start from epoch 17\n        print(f\"Epoch {epoch}/{NUM_EPOCHS}\")\n        train_fn(\n            disc_H,\n            disc_P,\n            gen_P,\n            gen_H,\n            loader,\n            opt_disc,\n            opt_gen,\n            L1,\n            mse,\n            d_scaler,\n            g_scaler,\n            epoch\n        )\n        print(f\"Saving checkpoints for epoch {epoch}...\")\n        save_checkpoint(gen_H, opt_gen, filename=f\"/kaggle/working/gen_H_epoch_{epoch}.pth\")\n        save_checkpoint(gen_P, opt_gen, filename=f\"/kaggle/working/gen_P_epoch_{epoch}.pth\")\n        save_checkpoint(disc_H, opt_disc, filename=f\"/kaggle/working/disc_H_epoch_{epoch}.pth\")\n        save_checkpoint(disc_P, opt_disc, filename=f\"/kaggle/working/disc_P_epoch_{epoch}.pth\")\n\n        # Generate and display an image\n        with torch.no_grad():\n            random_human_image = dataset[0][0].unsqueeze(0).to(DEVICE)  # Take the first image from the dataset\n            fake_portrait = gen_P(random_human_image)\n            display_image(fake_portrait[0])\n\n    print(\"Training completed.\")\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:58.587126Z","iopub.status.idle":"2025-02-05T18:48:58.58749Z","shell.execute_reply.started":"2025-02-05T18:48:58.587334Z","shell.execute_reply":"2025-02-05T18:48:58.58735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"            random_human_image = dataset[0][0].unsqueeze(0).to(DEVICE)  # Take the first image from the dataset\n            fake_portrait = gen_P(random_human_image)\n            display_image(fake_portrait[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:58.588432Z","iopub.status.idle":"2025-02-05T18:48:58.588709Z","shell.execute_reply.started":"2025-02-05T18:48:58.588576Z","shell.execute_reply":"2025-02-05T18:48:58.588589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Sau khi huấn luyện mô hình\nmodel.save('/kaggle/working/model.h5')  # Lưu mô hình vào thư mục /kaggle/working/\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T18:48:58.589886Z","iopub.status.idle":"2025-02-05T18:48:58.590207Z","shell.execute_reply.started":"2025-02-05T18:48:58.590036Z","shell.execute_reply":"2025-02-05T18:48:58.590051Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Using Model","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch\nimport matplotlib.pyplot as plt\nfrom torchvision import transforms\nfrom PIL import Image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T04:49:29.234279Z","iopub.execute_input":"2025-02-09T04:49:29.234995Z","iopub.status.idle":"2025-02-09T04:49:29.238915Z","shell.execute_reply.started":"2025-02-09T04:49:29.234962Z","shell.execute_reply":"2025-02-09T04:49:29.238064Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Generator**","metadata":{}},{"cell_type":"code","source":"# Cấu hình thiết bị (CUDA nếu có, CPU nếu không)\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Đường dẫn tới các file checkpoint\nCHECKPOINT_PATH = '/kaggle/input/cyclegan-final/pytorch/default/1/'\n\nclass ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity(),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\nclass Generator(nn.Module):\n    def __init__(self, img_channels, num_features=64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                img_channels,\n                num_features,\n                kernel_size=7,\n                stride=1,\n                padding=3,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n\n        self.down_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features, num_features * 2, kernel_size=3, stride=2, padding=1\n                ),\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 4,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n            ]\n        )\n\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n        )\n\n        self.up_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features * 4,\n                    num_features * 2,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 1,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n            ]\n        )\n\n        self.last = nn.Conv2d(\n            num_features * 1,\n            img_channels,\n            kernel_size=7,\n            stride=1,\n            padding=3,\n            padding_mode=\"reflect\",\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T04:49:33.678529Z","iopub.execute_input":"2025-02-09T04:49:33.679335Z","iopub.status.idle":"2025-02-09T04:49:33.691991Z","shell.execute_reply.started":"2025-02-09T04:49:33.679299Z","shell.execute_reply":"2025-02-09T04:49:33.691236Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Discriminator**","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                out_channels,\n                4,\n                stride,\n                1,\n                bias=True,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                features[0],\n                kernel_size=4,\n                stride=2,\n                padding=1,\n                padding_mode=\"reflect\",\n            ),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(\n                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)\n            )\n            in_channels = feature\n        layers.append(\n            nn.Conv2d(\n                in_channels,\n                1,\n                kernel_size=4,\n                stride=1,\n                padding=1,\n                padding_mode=\"reflect\",\n            )\n        )\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T04:49:39.419955Z","iopub.execute_input":"2025-02-09T04:49:39.420753Z","iopub.status.idle":"2025-02-09T04:49:39.428522Z","shell.execute_reply.started":"2025-02-09T04:49:39.420715Z","shell.execute_reply":"2025-02-09T04:49:39.427581Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Load Model**","metadata":{}},{"cell_type":"code","source":"# Tải checkpoint và khôi phục mô hình\ndef load_model(model, checkpoint_path):\n    checkpoint = torch.load(checkpoint_path, map_location=DEVICE)\n\n    # Loại bỏ tiền tố 'module.' nếu có\n    state_dict = checkpoint['state_dict']\n    new_state_dict = {}\n    for key, value in state_dict.items():\n        new_key = key.replace('module.', '')  # Loại bỏ tiền tố 'module.'\n        new_state_dict[new_key] = value\n    \n    # Tải lại state_dict đã được chỉnh sửa vào mô hình\n    model.load_state_dict(new_state_dict)\n    model.eval()  # Đặt mô hình vào chế độ inference (chế độ dự đoán)\n    return model\n\n# Hàm hiển thị ảnh\ndef display_image(image_tensor):\n    image = image_tensor.cpu().detach().numpy().transpose(1, 2, 0)\n    image = (image * 0.5) + 0.5  # Unnormalize\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n\n# Định nghĩa các mô hình\ndisc_H = Discriminator(in_channels=3).to(DEVICE)\ndisc_P = Discriminator(in_channels=3).to(DEVICE)\ngen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\ngen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n# Load các mô hình từ checkpoint\ngen_H = load_model(gen_H, CHECKPOINT_PATH + 'gen_H_epoch_18.pth')\ngen_P = load_model(gen_P, CHECKPOINT_PATH + 'gen_P_epoch_18.pth')\n\n# Đọc và tiền xử lý ảnh đầu vào\ndef process_image(image_path):\n    image = Image.open(image_path).convert(\"RGB\")\n    transform = transforms.Compose([\n        transforms.Resize((256, 256)),  # Resize ảnh nếu cần\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Chuẩn hóa ảnh\n    ])\n    image = transform(image).unsqueeze(0).to(DEVICE)  # Thêm batch dimension và chuyển sang device\n    return image\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T04:51:59.394981Z","iopub.execute_input":"2025-02-09T04:51:59.395332Z","iopub.status.idle":"2025-02-09T04:51:59.999687Z","shell.execute_reply.started":"2025-02-09T04:51:59.395301Z","shell.execute_reply":"2025-02-09T04:51:59.998729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Đường dẫn tới ảnh chân dung đầu vào\ninput_image_path = '/kaggle/input/my-avatar/vanmanh.pro.jpg'  # Thay đường dẫn ảnh của bạn tại đây\n\n# Tiền xử lý ảnh đầu vào\ninput_image = process_image(input_image_path)\n\n# Tạo ảnh vẽ từ ảnh chân dung bằng mô hình gen_P\nwith torch.no_grad():\n    output_image = gen_P(input_image)\n\n# Hiển thị ảnh kết quả\ndisplay_image(output_image[0])  # output_image có thể là batch, lấy phần tử đầu tiên","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T05:11:19.286448Z","iopub.execute_input":"2025-02-09T05:11:19.287156Z","iopub.status.idle":"2025-02-09T05:11:19.533842Z","shell.execute_reply.started":"2025-02-09T05:11:19.287121Z","shell.execute_reply":"2025-02-09T05:11:19.533074Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport matplotlib.pyplot as plt\nfrom torchvision import transforms\nfrom PIL import Image\n\n# Định nghĩa hàm tiền xử lý ảnh\ndef process_image(image_path):\n    transform = transforms.Compose([\n        transforms.Resize((256, 256)),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5,), (0.5,))\n    ])\n    image = Image.open(image_path).convert(\"RGB\")\n    image = transform(image).unsqueeze(0)  # Thêm batch dimension\n    return image.to(DEVICE)\n\n# Định nghĩa hàm hiển thị ảnh\ndef display_images(original, transformed):\n    fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n    \n    # Ảnh gốc\n    ax[0].imshow(original.permute(1, 2, 0).cpu().numpy() * 0.5 + 0.5)  # Chuẩn hóa ngược\n    ax[0].set_title(\"Ảnh Gốc\")\n    ax[0].axis(\"off\")\n    \n    # Ảnh sau chuyển đổi\n    ax[1].imshow(transformed.permute(1, 2, 0).cpu().detach().numpy() * 0.5 + 0.5)\n    ax[1].set_title(\"Ảnh Chuyển Đổi\")\n    ax[1].axis(\"off\")\n    \n    plt.show()\n\n# Đường dẫn tới thư mục chứa ảnh đầu vào\ninput_folder = '/kaggle/input/human-faces/Humans/'  # Thay bằng đường dẫn của bạn\n\n# Lặp qua tất cả các ảnh trong thư mục\nfor filename in os.listdir(input_folder):\n    if filename.endswith(('.jpg', '.jpeg', '.png')):\n        image_path = os.path.join(input_folder, filename)\n        input_image = process_image(image_path)\n        \n        # Chuyển đổi ảnh bằng mô hình\n        with torch.no_grad():\n            output_image = gen_P(input_image)[0]  # Lấy ảnh đầu tiên trong batch\n        \n        # Hiển thị ảnh gốc và ảnh đã chuyển đổi\n        display_images(input_image[0], output_image)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-09T04:53:12.699726Z","iopub.execute_input":"2025-02-09T04:53:12.700094Z","iopub.status.idle":"2025-02-09T04:53:42.075806Z","shell.execute_reply.started":"2025-02-09T04:53:12.70006Z","shell.execute_reply":"2025-02-09T04:53:42.074331Z"}},"outputs":[],"execution_count":null}]}