{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:45:35.685027Z","iopub.execute_input":"2024-11-23T19:45:35.685422Z","iopub.status.idle":"2024-11-23T19:45:44.186365Z","shell.execute_reply.started":"2024-11-23T19:45:35.685361Z","shell.execute_reply":"2024-11-23T19:45:44.184319Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#import important libraries\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport os\n#from google.colab import drive\nimport shutil\nimport sys\nimport cv2\nimport random\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom collections import defaultdict\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:46:40.900334Z","iopub.execute_input":"2024-11-23T19:46:40.900756Z","iopub.status.idle":"2024-11-23T19:46:43.90978Z","shell.execute_reply.started":"2024-11-23T19:46:40.900719Z","shell.execute_reply":"2024-11-23T19:46:43.908892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"KAGGLE_INPUT_PATH='/kaggle/input'\nKAGGLE_WORKING_PATH='/kaggle/working'\nKAGGLE_SYMLINK='kaggle'\n\n!umount /kaggle/input/ 2> /dev/null\nshutil.rmtree('/kaggle/input', ignore_errors=True)\nos.makedirs(KAGGLE_INPUT_PATH, 0o777, exist_ok=True)\nos.makedirs(KAGGLE_WORKING_PATH, 0o777, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:47:30.843318Z","iopub.execute_input":"2024-11-23T19:47:30.843818Z","iopub.status.idle":"2024-11-23T19:47:34.286089Z","shell.execute_reply.started":"2024-11-23T19:47:30.843781Z","shell.execute_reply":"2024-11-23T19:47:34.284415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set seed for reproducibility\ndef seed_everything(seed=888):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:48:03.409416Z","iopub.execute_input":"2024-11-23T19:48:03.409782Z","iopub.status.idle":"2024-11-23T19:48:03.418267Z","shell.execute_reply.started":"2024-11-23T19:48:03.409755Z","shell.execute_reply":"2024-11-23T19:48:03.416832Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"root_path = \"/kaggle/input/gan-getting-started\"\n# Reading img in rgb\nread_img = lambda path: cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:49:22.584206Z","iopub.execute_input":"2024-11-23T19:49:22.584674Z","iopub.status.idle":"2024-11-23T19:49:22.591474Z","shell.execute_reply.started":"2024-11-23T19:49:22.584629Z","shell.execute_reply":"2024-11-23T19:49:22.589898Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Show the number of images\nprint('Number of images in Monet directory: ', len(os.listdir(f'{root_path}/monet_jpg/')))\nprint('Number of images in photo directory: ', len(os.listdir(f'{root_path}/photo_jpg/')))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:49:27.168751Z","iopub.execute_input":"2024-11-23T19:49:27.169111Z","iopub.status.idle":"2024-11-23T19:49:27.177535Z","shell.execute_reply.started":"2024-11-23T19:49:27.169083Z","shell.execute_reply":"2024-11-23T19:49:27.176509Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load a sample image of monet and photo\ndata_path = f\"{root_path}/photo_jpg\"\nsample_photo_path = os.path.join(data_path, os.listdir(data_path)[20])\nprint(sample_photo_path)\nsample_photo = read_img(sample_photo_path)\n\ndata_path = f\"{root_path}/monet_jpg\"\nsample_monet_path = os.path.join(data_path, os.listdir(data_path)[20])\nsample_monet = read_img(sample_monet_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:56:29.467713Z","iopub.execute_input":"2024-11-23T19:56:29.468189Z","iopub.status.idle":"2024-11-23T19:56:29.518085Z","shell.execute_reply.started":"2024-11-23T19:56:29.468153Z","shell.execute_reply":"2024-11-23T19:56:29.5167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Photo\nplt.subplot(121)\nplt.title(\"Photo\")\nplt.imshow(sample_photo)\n\n# Monet\nplt.subplot(122)\nplt.title(\"Monet\")\nplt.imshow(sample_monet)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:56:56.242416Z","iopub.execute_input":"2024-11-23T19:56:56.24284Z","iopub.status.idle":"2024-11-23T19:56:56.731821Z","shell.execute_reply.started":"2024-11-23T19:56:56.242807Z","shell.execute_reply":"2024-11-23T19:56:56.730801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to display sample images\ndef display_sample_images_chanels(file_path, num_samples=1, title=\"\"):\n    fig, ax = plt.subplots(2*num_samples, 4, figsize=(20, 4*2*num_samples))\n    sample_image_path = file_path\n    image_decoded = cv2.imread(sample_image_path)\n    image_decoded = cv2.cvtColor(image_decoded, cv2.COLOR_BGR2RGB )\n    # Original image\n    ax[0][0].imshow(image_decoded)\n    ax[0][0].axis('off')\n    ax[0][0].set_title(f'{title}',color='#8d4908')\n\n    for j in range(3):\n        image_inChanel = image_decoded[:,:,j].reshape(-1)\n        ax[0][j+1].hist(image_inChanel,bins=25, color=['#ad2004', '#007446', '#13358f'][j], alpha=0.3)\n        ax[0][j+1].set_title(f'{title}',color='#8d4908')\n\n    # Gray scale image\n    ax[1][0].imshow(np.mean(image_decoded, axis=2),cmap=\"gray\")\n    ax[1][0].set_title(f'{title}', color='#8d4908')\n\n    # Red, Green, and Blue Chanels Image\n    for k in range(3):\n        image_individual_Chanel= np.full_like(image_decoded, 0)\n        image_individual_Chanel[:, :, k] = image_decoded[:, :, k]\n        ax[1][k+1].imshow(image_individual_Chanel)\n        ax[1][k+1].set_title(f'{title}', color='#8d4908')\n\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:57:26.607087Z","iopub.execute_input":"2024-11-23T19:57:26.60743Z","iopub.status.idle":"2024-11-23T19:57:26.616015Z","shell.execute_reply.started":"2024-11-23T19:57:26.607402Z","shell.execute_reply":"2024-11-23T19:57:26.614652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"display_sample_images_chanels(sample_monet_path, title=\"Monet\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:57:46.896855Z","iopub.execute_input":"2024-11-23T19:57:46.897372Z","iopub.status.idle":"2024-11-23T19:57:48.16176Z","shell.execute_reply.started":"2024-11-23T19:57:46.897332Z","shell.execute_reply":"2024-11-23T19:57:48.160507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# img params\nIMG_H, IMG_W, IMG_C = sample_photo.shape\nprint(f\"IMG_H: {IMG_H}, IMG_W: {IMG_W}, IMG_C: {IMG_C}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:58:27.080415Z","iopub.execute_input":"2024-11-23T19:58:27.080813Z","iopub.status.idle":"2024-11-23T19:58:27.088844Z","shell.execute_reply.started":"2024-11-23T19:58:27.080784Z","shell.execute_reply":"2024-11-23T19:58:27.087059Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Downsample(nn.Module):\n    \"\"\" A simple convolutional block that downsamples feature maps by 2 (stride=2)\n\n    For example\n        input_tensor [256, 256] -> output_tensor [128, 128]\n\n    Padding = same\n    Bias = False (although this will be checked)\n\n    Model\n        Conv2d -> (opt) InstanceNorm2d -> LeakyReLU\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 in_channels,\n                 out_channels,\n                 kernel_size,\n                 norm=True):\n\n        \"\"\" Init downsample module\n\n        ---\n        Parameters\n            in_channels: Number of input channels\n            out_channels: Number of output channels\n            kernel_size: Kernel size\n            norm: Whether to apply normalisation layer\n        \"\"\"\n\n        super().__init__()\n\n        # --- Layers\n        # Convolution Layer\n        self.conv = nn.Conv2d(in_channels=in_channels,\n                              out_channels=out_channels,\n                              kernel_size=kernel_size,\n                              stride=2,  # Downsample by 2\n                              padding=(kernel_size - 1) // 2,  # Same padding\n                              bias=False)  # ToDo: Check if this should be conditioned on norm\n\n        # Normalisation layer (equivalent of GroupNorm where num_groups=out_channels\n        if norm:\n            self.norm = nn.InstanceNorm2d(num_features=out_channels)\n        else:\n            self.norm = None\n\n        # Activation layer\n        self.act = nn.LeakyReLU()\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n\n        \"\"\" Downsample input tensor\n\n        ---\n        Parameters\n            x: Input tensor\n\n        ---\n        Returns\n            out: x downsampled by 2\n\n        \"\"\"\n\n        out = self.conv(x)\n\n        if self.norm:\n            out = self.norm(out)\n\n        out = self.act(out)\n\n        return out\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:58:40.986086Z","iopub.execute_input":"2024-11-23T19:58:40.986421Z","iopub.status.idle":"2024-11-23T19:58:40.994217Z","shell.execute_reply.started":"2024-11-23T19:58:40.986393Z","shell.execute_reply":"2024-11-23T19:58:40.993012Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Upsample(nn.Module):\n    \"\"\" A simple convolutional block that upsamples feature maps by 2 (kernel_size=4, stride=2)\n\n        For example\n            input_tensor [128, 128] -> output_tensor [256, 256]\n\n        Padding = same\n        Bias = False\n        Dropout = 0.5\n\n        Model\n            Conv2d -> InstanceNorm2d -> (opt) Dropout -> ReLU\n\n        \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 in_channels,\n                 out_channels,\n                 kernel_size=4,\n                 dropout=True):\n        \"\"\" Init upsample module\n\n        ---\n        Parameters\n            in_channels: Number of input channels\n            out_channels: Number of output channels\n            kernel_size: Kernel size\n            dropout: Whether to apply dropout\n        \"\"\"\n\n        super().__init__()\n\n        # --- Layers\n        # Transpose Convolution Layer\n        self.conv = nn.ConvTranspose2d(in_channels=in_channels,\n                                       out_channels=out_channels,\n                                       kernel_size=kernel_size,\n                                       stride=2,  # Upsample by 2\n                                       padding=(kernel_size - 1) // 2,  # Same padding\n                                       bias=False)\n\n        # Normalisation layer (equivalent of GroupNorm where num_groups=out_channels\n        self.norm = nn.InstanceNorm2d(num_features=out_channels)\n\n        # Dropout\n        if dropout:\n            self.dropout = nn.Dropout(0.5)\n        else:\n            self.dropout = None\n\n        # Activation layer\n        self.act = nn.ReLU()\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n        \"\"\" Upsample input tensor\n\n        ---\n        Parameters\n            x: Input tensor\n\n        ---\n        Returns\n            out: x upsampled by 2\n\n        \"\"\"\n\n        out = self.conv(x)\n        out = self.norm(out)\n\n        if self.dropout:\n            out = self.dropout(out)\n\n        out = self.act(out)\n\n        return out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:58:59.777086Z","iopub.execute_input":"2024-11-23T19:58:59.777404Z","iopub.status.idle":"2024-11-23T19:58:59.78596Z","shell.execute_reply.started":"2024-11-23T19:58:59.777379Z","shell.execute_reply":"2024-11-23T19:58:59.784465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def conv_weight_init(m):\n    \"\"\" Initialise Conv2D and ConvTranspose2D with N(0, 0.02)\n    \"\"\"\n    if any(isinstance(m, _m) for _m in [nn.Conv2d, nn.ConvTranspose2d]):\n        nn.init.normal_(m.weight, mean=0.0, std=0.02)\n        if m.bias is not None:\n            nn.init.constant_(m.bias, 0)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:59:07.544945Z","iopub.execute_input":"2024-11-23T19:59:07.545366Z","iopub.status.idle":"2024-11-23T19:59:07.552022Z","shell.execute_reply.started":"2024-11-23T19:59:07.545335Z","shell.execute_reply":"2024-11-23T19:59:07.550656Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Generator(nn.Module):\n    \"\"\" A Generator Network\n\n    ---\n    Structure\n        Img -> Encoder -> Decoder -> Img_Pred\n\n        Encoder\n            7 x Downsample blocks\n            Output -> Latent Tensor [512, IMG_H / 128, IMG_W / 128]\n        Decoder\n            6 x Upsample blocks\n            Output -> Upsampled Latent Tensor [64, IMG_H / 2, IMG_W / 2]\n        Final Upsample\n            ConvTranspose2D -> TanH\n            Output -> Target Image [IMG_C, IMG_H, IMG_W]\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self):\n        \"\"\" Init CycleGAN's Generator\n\n        \"\"\"\n        super().__init__()\n\n        # --- Encoder and Decoder\n        self.encoder = self._init_encoder()\n        self.decoder = self._init_decoder()\n\n        # Final Upsample\n        self.out = nn.ConvTranspose2d(in_channels=128,\n                                      out_channels=3,\n                                      kernel_size=4,\n                                      stride=2, padding=1)\n\n        # Maybe softsign could be better than TanH\n        self.act = nn.Tanh()\n\n        # Initialise conv wights with N(0, 0.02)\n        self.apply(conv_weight_init)\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n        \"\"\" Generate x_hat\n\n        ---\n        Parameters\n            x: Input tensor (img)\n\n        ---\n        Returns\n            torch.Tensor: Output tensor (img pred)\n\n        \"\"\"\n\n        # Encode to latent space / 128 #512\n        skips = []  # Skip connections U-Net Style\n        for layer in self.encoder:\n            x = layer(x)\n            skips.append(x)\n\n        # Skip last one, bottom bit, latent space\n        skips = reversed(skips[:-1])\n\n        # Decode form latent space *64 #128\n        for layer, skip in zip(self.decoder, skips):\n            x = layer(x)\n            x = torch.cat((x, skip), dim=1)\n\n        # Upsample so that out.shape == x.shape\n        out = self.act(self.out(x))\n\n        return out\n\n    # --------------------------------------------------------------------------------\n    def _init_encoder(self):\n        \"\"\" A Sequential Encoder with 7 Downsample Blocks each downsampling by 2\n\n        ---\n        Returns\n            nn.Sequential: Encoder\n\n        \"\"\"\n\n        # Settings\n        kernel_size = 4\n        in_channels = 64\n\n        # Final layer will have 512 feature maps and it's final output size is x / 128\n        out_channels = [128, 256, 512, 512, 512, 512]\n\n        # First layer doesn't use norm so just add it now\n        encoder = [Downsample(in_channels=3, out_channels=in_channels,\n                              kernel_size=kernel_size, norm=False)]\n\n        for out_ch in out_channels:\n            # Add layer\n            encoder.append(Downsample(in_channels=in_channels, out_channels=out_ch,\n                                      kernel_size=kernel_size))\n\n            # Update in_channels\n            in_channels = out_ch\n\n        return nn.Sequential(*encoder)\n\n    # --------------------------------------------------------------------------------\n    def _init_decoder(self):\n        \"\"\" A Sequential Decoder with 6 Upsample Blocks each upscaling by 2\n\n        ---\n        Returns\n            nn.Sequential: Decoder\n\n        \"\"\"\n\n        # Settings\n        kernel_size = 4\n        in_channels = [512, 1024, 1024, 1024, 512, 256]\n\n        # Final layer will have the same number of channels as the output of first downsample from encoder\n        out_channels = [512, 512, 512, 256, 128, 64]\n        dropout = [True, True, True, False, False, False]\n\n        decoder = [Upsample(in_channels=in_ch, out_channels=out_ch,\n                            kernel_size=kernel_size, dropout=drop)\n\n                   for in_ch, out_ch, drop in zip(in_channels, out_channels, dropout)]\n\n        return nn.Sequential(*decoder)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:59:16.99129Z","iopub.execute_input":"2024-11-23T19:59:16.991681Z","iopub.status.idle":"2024-11-23T19:59:17.004373Z","shell.execute_reply.started":"2024-11-23T19:59:16.99165Z","shell.execute_reply":"2024-11-23T19:59:17.003212Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass DiscriminatorHistory:\n    \"\"\" Keep a history of fake samples that Discriminator can use to improve its robustness\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 buffer_size=50):\n        \"\"\" Init History module\n\n        ---\n        Parameters\n            buffer_size: Size of the buffer, as per paper the defualt is 50\n\n        The paper says \"We keep an image buffer that stores the 50 previously created images\",\n        however, what's worth noting is that the authors used a batch size of 1 for training,\n        meaning that the network did 50 updates before sampling from history. Worth rethinking\n        whether the buffer_size should be dependant on the batch_size\n\n        \"\"\"\n\n        self.buffer_size = buffer_size\n\n        # Empty Tensor\n        self.buffer = torch.Tensor()\n\n\n    # --------------------------------------------------------------------------------\n    def __call__(self,\n                 x):\n        \"\"\" Forward\n\n        Sample randomly from history buffer and add new sampels\n\n        ---\n        Parameters\n            x: Input Tensor\n\n        ---\n        Returns\n            Tensor with mixed samples from input and history\n\n        \"\"\"\n\n        # Make sure that the buffer is on the same device as x\n        if self.buffer.device != x.device:\n            self.buffer = self.buffer.to(x.device)\n\n        # Check if buffer is filled\n        if len(self.buffer) >= self.buffer_size:\n            # Sample from the buffer\n            sample_size = max(len(x) // 2, 1)\n            history_indices = list(range(0, self.buffer_size))\n            history_get_indices = random.sample(history_indices, k=sample_size)\n            x1 = self.buffer[history_get_indices]\n\n            if sample_size == 1:\n                return x1\n\n            # Sample from the input\n            x_indices = list(range(0, len(x)))\n            x_get_indices = random.sample(x_indices, k=sample_size)\n            x_set_indices = list(set(x_indices) - set(x_get_indices))\n            x2 = x[x_get_indices]\n\n            # Replace buffer with other set of samples\n            self.buffer[history_get_indices] = x[x_set_indices]\n\n            # Concat\n            x = torch.cat((x1, x2), 0)\n\n        # Fill the buffer\n        else:\n            self.buffer = torch.cat((self.buffer, x), 0)\n\n        # Return\n        return x\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:59:32.818768Z","iopub.execute_input":"2024-11-23T19:59:32.819194Z","iopub.status.idle":"2024-11-23T19:59:32.828998Z","shell.execute_reply.started":"2024-11-23T19:59:32.819162Z","shell.execute_reply":"2024-11-23T19:59:32.827659Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Discriminator(nn.Module):\n    \"\"\" A Discriminator Network\n\n    ---\n    Structure\n        Img -> Discriminator -> Grid of predictions\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self):\n        \"\"\" Init CycleGAN's Discriminator\n\n        \"\"\"\n\n        super().__init__()\n\n        # --- Default settings\n        base_channels = 64\n        kernel_size = 4\n\n        # --- History\n        self.history = DiscriminatorHistory()\n\n        # --- Discriminator network\n        self.discriminator = nn.Sequential(\n            Downsample(in_channels=3, out_channels=base_channels,\n                       kernel_size=kernel_size, norm=False),\n            Downsample(in_channels=base_channels, out_channels=base_channels * 2,\n                       kernel_size=kernel_size),\n            Downsample(in_channels=base_channels * 2, out_channels=base_channels * 4,\n                       kernel_size=kernel_size),\n            nn.ZeroPad2d(padding=(0, 2, 0, 2)),\n            nn.Conv2d(in_channels=base_channels * 4, out_channels=base_channels * 8,\n                      kernel_size=4, stride=1, bias=False),\n            nn.InstanceNorm2d(num_features=base_channels * 8),\n            nn.ZeroPad2d(padding=(0, 2, 0, 2)),\n            nn.Conv2d(in_channels=base_channels * 8, out_channels=1,\n                      kernel_size=4, stride=1),\n            #             nn.Sigmoid()\n        )\n\n        # Initialise conv wights with N(0, 0.02)\n        self.apply(conv_weight_init)\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x,\n                sample_history=False):\n        \"\"\" Discriminate between real/fake img\n\n        ---\n        Parameters\n            x: Input tensor (img)\n            sample_history: If true it will mix x with samples from buffer\n\n        ---\n        Returns\n            torch.Tensor: Output tensor (pred grid)\n\n        \"\"\"\n\n        # Update x based on history\n        if sample_history:\n            x = self.history(x)\n\n\n        return self.discriminator(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T19:59:55.077503Z","iopub.execute_input":"2024-11-23T19:59:55.077907Z","iopub.status.idle":"2024-11-23T19:59:55.086932Z","shell.execute_reply.started":"2024-11-23T19:59:55.077876Z","shell.execute_reply":"2024-11-23T19:59:55.08581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define CycleGAN\n# --------------------------------------------------------------------------------\nclass CycleGAN(nn.Module):\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 lr=2e-4,\n                 lambda_cycle=10):  # Not sure what that is yet\n        \"\"\" Cycle GAN model\n\n        ---\n        Structure:\n            - 2 Generators\n                - Monet Generator - To go from Photo to Monet (G)\n                - Photo Generator - To got from Monet to Photo (F)\n\n            - 2 Discriminators\n                - Monet Discriminator - To say whether input is a real/fake monet (Dy)\n                - Photo Discriminator - To say whether input is a real/fake photo (Dx)\n\n\n        \"\"\"\n        super().__init__()\n\n        # Define Generators and Discriminators\n        self.gen_monet = Generator()\n        self.optim_gen_monet = optim.Adam(self.gen_monet.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        self.gen_photo = Generator()\n        self.optim_gen_photo = optim.Adam(self.gen_photo.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n\n        self.disc_monet = Discriminator()\n        self.optim_disc_monet = optim.Adam(self.disc_monet.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        self.disc_photo = Discriminator()\n        self.optim_disc_photo = optim.Adam(self.disc_photo.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n\n        # Loss functions\n        self.loss_l2 = nn.MSELoss()\n        self.loss_l1 = nn.L1Loss()\n\n        # Attribs\n        self.lambda_cycle = lambda_cycle\n\n    # --------------------------------------------------------------------------------\n    def forward(self, x):\n        \"\"\" Main point of the model is to go from photo -> monet so forward does just that\n\n        Expects a single tensor (photo data)\n            x: [btach, 3, 256, 256] - Real photo\n        \"\"\"\n\n        return self.gen_monet(x)\n\n    # --------------------------------------------------------------------------------\n    def forward_apply_style(self,\n                            real_A,\n                            style=\"monet\"):\n        \"\"\" Apply one of the styles either monet or photo\n\n        In addition the function returns more than just styled photo e.g.\n        Disciminators on Fake, Real and Cycled Photo\n\n        ---\n        Parameters\n            real_A: A batch of real samples from domain A being transfered to domain B\n            style: Domain B, Applies either monet or photo style\n\n        ---\n        Returns\n            dict: Dictionary of outputs from model\n\n        \"\"\"\n\n        with torch.inference_mode(mode=True):\n            self.eval()\n\n            # Set networks and Domain\n            if style == \"monet\":\n                gen_A = self.gen_photo\n                gen_B = self.gen_monet\n\n                disc_A = self.disc_photo\n                disc_B = self.disc_monet\n            elif style == \"photo\":\n                gen_A = self.gen_photo\n                gen_B = self.gen_monet\n\n                disc_A = self.disc_photo\n                disc_B = self.disc_monet\n            else:\n                raise ValueError(f\"Style can be either monet or photo, given: {style}\")\n\n            # ---\n            # Forward, notation uses P=Photo and M=Monet to easier demonstrate the flow\n            # x_P -> G_M -> \\hat{x}_M\n            fake_B = gen_B(real_A)\n\n            # \\hat{x}_M -> G_P -> \\hat{x}_P\n            cycled_A = gen_A(fake_B)\n\n            # \\hat{x}_M -> D_M -> \\hat{y}_M\n            disc_fake_B = disc_B(fake_B)\n\n            # x_P -> D_P -> y_P\n            disc_real_A = disc_A(real_A)\n\n            # \\hat{x}_P -> D_M -> \\hat{y}_P\n            disc_cycled_A = disc_A(cycled_A)\n\n        return {\"real_A\" : real_A,\n                \"gen_fake_B\" : fake_B,\n                \"gen_cycled_A\" : cycled_A,\n                \"disc_fake_B\" : disc_fake_B,\n                \"disc_real_A\" : disc_real_A,\n                \"disc_cycled_A\" : disc_cycled_A}\n\n    # --------------------------------------------------------------------------------\n    def forward_step(self, x):\n        \"\"\" A single forward step for the model\n\n        Forward -> Loss -> Backward -> Optimise\n\n        Expects a tuple of 2 tensors\n            (tensor, tensor): [batch, 3, 256, 256] - Real monet and photo\n        \"\"\"\n\n        # Unpack data\n        x_monet, x_photo = x\n\n        # Zero Grad Optimisers\n        self.optim_zero_grad()\n\n        # --- Generate fake data\n        x_fake_monet = self.gen_monet(x_photo)  # Real Photo -> Gm -> Fake Monet\n        x_fake_photo = self.gen_photo(x_monet)  # Real Monet -> Gp -> Fake Photo\n\n        # --- Discriminators\n        # Monet\n        d_real_monet = self.disc_monet(x_monet)  # How real does this real monet seems\n        d_fake_monet = self.disc_monet(x_fake_monet.detach(), sample_history=True)  # How fake does this fake seems\n\n        # Photo\n        d_real_photo = self.disc_photo(x_photo)\n        d_fake_photo = self.disc_photo(x_fake_photo.detach(), sample_history=True)\n\n        # Discriminators Loss\n        d_loss_monet = self.disc_loss(x_real=d_real_monet,\n                                      x_fake=d_fake_monet)\n\n        d_loss_photo = self.disc_loss(x_real=d_real_photo,\n                                      x_fake=d_fake_photo)\n\n        d_loss = d_loss_monet + d_loss_photo\n\n        # Backward\n        d_loss.backward()\n\n        # Optimise\n        self.optim_disc_monet.step()\n        self.optim_disc_photo.step()\n\n        # --- Generators\n        # Adversarial\n        x_adv_monet = self.disc_monet(x_fake_monet)  # How real does this fake monet seems\n        x_adv_photo = self.disc_photo(x_fake_photo)\n\n        # Cycle\n        x_cycled_monet = self.gen_monet(x_fake_photo)  # Fake Photo -> Gm -> Cycled Monet\n        x_cycled_photo = self.gen_photo(x_fake_monet)  # Fake Monet -> Gp -> Cycled Photo\n\n        # Adversarial Loss\n        g_loss_adv_monet = self.loss_l2(torch.ones_like(x_adv_monet), x_adv_monet)\n        g_loss_adv_photo = self.loss_l2(torch.ones_like(x_adv_photo), x_adv_photo)\n\n        # Cycle Loss\n        g_loss_cycle_monet = self.loss_l1(x_monet, x_cycled_monet) * self.lambda_cycle\n        g_loss_cycle_photo = self.loss_l1(x_photo, x_cycled_photo) * self.lambda_cycle\n\n        # Identity Loss\n#         g_loss_idt_monet = self.loss_l1(x_monet, x_fake_monet) * self.lambda_cycle\n#         g_loss_idt_photo = self.loss_l1(x_photo, x_fake_photo) * self.lambda_cycle\n\n        # Total Loss\n        g_loss_monet = g_loss_adv_monet + g_loss_cycle_monet #+ g_loss_idt_monet\n        g_loss_photo = g_loss_adv_photo + g_loss_cycle_photo #+ g_loss_idt_photo\n        g_loss = g_loss_monet + g_loss_photo\n\n        # Backward\n        g_loss.backward()\n\n        # Optimise\n        self.optim_gen_monet.step()\n        self.optim_gen_photo.step()\n\n        # Get Losses\n        return {\n            \"g_monet\" : g_loss_monet,\n            \"g_photo\" : g_loss_photo,\n            \"d_monet\" : d_loss_monet,\n            \"d_photo\" : d_loss_photo\n        }\n\n    # --------------------------------------------------------------------------------\n    def disc_loss(self,\n                  x_real,\n                  x_fake):\n        \"\"\" Compute loss for discriminator\n\n        Loss = 1/2(Real loss + Fake loss)\n\n        Expects 2 tensor\n            x_real: [batch, 3, 256, 256] - A real photo\n            x_fake: [batch, 3, 256, 256] - A fake photo from generator\n\n        \"\"\"\n\n        loss_real = self.loss_l2(torch.ones_like(x_real), x_real)\n        loss_fake = self.loss_l2(torch.zeros_like(x_fake), x_fake)\n\n        return (loss_real + loss_fake) * 0.5\n\n    # --------------------------------------------------------------------------------\n    def optim_zero_grad(self):\n        \"\"\" Apply zero grad on all optimisers\n        \"\"\"\n\n        # Generators\n        self.optim_gen_monet.zero_grad()\n        self.optim_gen_photo.zero_grad()\n\n        # Discriminators\n        self.optim_disc_monet.zero_grad()\n        self.optim_disc_photo.zero_grad()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T20:00:00.478093Z","iopub.execute_input":"2024-11-23T20:00:00.478538Z","iopub.status.idle":"2024-11-23T20:00:00.505075Z","shell.execute_reply.started":"2024-11-23T20:00:00.4785Z","shell.execute_reply":"2024-11-23T20:00:00.503095Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MonetDataset(Dataset):\n    def __init__(self,\n                 root_path):\n\n        # Paths\n        self.root_path = root_path\n        self.photo_paths = os.listdir(f\"{root_path}/photo_jpg\")\n        self.monet_paths = os.listdir(f\"{root_path}/monet_jpg\")\n\n    def __len__(self):\n        return len(self.photo_paths)\n\n    def __getitem__(self, idx):\n        \"\"\" Get one of the photos and a random monet\n        \"\"\"\n\n        # Get a photo path\n        path_photo = f\"{root_path}/photo_jpg/{self.photo_paths[idx]}\"\n\n        # Get a random monet path\n        path_monet = f\"{root_path}/monet_jpg/{random.choices(self.monet_paths, k=1)[0]}\"\n\n        # Load imgs and transform\n        x_photo = self.to_tensor(self.read_img(path_photo))\n        x_monet = self.to_tensor(self.read_img(path_monet))\n\n        return x_monet, x_photo, path_photo\n\n    def to_tensor(self, x):\n        \"\"\" Transform uint8 image [0,255] to torch float32 [-1,1]\n        \"\"\"\n        return torch.tensor(x, dtype=torch.float32).permute(2,0,1) / 127.5 - 1.0\n\n    def read_img(self, path):\n        \"\"\" Read img with cv2 and transform to RGB\n        \"\"\"\n        return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T20:00:16.260203Z","iopub.execute_input":"2024-11-23T20:00:16.260567Z","iopub.status.idle":"2024-11-23T20:00:16.269067Z","shell.execute_reply.started":"2024-11-23T20:00:16.260536Z","shell.execute_reply":"2024-11-23T20:00:16.267936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training params\nNUM_EPOCHS = 10\nBATCH_SIZE = 32\n\n\n# Device\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n\n# Keep losses per epoch\nloss_epoch = defaultdict(list)\n\n# Dataset and data loader\ndataset = MonetDataset(root_path=root_path)\ndata_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True, num_workers=4)\n\n# Model\ncycle_gan = CycleGAN(lr=2e-4,).to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T20:00:20.794151Z","iopub.execute_input":"2024-11-23T20:00:20.794555Z","iopub.status.idle":"2024-11-23T20:00:22.231265Z","shell.execute_reply.started":"2024-11-23T20:00:20.794514Z","shell.execute_reply":"2024-11-23T20:00:22.230004Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train for N epochs\nfor i in range(1, NUM_EPOCHS + 1):\n\n    # Tqdm\n    with tqdm(total=len(data_loader),\n          desc=f\"Epoch {i}/{NUM_EPOCHS}\",\n          ascii=True,\n          colour=\"green\") as pbar:\n\n        # Per batch\n        loss_batch = defaultdict(list)\n        for x in data_loader:\n            x_monet = x[0].to(device)\n            x_photo = x[1].to(device)\n\n            # Batch of data\n            loss = cycle_gan.forward_step((x_monet, x_photo))\n\n            # Update dict\n            for k, v in loss.items():\n                loss_batch[k].append(v.item())\n\n            # Bar Update\n            pbar.set_postfix({k:v.item() for k, v in loss.items()})\n            pbar.update()\n\n        # Update mean loss for epoch\n        for k, v in loss_batch.items():\n            loss_epoch[k].append(np.mean(v))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-23T20:00:58.224055Z","iopub.execute_input":"2024-11-23T20:00:58.224438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import PIL\n! mkdir images","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"i = 1\ndata_loader = DataLoader(dataset=dataset, batch_size=1)\nwith torch.inference_mode(mode=True):\n    cycle_gan.eval()\n\n    for x in data_loader:\n        photo, _, path = x\n\n        pred = cycle_gan(photo.to(device)).cpu().squeeze().numpy()\n        pred = np.transpose(pred, (1,2,0))\n\n        pred = (pred * 127.5 + 127.5).astype(np.uint8)\n        path = path[0].split(\"/\")[-1]\n        pred_img = PIL.Image.fromarray(pred)\n        #pred_img.save(f'images/{path}')\n        pred_img.save(f'/images/{path}')\n\n        i += 1","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport zipfile\n\n# Directory containing the images\noutput_dir = \"images\"  # Directory where images are stored\nzip_file_path = \"images.zip\"  # Path to save the ZIP file\n\n# Create a ZIP file\nwith zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:\n    for root, dirs, files in os.walk(output_dir):\n        for file in files:\n            file_path = os.path.join(root, file)  # Full path to the file\n            arcname = os.path.relpath(file_path, output_dir)  # Relative path inside the ZIP\n            zipf.write(file_path, arcname)  # Add file to ZIP archive\n\nprint(f\"Zipped images saved to {zip_file_path}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}