{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.12","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":284644,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":243933,"modelId":265550}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\n\nbase_dir = \"/kaggle/\"\n\ninput_dir = os.path.join(base_dir, \"input/gan-getting-started/\")\noutput_dir = os.path.join(base_dir, \"data/painter/output\")\nworking_dir = os.path.join(base_dir, \"data/painter/working\")\n\nphotos_path = os.path.join(input_dir, \"photo_jpg\")\nmonet_path = os.path.join(input_dir, \"monet_jpg\")","metadata":{"execution":{"iopub.status.busy":"2025-03-13T09:07:12.722778Z","iopub.execute_input":"2025-03-13T09:07:12.723098Z","iopub.status.idle":"2025-03-13T09:07:12.727593Z","shell.execute_reply.started":"2025-03-13T09:07:12.723076Z","shell.execute_reply":"2025-03-13T09:07:12.726629Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import random\nimport re\nimport shutil\nimport zipfile\nfrom itertools import chain\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.utils as vutils\nfrom PIL import Image\nfrom torch.optim import lr_scheduler\nfrom torch.utils.data import DataLoader\nfrom torchvision import transforms\n\nimport os\nfrom typing import Optional\n\nimport torch\nfrom PIL import Image\nfrom torch.utils.data import DataLoader, Dataset, Subset\n\nimport os\nimport time\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.utils as vutils\nfrom torch.profiler import ProfilerActivity, profile, record_function\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2025-03-13T09:09:36.090502Z","iopub.execute_input":"2025-03-13T09:09:36.090945Z","iopub.status.idle":"2025-03-13T09:09:36.096801Z","shell.execute_reply.started":"2025-03-13T09:09:36.090916Z","shell.execute_reply":"2025-03-13T09:09:36.095776Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Setting up","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 1\nIMG_SIZE = 256\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nNUM_WORKERS = 4\nRANDOM_SEED = 42","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:12.748377Z","iopub.execute_input":"2025-03-13T09:07:12.748731Z","iopub.status.idle":"2025-03-13T09:07:12.763605Z","shell.execute_reply.started":"2025-03-13T09:07:12.748703Z","shell.execute_reply":"2025-03-13T09:07:12.762878Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Random seeds","metadata":{}},{"cell_type":"code","source":"torch.manual_seed(RANDOM_SEED)\ntorch.cuda.manual_seed_all(RANDOM_SEED)\nnp.random.seed(RANDOM_SEED)\nrandom.seed(RANDOM_SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:12.764801Z","iopub.execute_input":"2025-03-13T09:07:12.764994Z","iopub.status.idle":"2025-03-13T09:07:12.779902Z","shell.execute_reply.started":"2025-03-13T09:07:12.764977Z","shell.execute_reply":"2025-03-13T09:07:12.779271Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"Let\\`s inspect the data a little bit. In `monet_jpg` directory we have jpg images of Monet\\`s paintings. Also we have different real photos in `photo_jpg` directory.","metadata":{}},{"cell_type":"code","source":"image = Image.open(os.path.join(monet_path, \"0bd913dbc7.jpg\"))\nplt.imshow(image)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:12.781110Z","iopub.execute_input":"2025-03-13T09:07:12.781388Z","iopub.status.idle":"2025-03-13T09:07:13.014763Z","shell.execute_reply.started":"2025-03-13T09:07:12.781361Z","shell.execute_reply":"2025-03-13T09:07:13.013923Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Let\\`s take a look at our image more closely. We have 256x256x3 arrays (256 pixels in height, 256 pixels in width and 3 channels - RGB). Also we can see, that images have values from 0 to 255 (8-bit images).","metadata":{}},{"cell_type":"code","source":"image_data = np.asarray(image)\nprint(f\"Shape: {image_data.shape}\")\nprint(f\"Data type: {image_data.dtype}\")\nprint(f\"Min: {image_data.min()}, Max: {image_data.max()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:13.015583Z","iopub.execute_input":"2025-03-13T09:07:13.015813Z","iopub.status.idle":"2025-03-13T09:07:13.020906Z","shell.execute_reply.started":"2025-03-13T09:07:13.015787Z","shell.execute_reply":"2025-03-13T09:07:13.020041Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"So let\\`s take a look at 25 random images from each domain.","metadata":{}},{"cell_type":"code","source":"def show_images_from_path(images_directory, num_images=25):\n    images = os.listdir(images_directory)\n    images = [os.path.join(images_directory, image) for image in images]\n    images = random.sample(images, num_images)\n\n    fig, axes = plt.subplots(5, 5, figsize=(20, 20))\n    axes = axes.flatten()\n    for i, image_path in enumerate(images):\n        image = Image.open(image_path)\n        axes[i].imshow(image)\n        axes[i].axis(\"off\")\n    plt.tight_layout()\n    plt.show()\n\nshow_images_from_path(monet_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:13.021765Z","iopub.execute_input":"2025-03-13T09:07:13.022068Z","iopub.status.idle":"2025-03-13T09:07:16.474117Z","shell.execute_reply.started":"2025-03-13T09:07:13.022038Z","shell.execute_reply":"2025-03-13T09:07:16.471591Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"For sure we can see, that Cloude Monet paintings are more colorful and have more brush strokes, than real photos. Also we can see, that Monet\\`s paintings contain a lot of landscapes.","metadata":{}},{"cell_type":"code","source":"show_images_from_path(photos_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:16.475302Z","iopub.execute_input":"2025-03-13T09:07:16.475662Z","iopub.status.idle":"2025-03-13T09:07:19.370533Z","shell.execute_reply.started":"2025-03-13T09:07:16.475636Z","shell.execute_reply":"2025-03-13T09:07:19.369101Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"As we can see in real photos dataset we also have mostly landscapes, but also we can sometimes observe people, animals, buildings and other objects.","metadata":{}},{"cell_type":"markdown","source":"## Transformation pipeline and data augmentation","metadata":{}},{"cell_type":"markdown","source":"For augmentation we will only use horizontal flip, because we don't want to distort the images to much. We also normalize our images to have values between -1 and 1. For validation set we will only normalize images.","metadata":{}},{"cell_type":"code","source":"train_transform = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.RandomHorizontalFlip(p=0.5),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]\n)\nval_transform = test_transform = transforms.Compose(\n    [\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.373421Z","iopub.execute_input":"2025-03-13T09:07:19.373694Z","iopub.status.idle":"2025-03-13T09:07:19.378777Z","shell.execute_reply.started":"2025-03-13T09:07:19.373671Z","shell.execute_reply":"2025-03-13T09:07:19.377792Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Let\\`s test our transformation on `image_data`.","metadata":{}},{"cell_type":"code","source":"transformed_image = train_transform(image)\nprint(f\"Shape: {transformed_image.shape}\")\nprint(f\"Data type: {transformed_image.dtype}\")\nprint(f\"Min: {transformed_image.min()}, Max: {transformed_image.max()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.380192Z","iopub.execute_input":"2025-03-13T09:07:19.380392Z","iopub.status.idle":"2025-03-13T09:07:19.397343Z","shell.execute_reply.started":"2025-03-13T09:07:19.380374Z","shell.execute_reply":"2025-03-13T09:07:19.396499Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Creating data loaders","metadata":{}},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    def __init__(\n        self,\n        x_images_paths: str,\n        y_images_paths: str,\n        transformations: Optional[callable] = None,\n    ):\n        super().__init__()\n        self.root_x = x_images_paths\n        self.root_y = y_images_paths\n\n        self.transform = transformations\n\n        self.x_images = os.listdir(x_images_paths)\n        self.y_images = os.listdir(y_images_paths)\n\n        self.x_length = len(self.x_images)\n        self.y_length = len(self.y_images)\n\n        # amount of images from both domains can be different\n        # we set dataset length to the maximum amount of images\n        self.length_dataset = min(self.x_length, self.y_length)\n\n    def __getitem__(self, idx):\n        # we should use modulo to avoid index out of range for\n        # domain where length < length_dataset\n        x_image = self.x_images[idx % self.x_length]\n        y_image = self.y_images[idx % self.y_length]\n\n        x_image = Image.open(os.path.join(self.root_x, x_image))\n        y_image = Image.open(os.path.join(self.root_y, y_image))\n\n        x_image = x_image.convert(\"RGB\")\n        y_image = y_image.convert(\"RGB\")\n\n        if self.transform:\n            x_image = self.transform(x_image)\n            y_image = self.transform(y_image)\n\n        return x_image, y_image\n\n    def __len__(self):\n        return self.length_dataset","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.398220Z","iopub.execute_input":"2025-03-13T09:07:19.398434Z","iopub.status.idle":"2025-03-13T09:07:19.419820Z","shell.execute_reply.started":"2025-03-13T09:07:19.398415Z","shell.execute_reply":"2025-03-13T09:07:19.418971Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_loaders(\n    x_path: str,\n    y_path: str,\n    train_transform: callable,\n    val_transform: callable,\n    batch_size: int = BATCH_SIZE,\n    train_split: float = 0.8,\n    num_workers: int = NUM_WORKERS,\n):\n    full_dataset = ImageDataset(x_path, y_path, transformations=None)\n\n    # Compute split sizes\n    train_size = int(train_split * len(full_dataset))\n    val_size = len(full_dataset) - train_size\n\n    # Ensure deterministic split\n    generator = torch.Generator().manual_seed(RANDOM_SEED)\n    train_indices, val_indices = torch.utils.data.random_split(\n        range(len(full_dataset)), [train_size, val_size], generator=generator\n    )\n\n    # Create separate train/val datasets with different transforms\n    train_dataset = Subset(\n        ImageDataset(x_path, y_path, transformations=train_transform), train_indices\n    )\n    val_dataset = Subset(\n        ImageDataset(x_path, y_path, transformations=val_transform), val_indices\n    )\n\n    # Create DataLoaders\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=True,\n        num_workers=num_workers,\n        pin_memory=True,\n    )\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=num_workers,\n        pin_memory=True,\n    )\n\n    return train_loader, val_loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.420583Z","iopub.execute_input":"2025-03-13T09:07:19.420830Z","iopub.status.idle":"2025-03-13T09:07:19.435313Z","shell.execute_reply.started":"2025-03-13T09:07:19.420807Z","shell.execute_reply":"2025-03-13T09:07:19.434609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_loader, val_loader = get_loaders(\n    photos_path,\n    monet_path,\n    train_transform,\n    val_transform,\n    batch_size=BATCH_SIZE,\n    train_split=0.9,\n    num_workers=NUM_WORKERS,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.436076Z","iopub.execute_input":"2025-03-13T09:07:19.436329Z","iopub.status.idle":"2025-03-13T09:07:19.480902Z","shell.execute_reply.started":"2025-03-13T09:07:19.436306Z","shell.execute_reply":"2025-03-13T09:07:19.479997Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Let\\`s test our loaders and see some images from them.","metadata":{}},{"cell_type":"code","source":"def show_sample(dataset, num_samples=1):\n    fig, axes = plt.subplots(num_samples, 2, figsize=(10, 5 * num_samples))\n\n    if num_samples == 1:\n        axes = [axes]\n\n    for i in range(num_samples):\n        idx = torch.randint(0, len(dataset), (1,)).item()\n\n        real_image, monet_image = dataset[idx]\n\n        real_image = (real_image + 1) / 2\n        monet_image = (monet_image + 1) / 2\n\n        real_image = real_image.permute(1, 2, 0).numpy()\n        monet_image = monet_image.permute(1, 2, 0).numpy()\n\n        axes[i][0].imshow(real_image)\n        axes[i][0].set_title(\"Real image\")\n        axes[i][0].axis(\"off\")\n\n        axes[i][1].imshow(monet_image)\n        axes[i][1].set_title(\"Monet image\")\n        axes[i][1].axis(\"off\")\n\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.481831Z","iopub.execute_input":"2025-03-13T09:07:19.482147Z","iopub.status.idle":"2025-03-13T09:07:19.487747Z","shell.execute_reply.started":"2025-03-13T09:07:19.482115Z","shell.execute_reply":"2025-03-13T09:07:19.487010Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"show_sample(train_loader.dataset, num_samples=2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:19.488486Z","iopub.execute_input":"2025-03-13T09:07:19.488747Z","iopub.status.idle":"2025-03-13T09:07:20.472469Z","shell.execute_reply.started":"2025-03-13T09:07:19.488717Z","shell.execute_reply":"2025-03-13T09:07:20.471510Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"## CycleGAN","metadata":{}},{"cell_type":"markdown","source":"We will use CycleGAN model to generate monet paintings from photos.\nThis model consists of two generators and two discriminators.\nIdea behind this model is to learn mapping to transform images from one domain to another and vice versa. [Link to paper](https://arxiv.org/pdf/1703.10593).","metadata":{}},{"cell_type":"markdown","source":"### Generator","metadata":{}},{"cell_type":"markdown","source":"We have generator for our model. The architecture of the generator is the following:\n1. Initial convolutional layer with kernel size 7x7 and stride 1.\n2. Several convolutional blocks of downsampling layers.\n3. Several residual blocks.\n4. Several convolutional blocks of upsampling layers.\n5. Final convolutional layer with kernel size 7x7 and stride 1.","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        is_downsampling: bool = True,\n        use_activation: bool = True,\n        **kwargs,\n    ):\n        super().__init__()\n        layers = []\n\n        if is_downsampling:\n            layers.append(\n                nn.Conv2d(\n                    in_channels,\n                    out_channels,\n                    padding_mode=\"reflect\",\n                    **kwargs,\n                )\n            )\n        else:\n            layers.append(\n                nn.ConvTranspose2d(\n                    in_channels,\n                    out_channels,\n                    **kwargs,\n                )\n            )\n\n        layers.append(nn.InstanceNorm2d(out_channels))\n        layers.append(nn.ReLU(inplace=True) if use_activation else nn.Identity())\n\n        self.conv = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass ResnetBlock(nn.Module):\n\n    def __init__(self, channels: int):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(\n                channels,\n                channels,\n                kernel_size=3,\n                padding=1,\n            ),\n            ConvBlock(\n                channels,\n                channels,\n                use_activation=False,\n                kernel_size=3,\n                padding=1,\n            ),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\n\nclass Generator(nn.Module):\n\n    def __init__(\n        self,\n        img_channels: int = 3,\n        num_of_downsampling: int = 2,\n        num_resnet_blocks=9,\n    ):\n        super().__init__()\n\n        num_features = 64\n\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.ReLU(inplace=True),\n        )\n\n        # Downsampling\n        self.down_blocks = nn.ModuleList()\n        for _ in range(num_of_downsampling):\n            self.down_blocks.append(\n                ConvBlock(\n                    num_features,\n                    num_features * 2,\n                    is_downsampling=True,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                )\n            )\n            num_features *= 2\n\n        # ResNet blocks\n        self.resnet_blocks = nn.Sequential(\n            *[ResnetBlock(num_features) for _ in range(num_resnet_blocks)]\n        )\n\n        # Upsampling\n        self.up_blocks = nn.ModuleList()\n        for _ in range(num_of_downsampling):\n            self.up_blocks.append(\n                ConvBlock(\n                    num_features,\n                    num_features // 2,\n                    is_downsampling=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                )\n            )\n            num_features //= 2\n\n        self.final = nn.Conv2d(\n            num_features,\n            img_channels,\n            kernel_size=7,\n            padding=3,\n            stride=1,\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.resnet_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        x = self.final(x)\n        return torch.tanh(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:20.473370Z","iopub.execute_input":"2025-03-13T09:07:20.473632Z","iopub.status.idle":"2025-03-13T09:07:20.484402Z","shell.execute_reply.started":"2025-03-13T09:07:20.473611Z","shell.execute_reply":"2025-03-13T09:07:20.483579Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"gen = Generator()\nx = torch.randn((1, 3, IMG_SIZE, IMG_SIZE))\ngen_out = gen(x)\nprint(f\"gen_out shape: {gen_out.shape}, gen_in shape: {x.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:20.485205Z","iopub.execute_input":"2025-03-13T09:07:20.485455Z","iopub.status.idle":"2025-03-13T09:07:21.298299Z","shell.execute_reply.started":"2025-03-13T09:07:20.485435Z","shell.execute_reply":"2025-03-13T09:07:21.297506Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We can see, that output of the generator is the same shape as input, because we want generator to change the style of image. We could look at generated image, but for now it will be just noise.","metadata":{}},{"cell_type":"markdown","source":"### Discriminator","metadata":{}},{"cell_type":"markdown","source":"For discriminator we will use PatchGAN discriminator. The architecture of the discriminator is the following:\n1. Initial convolutional layer with kernel size 4x4 and stride 2.\n2. Several convolutional blocks.\n3. Final convolutional layer with kernel size 4x4 and stride 2.\n\nIn the end we will have 1x15x15 output for each image which will be used to calculate adversarial loss.","metadata":{}},{"cell_type":"code","source":"class DiscConvBlock(nn.Module):\n\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int,\n        use_instance_norm: bool = True,\n        relu_slope: float = 0.2,\n    ):\n        super().__init__()\n        layers = [\n            nn.Conv2d(\n                in_channels,\n                out_channels,\n                kernel_size=4,\n                stride=stride,\n                padding=1,\n                bias=True,\n                padding_mode=\"reflect\",\n            )\n        ]\n        if use_instance_norm:\n            layers.append(nn.InstanceNorm2d(out_channels))\n        layers.append(nn.LeakyReLU(relu_slope))\n\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.model(x)\n\n\nclass Discriminator(nn.Module):\n\n    def __init__(\n        self,\n        in_channels=3,\n        num_of_blocks=3,\n        noise_std=0.05,\n    ):\n        super().__init__()\n        self.noise_std = noise_std\n\n        num_features = 64\n\n        self.initial = DiscConvBlock(\n            in_channels,\n            num_features,\n            stride=2,\n            use_instance_norm=False,\n        )\n\n        self.layers = nn.ModuleList()\n        for i in range(num_of_blocks):\n            stride = 1 if i == num_of_blocks else 2\n\n            self.layers.append(\n                DiscConvBlock(\n                    in_channels=num_features,\n                    out_channels=num_features * 2,\n                    stride=stride,\n                    use_instance_norm=True,\n                )\n            )\n\n            num_features *= 2\n\n        self.final = nn.Conv2d(\n            num_features,\n            1,\n            kernel_size=4,\n            stride=1,\n            padding=1,\n            padding_mode=\"reflect\",\n        )\n\n    def forward(self, x):\n        if self.training:\n            noise = torch.randn_like(x) * self.noise_std\n            x = x + noise\n\n        x = self.initial(x)\n        for layer in self.layers:\n            x = layer(x)\n        x = self.final(x)\n        return torch.sigmoid(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.299307Z","iopub.execute_input":"2025-03-13T09:07:21.299660Z","iopub.status.idle":"2025-03-13T09:07:21.307629Z","shell.execute_reply.started":"2025-03-13T09:07:21.299627Z","shell.execute_reply":"2025-03-13T09:07:21.306707Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"disc = Discriminator()\nx = torch.randn((1, 3, IMG_SIZE, IMG_SIZE))\ndisc_out = disc(x)\nprint(f\"disc_out shape: {disc_out.shape}, disc_in shape: {x.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.308433Z","iopub.execute_input":"2025-03-13T09:07:21.308693Z","iopub.status.idle":"2025-03-13T09:07:21.375168Z","shell.execute_reply.started":"2025-03-13T09:07:21.308663Z","shell.execute_reply":"2025-03-13T09:07:21.374211Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Image buffer","metadata":{}},{"cell_type":"markdown","source":"Also authors of the paper used ImageBuffer to store history of generated images rather than the ones produced by the latest generators. Image buffer stores latest images and with some probability it will return old image instead of new one (changing the old image to new inside of buffer). It should help to stabilize training.","metadata":{}},{"cell_type":"code","source":"class ImageBuffer:\n    def __init__(self, max_size=50, p=0.5):\n        self.max_size = max_size\n        self.buffer = []\n        self.probability = p\n\n    def get_images(self, generated_images: torch.Tensor):\n        return_images = []\n        for img in generated_images:\n            img = torch.unsqueeze(img.data, 0)\n\n            # if buffer is not full, add images to it\n            if len(self.buffer) < self.max_size:\n                self.buffer.append(img)\n                return_images.append(img)\n            else:\n                if random.uniform(0, 1) > self.probability:\n                    idx = random.randint(0, self.max_size - 1)\n                    output_image = self.buffer[idx].clone()\n                    self.buffer[idx] = img\n                    return_images.append(output_image)\n                else:\n                    return_images.append(img)\n\n        return_images = torch.cat(return_images, dim=0)\n        return return_images","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.376059Z","iopub.execute_input":"2025-03-13T09:07:21.376328Z","iopub.status.idle":"2025-03-13T09:07:21.382648Z","shell.execute_reply.started":"2025-03-13T09:07:21.376307Z","shell.execute_reply":"2025-03-13T09:07:21.381835Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"buffer = ImageBuffer(2)\nx = torch.randn((10, 3, IMG_SIZE, IMG_SIZE))\nres = buffer.get_images(x)\n\n(res == x).all()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.383407Z","iopub.execute_input":"2025-03-13T09:07:21.383724Z","iopub.status.idle":"2025-03-13T09:07:21.417475Z","shell.execute_reply.started":"2025-03-13T09:07:21.383691Z","shell.execute_reply":"2025-03-13T09:07:21.416831Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Model itself","metadata":{}},{"cell_type":"code","source":"class CycleGAN(nn.Module):\n\n    def __init__(\n        self,\n        generator_X: Generator,\n        generator_Y: Generator,\n        discriminator_X: Discriminator,\n        discriminator_Y: Discriminator,\n        gen_optimizer: optim.Optimizer,\n        disc_optimizer: optim.Optimizer,\n        discriminator_loss_factor: float = 0.5,\n        lambda_cycle: float = 10,\n        lambda_identity: float = 0.5,\n    ):\n        super().__init__()\n        \"\"\"\n        Naming conventions:\n        - Generators: The letter indicates the target domain.\n        - generator_X: Converts images from domain Y to domain X (Y → X)\n        - generator_Y: Converts images from domain X to domain Y (X → Y)\n\n        - Discriminators: The letter indicates the domain of the input image.\n        - discriminator_X: Evaluates images in domain X (real or generated)\n        - discriminator_Y: Evaluates images in domain Y (real or generated)\n        \"\"\"\n        # Generators\n        self.generator_X = generator_X\n        self.generator_Y = generator_Y\n\n        # Discriminators\n        self.discriminator_X = discriminator_X\n        self.discriminator_Y = discriminator_Y\n\n        # Image Buffers\n        self.X_buffer = ImageBuffer(50)\n        self.Y_buffer = ImageBuffer(50)\n\n        # Optimizers\n        self.gen_optimizer = gen_optimizer\n        self.disc_optimizer = disc_optimizer\n\n        self.l1 = nn.L1Loss()\n        self.mse = nn.MSELoss()\n\n        self.gen_scaler = torch.amp.GradScaler()\n        self.disc_scaler = torch.amp.GradScaler()\n\n        self.disc_loss_factor = discriminator_loss_factor\n        self.lambda_cycle = lambda_cycle\n        self.lambda_identity = lambda_cycle * lambda_identity\n\n        self.gen_loss = 0\n        self.disc_loss = 0\n\n    def _train_discriminator(\n        self, real_x: torch.Tensor, real_y: torch.Tensor, device: torch.device\n    ):\n        with torch.amp.autocast(device):\n            # Fake generated images to domain Y\n            fake_y = self.generator_Y(real_x)\n\n            # Getting predictions for real and fake images\n            disc_y_real = self.discriminator_Y(real_y)\n            disc_y_fake = self.discriminator_Y(fake_y.detach())\n\n            # Calculating losses. Real images should be classified as 1 and fake as 0\n            disc_y_real_loss = self.mse(disc_y_real, torch.ones_like(disc_y_real))\n            disc_y_fake_loss = self.mse(disc_y_fake, torch.zeros_like(disc_y_fake))\n\n            disc_y_loss = disc_y_real_loss + disc_y_fake_loss\n\n            # Fake generated images to domain X\n            fake_x = self.generator_X(real_y)\n\n            # Getting predictions for real and fake images\n            disc_x_real = self.discriminator_X(real_x)\n            disc_x_fake = self.discriminator_X(fake_x.detach())\n\n            # Calculating losses. Real images should be classified as 1 and fake as 0\n            disc_x_real_loss = self.mse(disc_x_real, torch.ones_like(disc_x_real))\n            disc_x_fake_loss = self.mse(disc_x_fake, torch.zeros_like(disc_x_fake))\n\n            disc_x_loss = disc_x_real_loss + disc_x_fake_loss\n\n            # Total discriminator loss\n            total_disc_loss = (disc_y_loss + disc_x_loss) * self.disc_loss_factor\n\n        return total_disc_loss, fake_x, fake_y\n\n    def _train_generator(\n        self,\n        real_x: torch.Tensor,\n        real_y: torch.Tensor,\n        fake_x: torch.Tensor,\n        fake_y: torch.Tensor,\n        device: torch.device,\n    ):\n        with torch.amp.autocast(device):\n            # Getting predictions for fake images\n            disc_y_fake = self.discriminator_Y(fake_y)\n            disc_x_fake = self.discriminator_X(fake_x)\n\n            # Adversiral losses: Generators should fool discriminators\n            gen_y_loss = self.mse(disc_y_fake, torch.ones_like(disc_y_fake))\n            gen_x_loss = self.mse(disc_x_fake, torch.ones_like(disc_x_fake))\n\n            # Total adversarial loss\n            total_adv_loss = gen_y_loss + gen_x_loss\n\n            # Cycle consistency loss: Images should be reconstructed to the original domain\n            cycle_y = self.generator_Y(fake_x)\n            cycle_x = self.generator_X(fake_y)\n\n            cycle_y_loss = self.l1(real_y, cycle_y) * self.lambda_cycle\n            cycle_x_loss = self.l1(real_x, cycle_x) * self.lambda_cycle\n\n            # Identity mapping loss: Images should not be changed if they are already in the target domain\n            identity_y = self.generator_Y(real_y)\n            identity_x = self.generator_X(real_x)\n\n            identity_y_loss = self.l1(real_y, identity_y) * self.lambda_identity\n            identity_x_loss = self.l1(real_x, identity_x) * self.lambda_identity\n\n            # add everything together\n            total_gen_loss = (\n                total_adv_loss\n                + cycle_y_loss\n                + cycle_x_loss\n                + identity_y_loss\n                + identity_x_loss\n            )\n\n        return total_gen_loss\n\n    def backward(self):\n        self.disc_optimizer.zero_grad()\n        self.disc_scaler.scale(self.disc_loss).backward()\n        self.disc_scaler.step(self.disc_optimizer)\n        self.disc_scaler.update()\n\n        self.gen_optimizer.zero_grad()\n        self.gen_scaler.scale(self.gen_loss).backward()\n        self.gen_scaler.step(self.gen_optimizer)\n        self.gen_scaler.update()\n\n    def forward(self, real_x, real_y, device):\n        disc_loss, fake_x, fake_y = self._train_discriminator(real_x, real_y, device)\n        gen_loss = self._train_generator(real_x, real_y, fake_x, fake_y, device)\n\n        self.disc_loss = disc_loss\n        self.gen_loss = gen_loss\n\n        return {\n            \"gen_loss\": gen_loss.item(),\n            \"disc_loss\": disc_loss.item(),\n            \"fake_x\": fake_x,\n            \"fake_y\": fake_y,\n        }\n\n    def save_checkpoint(self, epoch, checkpoint_dir):\n        os.makedirs(checkpoint_dir, exist_ok=True)\n        checkpoint = {\n            \"epoch\": epoch,\n            \"generator_X\": self.generator_X.state_dict(),\n            \"generator_Y\": self.generator_Y.state_dict(),\n            \"discriminator_X\": self.discriminator_X.state_dict(),\n            \"discriminator_Y\": self.discriminator_Y.state_dict(),\n            \"gen_optimizer\": self.gen_optimizer.state_dict(),\n            \"disc_optimizer\": self.disc_optimizer.state_dict(),\n        }\n        torch.save(checkpoint, f\"{checkpoint_dir}/cyclegan_epoch_{epoch + 1}.pth\")\n\n    def load_checkpoint(self, checkpoint_path, device=\"cuda\"):\n        checkpoint = torch.load(checkpoint_path, map_location=device)\n\n        self.generator_X.load_state_dict(checkpoint[\"generator_X\"])\n        self.generator_Y.load_state_dict(checkpoint[\"generator_Y\"])\n        self.discriminator_X.load_state_dict(checkpoint[\"discriminator_X\"])\n        self.discriminator_Y.load_state_dict(checkpoint[\"discriminator_Y\"])\n\n        self.gen_optimizer.load_state_dict(checkpoint[\"gen_optimizer\"])\n        self.disc_optimizer.load_state_dict(checkpoint[\"disc_optimizer\"])\n\n        print(\n            f\"Checkpoint loaded from {checkpoint_path}, starting from epoch {checkpoint['epoch'] + 1}\"\n        )\n        return checkpoint[\"epoch\"] + 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.418197Z","iopub.execute_input":"2025-03-13T09:07:21.418401Z","iopub.status.idle":"2025-03-13T09:07:21.435358Z","shell.execute_reply.started":"2025-03-13T09:07:21.418382Z","shell.execute_reply":"2025-03-13T09:07:21.434602Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train(\n    model: CycleGAN,\n    device: torch.device,\n    train_loader: torch.utils.data.DataLoader,\n    val_loader: torch.utils.data.DataLoader,\n    checkpoint_dir: str,\n    num_epochs: int = 100,\n    save_every: int = 10,\n    start_epoch: int = 0,\n    scheduler_gen: torch.optim.lr_scheduler._LRScheduler = None,\n    scheduler_disc: torch.optim.lr_scheduler._LRScheduler = None,\n):\n    for epoch in range(start_epoch, num_epochs):\n        model.train()\n        start_time = time.time()\n\n        loop = tqdm(train_loader, leave=True)\n\n        for idx, (real_X, real_Y) in enumerate(loop):\n            real_X = real_X.to(device)\n            real_Y = real_Y.to(device)\n\n            train_results = model(real_X, real_Y, device)\n            model.backward()\n\n            loop.set_postfix(\n                {\n                    \"Gen Loss\": train_results[\"gen_loss\"],\n                    \"Disc Loss\": train_results[\"disc_loss\"],\n                }\n            )\n            loop.set_description(f\"Epoch [{epoch+1}/{num_epochs}]\")\n\n        model.eval()\n        val_gen_loss, val_disc_loss = 0, 0\n        with torch.no_grad():\n            for real_X, real_Y in val_loader:\n                real_X = real_X.to(device)\n                real_Y = real_Y.to(device)\n\n                val_losses = model(real_X, real_Y, device)\n\n                # save losses for visualization\n                val_gen_loss += val_losses[\"gen_loss\"]\n                val_disc_loss += val_losses[\"disc_loss\"]\n\n            val_gen_loss /= len(val_loader)\n            val_disc_loss /= len(val_loader)\n\n        time_diff = time.time() - start_time\n        print(\n            f\"Epoch [{epoch+1}/{num_epochs}] - {time_diff:.2f}s. \"\n            f\"- Train Gen Loss: {train_results['gen_loss']:.4f},\"\n            f\" Train Disc Loss: {train_results['disc_loss']:.4f},\"\n            f\" Val Gen Loss: {val_gen_loss:.4f},\"\n            f\" Val Disc Loss: {val_disc_loss:.4f}\"\n        )\n\n        if scheduler_gen is not None:\n            scheduler_gen.step(epoch=epoch)\n\n        if scheduler_disc is not None:\n            scheduler_disc.step(epoch=epoch)\n\n        if (epoch + 1) % save_every == 0:\n            model.save_checkpoint(epoch, checkpoint_dir)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.438208Z","iopub.execute_input":"2025-03-13T09:07:21.438406Z","iopub.status.idle":"2025-03-13T09:07:21.458059Z","shell.execute_reply.started":"2025-03-13T09:07:21.438389Z","shell.execute_reply":"2025-03-13T09:07:21.457319Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"markdown","source":"Authors of paper Adam optimizer and trained their network for 100 epochs with learning rate 2e-4. Then they linearly decayed learning rate to 0 over next 100 epochs. We going to do same thing.","metadata":{}},{"cell_type":"code","source":"def lr_lambda(epoch):\n    if epoch < 100:\n        return 1.0\n    else:\n        return 1.0 - (epoch - 100) / 100","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.459079Z","iopub.execute_input":"2025-03-13T09:07:21.459315Z","iopub.status.idle":"2025-03-13T09:07:21.475333Z","shell.execute_reply.started":"2025-03-13T09:07:21.459295Z","shell.execute_reply":"2025-03-13T09:07:21.474603Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"So here is creation of our model. We start with initialization of generators, discriminators and optimizers. And also we create learning rate schedulers.","metadata":{}},{"cell_type":"code","source":"generator_X = Generator()\ngenerator_Y = Generator()\ndiscriminator_X = Discriminator()\ndiscriminator_Y = Discriminator()\n\ngen_params = chain(generator_X.parameters(), generator_Y.parameters())\ndisc_params = chain(discriminator_X.parameters(), discriminator_Y.parameters())\n\ngen_opt = optim.Adam(gen_params, lr=0.0002, betas=(0.5, 0.999))\ndisc_opt = optim.Adam(disc_params, lr=0.0002, betas=(0.5, 0.999))\n\nscheduler_gen = lr_scheduler.LambdaLR(gen_opt, lr_lambda)\nscheduler_disc = lr_scheduler.LambdaLR(disc_opt, lr_lambda)\n\nmodel = CycleGAN(\n    generator_X=generator_X,\n    generator_Y=generator_Y,\n    discriminator_X=discriminator_X,\n    discriminator_Y=discriminator_Y,\n    gen_optimizer=gen_opt,\n    disc_optimizer=disc_opt,\n    discriminator_loss_factor=0.5,\n    lambda_cycle=10,\n    lambda_identity=0.5,\n).to(DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.476120Z","iopub.execute_input":"2025-03-13T09:07:21.476333Z","iopub.status.idle":"2025-03-13T09:07:21.712498Z","shell.execute_reply.started":"2025-03-13T09:07:21.476314Z","shell.execute_reply":"2025-03-13T09:07:21.711877Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"I tried different configurations and found out, that model generates more diverse images if we lowers our `lambda_cycle` parameter at the end. So for this I created function which takes `lambda_cycle` (and count of epochs) as parameter and changes it for model. Also we use checkpoints while training so we can load our model from last checkpoint and continue training.","metadata":{}},{"cell_type":"code","source":"def get_latest_checkpoint(checkpoint_dir, pattern=\"cyclegan_epoch_(\\d+).pth\"):\n    checkpoint_files = [f for f in os.listdir(checkpoint_dir) if re.match(pattern, f)]\n    if not checkpoint_files:\n        return None  # No checkpoint found\n\n    # Extract epoch numbers and find the latest one\n    latest_checkpoint = max(\n        checkpoint_files, key=lambda f: int(re.search(pattern, f).group(1))\n    )\n    return os.path.join(checkpoint_dir, latest_checkpoint)\n\n\ndef train_with_lambda_cycle(\n    model: CycleGAN,\n    lambda_cycle_values,\n    checkpoint_dir,\n    start_epoch=0,\n    checkpoint_path=None,\n    **kwargs,\n):\n    \"\"\"Handles training across multiple lambda_cycle values with automatic checkpoint loading.\"\"\"\n\n    # Load latest checkpoint if none is specified\n    if checkpoint_path is None:\n        checkpoint_path = get_latest_checkpoint(checkpoint_dir)\n\n    if checkpoint_path and os.path.exists(checkpoint_path):\n        start_epoch = model.load_checkpoint(checkpoint_path)\n        print(f\"Resuming from checkpoint: {checkpoint_path} (Epoch {start_epoch})\")\n    else:\n        start_epoch = 0\n        print(\"No checkpoint found. Starting from scratch.\")\n    \n    # Iterate through lambda_cycle values and train\n    for lambda_cycle, next_epoch in lambda_cycle_values:\n        if start_epoch >= next_epoch:\n            print(f\"Skipping lambda_cycle={lambda_cycle} (already trained)\")\n            continue\n        model.lambda_cycle = lambda_cycle\n        print(f\"\\nStarting training with lambda_cycle={lambda_cycle}\")\n\n        train(\n            model=model,\n            checkpoint_dir=checkpoint_dir,\n            start_epoch=start_epoch,\n            num_epochs=next_epoch,\n            **kwargs,\n        )\n\n        start_epoch = next_epoch  # Update for next phase","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:07:21.713199Z","iopub.execute_input":"2025-03-13T09:07:21.713420Z","iopub.status.idle":"2025-03-13T09:07:21.719193Z","shell.execute_reply.started":"2025-03-13T09:07:21.713401Z","shell.execute_reply":"2025-03-13T09:07:21.718596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"lambda_cycle_values = [\n    (10, 180),\n    (5, 200),\n]\n\ntrain_with_lambda_cycle(\n    model=model,\n    train_loader=train_loader,\n    val_loader=val_loader,\n    save_every=5,\n    lambda_cycle_values=lambda_cycle_values,\n    device=DEVICE.type,\n    scheduler_gen=scheduler_gen,\n    scheduler_disc=scheduler_disc,\n    checkpoint_dir=\"checkpoints\",\n    checkpoint_path=os.path.join(base_dir, \"input/cyclegan-weights/pytorch/default/1/cyclegan_epoch_200.pth\")\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:15:54.829821Z","iopub.execute_input":"2025-03-13T09:15:54.830117Z","iopub.status.idle":"2025-03-13T09:15:58.423730Z","shell.execute_reply.started":"2025-03-13T09:15:54.830097Z","shell.execute_reply":"2025-03-13T09:15:58.422984Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Testing","metadata":{}},{"cell_type":"markdown","source":"Let\\`s see what our model able to generate. We will take some random images from validation set and generate monet paintings from them.","metadata":{}},{"cell_type":"code","source":"def generate_monet_images(generator, dataloader, device, num_images=5):\n    generator.eval()\n    images_shown = 0\n    fig, axes = plt.subplots(num_images, 2, figsize=(10, 5 * num_images))\n\n    if num_images == 1:\n        axes = [axes]\n\n    with torch.no_grad():\n        for real_images, monet_images in dataloader:\n            real_images = real_images.to(device)\n            generated_images = generator(real_images)\n\n            if generated_images.dim() == 3:\n                generated_images = generated_images.unsqueeze(0)\n                real_images = real_images.unsqueeze(0)\n\n            real_images = real_images.cpu().permute(0, 2, 3, 1)\n            generated_images = generated_images.cpu().permute(0, 2, 3, 1)\n\n            for j in range(real_images.shape[0]):\n                if images_shown >= num_images:\n                    break\n\n                ax_real = axes[images_shown][0]\n                ax_gen = axes[images_shown][1]\n\n                ax_real.imshow((real_images[j] * 0.5 + 0.5).numpy())\n                ax_real.set_title(\"Real image\")\n                ax_real.axis(\"off\")\n\n                ax_gen.imshow((generated_images[j] * 0.5 + 0.5).numpy())\n                ax_gen.set_title(\"Generated image\")\n                ax_gen.axis(\"off\")\n\n                images_shown += 1\n\n            if images_shown >= num_images:\n                break\n\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:02.029290Z","iopub.execute_input":"2025-03-13T09:16:02.029610Z","iopub.status.idle":"2025-03-13T09:16:02.036641Z","shell.execute_reply.started":"2025-03-13T09:16:02.029586Z","shell.execute_reply":"2025-03-13T09:16:02.035833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_subset = torch.utils.data.Subset(\n    val_loader.dataset, torch.randint(0, len(val_loader.dataset), (5,))\n)\ngenerate_monet_images(model.generator_Y, val_subset, device=\"cuda\", num_images=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:04.052330Z","iopub.execute_input":"2025-03-13T09:16:04.052667Z","iopub.status.idle":"2025-03-13T09:16:06.362419Z","shell.execute_reply.started":"2025-03-13T09:16:04.052643Z","shell.execute_reply":"2025-03-13T09:16:06.361440Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"While it is clearly not the best results we could get, we can see, that model is able to generate some paintings and it`s impressive.","metadata":{}},{"cell_type":"code","source":"def show_generated_images(generator, photo_loader, device, num_images=16):\n    generator.eval()\n\n    with torch.no_grad():\n        real_photos, monet_images = next(iter(photo_loader))\n        real_photos = real_photos[:num_images].to(device)\n        fake_monets = generator(real_photos).cpu()\n\n    grid = vutils.make_grid(fake_monets, nrow=4, normalize=True)\n\n    plt.figure(figsize=(12, 12))\n    plt.imshow(grid.permute(1, 2, 0))\n    plt.axis(\"off\")\n    plt.title(\"Generated Monet Paintings\")\n    plt.show()\n\n\nshow_generated_images(model.generator_Y, train_loader, device=DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:06.364148Z","iopub.execute_input":"2025-03-13T09:16:06.364639Z","iopub.status.idle":"2025-03-13T09:16:06.848674Z","shell.execute_reply.started":"2025-03-13T09:16:06.364590Z","shell.execute_reply":"2025-03-13T09:16:06.847704Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Generate submission","metadata":{}},{"cell_type":"code","source":"def generate_monet_and_save_to_zip(model, dataloader, device, output_zip):\n    model.eval()\n    temp_dir = \"/tmp/temp_monet_images\"\n    os.makedirs(temp_dir, exist_ok=True)\n\n    with torch.no_grad():\n        image_paths = []\n        for idx, real_images in enumerate(dataloader):\n            real_images = real_images.to(device)\n            generated_images = model(real_images)\n\n            if generated_images.dim() == 3:\n                generated_images = generated_images.unsqueeze(0)\n\n            generated_images = generated_images.cpu().permute(\n                0, 2, 3, 1\n            )  # Convert to (batch, H, W, C)\n\n            for j in range(generated_images.shape[0]):\n                image = generated_images[j]\n                image = (image * 0.5 + 0.5).numpy()  # Convert to [0, 1] range\n                image = (image * 255).astype(np.uint8)  # Convert to [0, 255]\n                image = Image.fromarray(image)\n\n                # Unique filename\n                filename = f\"image_{idx * generated_images.shape[0] + j}.jpg\"\n                filepath = os.path.join(temp_dir, filename)\n                image.save(filepath)\n                image_paths.append(filepath)\n\n    # Save all images to a ZIP file\n    with zipfile.ZipFile(output_zip, \"w\") as zipf:\n        for file in image_paths:\n            zipf.write(file, os.path.basename(file))\n\n    shutil.rmtree(temp_dir)\n\n    print(f\"Saved {len(image_paths)} images to {output_zip}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:16.196677Z","iopub.execute_input":"2025-03-13T09:16:16.197012Z","iopub.status.idle":"2025-03-13T09:16:16.203732Z","shell.execute_reply.started":"2025-03-13T09:16:16.196992Z","shell.execute_reply":"2025-03-13T09:16:16.202761Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SingleDomainImageDataset(Dataset):\n\n    def __init__(self, images_paths: str, transformations: Optional[callable] = None):\n        super().__init__()\n        self.root = images_paths\n        self.transform = transformations\n\n        self.images = os.listdir(images_paths)\n        self.length = len(self.images)\n\n    def __len__(self):\n        return self.length\n\n    def __getitem__(self, index):\n        image = self.images[index]\n        image = Image.open(os.path.join(self.root, image))\n        image = image.convert(\"RGB\")\n\n        if self.transform:\n            image = self.transform(image)\n\n        return image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:17.671705Z","iopub.execute_input":"2025-03-13T09:16:17.672011Z","iopub.status.idle":"2025-03-13T09:16:17.677499Z","shell.execute_reply.started":"2025-03-13T09:16:17.671991Z","shell.execute_reply":"2025-03-13T09:16:17.676551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataset = SingleDomainImageDataset(photos_path, transformations=val_transform)\ndataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:18.006268Z","iopub.execute_input":"2025-03-13T09:16:18.006591Z","iopub.status.idle":"2025-03-13T09:16:18.032239Z","shell.execute_reply.started":"2025-03-13T09:16:18.006567Z","shell.execute_reply":"2025-03-13T09:16:18.031579Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"generate_monet_and_save_to_zip(\n    model.generator_Y,\n    dataloader,\n    device=DEVICE,\n    output_zip=\"images.zip\",\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T09:16:18.421362Z","iopub.execute_input":"2025-03-13T09:16:18.421711Z","iopub.status.idle":"2025-03-13T09:16:20.812697Z","shell.execute_reply.started":"2025-03-13T09:16:18.421684Z","shell.execute_reply":"2025-03-13T09:16:20.811497Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}