{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30301,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision\nimport torchvision.transforms as transforms\nimport os\nimport torch.nn.functional as F\nimport PIL.Image as Image\nimport time\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport zipfile\n\ntrain_dir = '/kaggle/input/gan-getting-started/monet_jpg'\ntest_dir = '/kaggle/input/gan-getting-started/photo_jpg'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-12-31T07:11:26.684589Z","iopub.execute_input":"2022-12-31T07:11:26.684987Z","iopub.status.idle":"2022-12-31T07:11:26.692835Z","shell.execute_reply.started":"2022-12-31T07:11:26.684956Z","shell.execute_reply":"2022-12-31T07:11:26.69155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs = 5","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.694706Z","iopub.execute_input":"2022-12-31T07:11:26.695218Z","iopub.status.idle":"2022-12-31T07:11:26.703713Z","shell.execute_reply.started":"2022-12-31T07:11:26.695172Z","shell.execute_reply":"2022-12-31T07:11:26.702448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(torch.utils.data.Dataset):\n    def __init__(self, root, transform=None):\n        self.root = root\n        self.transform = transform\n        # Get the list of all image files in the folder\n        self.images = os.listdir(root)\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        # Load the image file\n        image = Image.open(os.path.join(self.root, self.images[idx]))\n        # Apply the transform (if specified)\n        if self.transform:\n            image = self.transform(image)\n        # Return the image and its label\n        return image, 0\n\ndef load_data():\n    # Define the transform to normalize the data\n    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\n    # Load the A and B datasets\n    A_data = CustomDataset(root=test_dir, transform=transform)\n    B_data = CustomDataset(root=train_dir, transform=transform)\n    \n    # Return the datasets as a dictionary\n    return {'A': A_data, 'B': B_data}","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.705021Z","iopub.execute_input":"2022-12-31T07:11:26.705775Z","iopub.status.idle":"2022-12-31T07:11:26.716325Z","shell.execute_reply.started":"2022-12-31T07:11:26.705741Z","shell.execute_reply":"2022-12-31T07:11:26.715138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define the generator and discriminator models\nclass Generator(nn.Module):\n    def __init__(self):\n        super(Generator, self).__init__()\n        # Define the layers of the generator model\n        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)\n        self.relu1 = nn.ReLU()\n        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n        self.relu2 = nn.ReLU()\n        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)\n        self.relu3 = nn.ReLU()\n        self.conv4 = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1)\n        self.relu4 = nn.ReLU()\n        self.deconv1 = nn.ConvTranspose2d(64, 128, kernel_size=3, stride=1, padding=1)\n        self.relu5 = nn.ReLU()\n        self.deconv2 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=1, padding=1)\n        self.relu6 = nn.ReLU()\n        self.deconv3 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=1, padding=1)\n        self.relu7 = nn.ReLU()\n        self.deconv4 = nn.ConvTranspose2d(32, 3, kernel_size=3, stride=1, padding=1)\n        \n    def forward(self, x):\n        # Define the forward pass through the network\n        x = self.conv1(x)\n        x = self.relu1(x)\n        x = self.conv2(x)\n        x = self.relu2(x)\n        x = self.conv3(x)\n        x = self.relu3(x)\n        x = self.conv4(x)\n        x = self.relu4(x)\n        x = self.deconv1(x)\n        x = self.relu5(x)\n        x = self.deconv2(x)\n        x = self.relu6(x)\n        x = self.deconv3(x)\n        x = self.relu7(x)\n        x = self.deconv4(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.719167Z","iopub.execute_input":"2022-12-31T07:11:26.720186Z","iopub.status.idle":"2022-12-31T07:11:26.734949Z","shell.execute_reply.started":"2022-12-31T07:11:26.720139Z","shell.execute_reply":"2022-12-31T07:11:26.733975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n        # Define the layers of the discriminator model\n        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)\n        self.relu1 = nn.ReLU()\n        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n        self.relu2 = nn.ReLU()\n        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)\n        self.relu3 = nn.ReLU()\n        self.conv4 = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1)\n        self.relu4 = nn.ReLU()\n        self.conv_out = nn.Conv2d(64, 3, kernel_size=1, stride=1)\n\n    def forward(self, x):\n        # Define the forward pass through the network\n        x = self.conv1(x)\n        x = self.relu1(x)\n        x = self.conv2(x)\n        x = self.relu2(x)\n        x = self.conv3(x)\n        x = self.relu3(x)\n        x = self.conv4(x)\n        x = self.relu4(x)\n        x = self.conv_out(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.736258Z","iopub.execute_input":"2022-12-31T07:11:26.736819Z","iopub.status.idle":"2022-12-31T07:11:26.750826Z","shell.execute_reply.started":"2022-12-31T07:11:26.736786Z","shell.execute_reply":"2022-12-31T07:11:26.74986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator = Generator()\ndiscriminator = Discriminator()\n\n# Define the loss function and optimizers\ncriterion = nn.BCEWithLogitsLoss()\nG_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\nD_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n\n# Load the data and split it into A and B datasets\ndata = load_data()\nA_data = data['A']\nB_data = data['B']","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.752224Z","iopub.execute_input":"2022-12-31T07:11:26.753256Z","iopub.status.idle":"2022-12-31T07:11:26.779694Z","shell.execute_reply.started":"2022-12-31T07:11:26.753223Z","shell.execute_reply":"2022-12-31T07:11:26.778656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_discriminator(discriminator, real_A, real_B):\n    # Zero the gradients\n    D_optimizer.zero_grad()\n    # Compute the output of the discriminator for the real A data\n    real_A_output = discriminator(real_A)\n    # Compute the loss for the real A data\n    real_A_loss = criterion(real_A_output, torch.ones(real_A_output.size()))\n    # Compute the output of the discriminator for the real B data\n    real_B_output = discriminator(real_B)\n    # Compute the loss for the real B data\n    real_B_loss = criterion(real_B_output, torch.ones(real_B_output.size()))\n    # Compute the total loss for the real data\n    loss = real_A_loss + real_B_loss\n    # Backpropagate the loss\n    loss.backward()\n    # Update the weights\n    D_optimizer.step()\n    return loss.item()","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.780886Z","iopub.execute_input":"2022-12-31T07:11:26.78173Z","iopub.status.idle":"2022-12-31T07:11:26.788579Z","shell.execute_reply.started":"2022-12-31T07:11:26.781694Z","shell.execute_reply":"2022-12-31T07:11:26.78756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_generator(generator, discriminator, A, B):\n    # Zero the gradients\n    G_optimizer.zero_grad()\n    # Generate fake A data\n    fake_A = generator(A)\n    # Generate fake B data\n    fake_B = generator(B)\n    # Compute the loss for the fake A data\n    # .permute(2, 1, 0)\n    fake_A_loss = criterion(discriminator(fake_A), torch.ones(fake_A.size()))\n    # Compute the loss for the fake B data\n    fake_B_loss = criterion(discriminator(fake_B), torch.ones(fake_B.size()))\n    # Compute the total loss\n    loss = fake_A_loss + fake_B_loss\n    # Backpropagate the loss and update the weights\n    loss.backward()\n    G_optimizer.step()\n    # Return the total loss\n    return loss.item()\n","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.789991Z","iopub.execute_input":"2022-12-31T07:11:26.790923Z","iopub.status.idle":"2022-12-31T07:11:26.804306Z","shell.execute_reply.started":"2022-12-31T07:11:26.790889Z","shell.execute_reply":"2022-12-31T07:11:26.803357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training loop\nfor epoch in range(num_epochs):\n    for i, (A, B) in enumerate(zip(A_data, B_data)):\n        start = time.time()\n        A = A[0]\n        B = B[0]\n        # Generate fake B images from A\n        fake_B = generator(A)\n        # Generate fake A images from B\n        fake_A = generator(B)\n        # Train the discriminator on real A and B images\n        D_real_loss = train_discriminator(discriminator, A, B)\n        # Train the discriminator on fake A and B images\n        D_fake_loss = train_discriminator(discriminator, fake_A, fake_B)\n        # Compute the total discriminator loss\n        D_loss = D_real_loss + D_fake_loss\n        # Train the generator on fake A and B images\n        G_loss = train_generator(generator, discriminator, A, B)\n        # Print the losses\n        print(\"Epoch: {}/{}, Step: {}/{}, D Loss: {:.10f}, G Loss: {:.10f}, T: {:.2f} Sec\".format(epoch+1, num_epochs, i+1, min(len(A_data), len(B_data)), D_loss, G_loss, time.time()-start))\n\n# Save the trained model\ntorch.save(generator.state_dict(), \"generator.pt\")\ntorch.save(discriminator.state_dict(), \"discriminator.pt\")","metadata":{"execution":{"iopub.status.busy":"2022-12-31T07:11:26.80735Z","iopub.execute_input":"2022-12-31T07:11:26.807723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def image_to_tensor(image):\n    # Convert the image to a tensor\n    image = np.transpose(image, (2, 0, 1)).astype(np.float32)\n    image = torch.from_numpy(image).float()\n    return image\n\ndef normalize_image(image):\n    mean = image.mean()\n    std = image.std()\n    # Normalize the image\n    return (image - mean) / std\n\ndef denormalize_image(image):\n    mean = image.mean()\n    std = image.std()\n    # Denormalize the image\n    image = image * std + mean\n    return image\n\ndef tensor_to_image(tensor):\n    # Convert the tensor to an image\n    tensor = tensor.detach().cpu()\n    image = tensor.numpy()\n    image = np.transpose(image, (1, 2, 0))\n    image = image.astype(np.uint8)\n    return image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"zip_obj = zipfile.ZipFile('submission.zip', 'w')\n\nfor i, image_file in enumerate(os.listdir(test_dir)):\n    input_image = np.array(Image.open(f'{test_dir}/{image_file}'))\n    # Convert the input image to a tensor\n    input_tensor = image_to_tensor(input_image)\n    # Normalize the input image\n    normalized = normalize_image(input_tensor)\n    # Generate the output image\n    output = generator(normalized)\n    # Denormalize the output image\n    denormalized = denormalize_image(output)\n    # Convert the output image tensor to a numpy array\n    output_tensor = tensor_to_image(denormalized)\n    # Save the output image\n    output_image = Image.fromarray(output_tensor)\n    filename = 'image_{}.jpg'.format(i)\n    file_path = os.path.join('./', filename)\n    output_image.save(filename)\n    zip_obj.write(file_path, arcname=filename)\n    os.remove(file_path)\n    print(f'image {i} generated')\n\nzip_obj.close()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create a figure with two subplots\nfig, axs = plt.subplots(1, 2)\n\n# Display the last images in the subplots\naxs[0].imshow(input_image)\naxs[0].set_title('input')\naxs[1].imshow(output_image)\naxs[1].set_title('output')\n\n# Show the plot\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}