{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from tqdm.auto import tqdm\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import make_grid\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "torch.manual_seed(0)\n",
    "\n",
    "def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):\n",
    "    '''\n",
    "    Function for visualizing images: Given a tensor of images, number of images, and\n",
    "    size per image, plots and prints the images in an uniform grid.\n",
    "    '''\n",
    "    image_shifted = image_tensor\n",
    "    image_unflat = image_shifted.detach().cpu().view(-1, *size)\n",
    "    image_grid = make_grid(image_unflat[:num_images], nrow=5)\n",
    "    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop(image, new_shape):\n",
    "    middle_height = image.shape[2] // 2\n",
    "    middle_width = image.shape[3] // 2\n",
    "    starting_height = middle_height - new_shape[2] // 2\n",
    "    final_height = starting_height + new_shape[2]\n",
    "    starting_width = middle_width - new_shape[3] // 2\n",
    "    final_width = starting_width + new_shape[3]\n",
    "    cropped_image = image[:, :, starting_height:final_height, starting_width:final_width]    \n",
    "    return cropped_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def  __init__(self,input_channels, use_bn=True, use_dropout = False):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(input_channels, 2 * input_channels, kernel_size=3)\n",
    "        self.conv2 = nn.Conv2d(input_channels * 2, 2 * input_channels, kernel_size = 3)\n",
    "        self.activation = nn.ReLU()\n",
    "        self.pooling = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        if use_bn:\n",
    "            self.bn = nn.BatchNorm2d(input_channels * 2)\n",
    "        self.use_bn = use_bn\n",
    "        if use_dropout:\n",
    "            self.dropout = nn.Dropout()\n",
    "        self.use_dp = use_dropout\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        if self.use_bn:\n",
    "            x = self.bn(x)\n",
    "        if self.use_dp:\n",
    "            x = self.dropout(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.conv2(x)\n",
    "        if self.use_bn:\n",
    "            x = self.bn(x)\n",
    "        if self.use_dp:\n",
    "            x = self.dropout(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.pooling(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, input_channels, use_bn=True, use_dropout = False):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.upsample = nn.Upsample(scale_factor = 2, mode='bilinear', align_corners=True)\n",
    "        self.conv1 = nn.Conv2d(input_channels, input_channels // 2, kernel_size = 2)\n",
    "        self.conv2 = nn.Conv2d(input_channels , input_channels //2 , kernel_size = 3,padding=1)\n",
    "        self.conv3 = nn.Conv2d(input_channels //2, input_channels // 2, kernel_size = 2,padding=1)\n",
    "        self.activation = nn.ReLU()\n",
    "        if use_bn:\n",
    "            self.bn = nn.BatchNorm2d(input_channels // 2)\n",
    "        self.use_bn = use_bn\n",
    "        if use_dropout:\n",
    "            self.dropout = nn.Dropout()\n",
    "        self.use_dp = use_dropout\n",
    "        \n",
    "\n",
    "    def forward(self, x, skip_con_x):\n",
    "        x = self.upsample(x)\n",
    "        x = self.conv1(x)\n",
    "        skip_con_x = crop(skip_con_x, x.shape)\n",
    "        x = torch.cat([x, skip_con_x], axis = 1)\n",
    "        x = self.conv2(x)\n",
    "        if self.use_bn:\n",
    "            x = self.bn(x)\n",
    "        if self.use_dp:\n",
    "            x = self.dropout(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.conv3(x)\n",
    "        if self.use_bn:\n",
    "            x = self.bn(x)\n",
    "        if self.use_dp:\n",
    "            x = self.dropout(x)\n",
    "        x = self.activation(x)\n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureMapBlock(nn.Module):\n",
    "    def __init__(self, input_channels, output_channels):\n",
    "        super(FeatureMapBlock, self).__init__()\n",
    "        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=1)\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "    def __init__(self, input_channels, output_channels, hidden_channels = 64):\n",
    "        super(UNet, self).__init__()\n",
    "        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\n",
    "        self.encoder1 = Encoder(hidden_channels, use_dropout=True)\n",
    "        self.encoder2 = Encoder(hidden_channels * 2,use_dropout=True)\n",
    "        self.encoder3 = Encoder(hidden_channels * 4, use_dropout=True)\n",
    "        self.encoder4 = Encoder(hidden_channels * 8)\n",
    "        self.encoder5 = Encoder(hidden_channels * 16)\n",
    "        self.encoder6 = Encoder(hidden_channels * 32)\n",
    "        self.decoder1 = Decoder(hidden_channels * 64)\n",
    "        self.decoder2 = Decoder(hidden_channels * 32)\n",
    "        self.decoder3 = Decoder(hidden_channels * 16)\n",
    "        self.decoder4 = Decoder(hidden_channels * 8)\n",
    "        self.decoder5 = Decoder(hidden_channels * 4 )\n",
    "        self.decoder6 = Decoder(hidden_channels * 2)\n",
    "        self.downfeature = FeatureMapBlock(hidden_channels , output_channels)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x0 = self.upfeature(x)\n",
    "        x1 = self.encoder1(x0)\n",
    "        x2 = self.encoder2(x1)\n",
    "        x3 = self.encoder3(x2)\n",
    "        x4 = self.encoder4(x3)\n",
    "        x5 = self.encoder5(x4)\n",
    "        x6 = self.encoder6(x5)\n",
    "        x7 = self.decoder1(x6, x5)\n",
    "        x8 = self.decoder2(x7, x4)\n",
    "        x9 = self.decoder3(x8, x3)\n",
    "        x10 = self.decoder4(x9, x2)\n",
    "        x11 = self.decoder5(x10, x1)\n",
    "        x12 = self.decoder6(x11, x0)\n",
    "        xn = self.downfeature(x12)\n",
    "\n",
    "        return self.sigmoid(xn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_channels, hidden_channels=8) -> None:\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.upfeature = FeatureMapBlock(input_channels,hidden_channels)\n",
    "        self.encoder1 = Encoder(hidden_channels, use_bn=False)\n",
    "        self.encoder2 = Encoder(hidden_channels*2)\n",
    "        self.encoder3 = Encoder(hidden_channels*4)\n",
    "        self.encoder4 = Encoder(hidden_channels*8)\n",
    "        self.final = nn.Conv2d(hidden_channels * 16,1,kernel_size=1)\n",
    "\n",
    "    def forward(self, x,y):\n",
    "        x = torch.cat([x,y],axis=1)\n",
    "        x0 = self.upfeature(x)\n",
    "        x1 = self.encoder1(x0)\n",
    "        x2 = self.encoder2(x1)\n",
    "        x3 = self.encoder3(x2)\n",
    "        x4 = self.encoder4(x3)\n",
    "        xn = self.final(x4)\n",
    "\n",
    "        return xn\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "# New parameters\n",
    "adv_criterion = nn.BCEWithLogitsLoss() \n",
    "recon_criterion = nn.L1Loss() \n",
    "lambda_recon = 200\n",
    "\n",
    "n_epochs = 20\n",
    "input_dim = 3\n",
    "real_dim = 3\n",
    "display_step = 200\n",
    "batch_size = 4\n",
    "lr = 0.0002\n",
    "target_shape = 256\n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "dataset = torchvision.datasets.ImageFolder(\"/datsets/maps\", transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = UNet(input_dim, real_dim).to(device)\n",
    "gen_opt = torch.optim.Adam(gen.parameters(),lr=lr)\n",
    "disc = Discriminator(input_dim+real_dim).to(device)\n",
    "disc_opt = torch.optim.Adam(disc.parameters(),lr=lr)\n",
    "\n",
    "def weights_init(m):\n",
    "    if isinstance(m,nn.Conv2d) or isinstance(m,nn.ConvTranspose2d):\n",
    "        torch.nn.init.normal_(m.weight,0.0,0.02)\n",
    "    if isinstance(m,nn.BatchNorm2d):\n",
    "        torch.nn.init.normal_(m.weight,0.0,0.02)\n",
    "        torch.nn.init.constant_(m.bias,0.0)\n",
    "\n",
    "gen =gen.apply(weights_init)\n",
    "disc = disc.apply(weights_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gen_loss(gen, disc, real, condition, dev_criterion, recon_criterion, lambda_recon):\n",
    "    gen_pred = gen(condition)\n",
    "    disc_pred = disc(gen_pred, real)\n",
    "    gen_adv_loss = dev_criterion(disc_pred, torch.ones_like(disc_pred))\n",
    "    gen_recon_loss = recon_criterion(real,gen_pred)\n",
    "    gen_loss = gen_adv_loss + lambda_recon * gen_recon_loss\n",
    "\n",
    "    return gen_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    mean_gen_loss = 0\n",
    "    mean_disc_loss = 0\n",
    "    dataloader = DataLoader(dataset,batch_size=batch_size,shuffle=True)\n",
    "\n",
    "    cur_step = 0\n",
    "    for epoch in range(n_epochs):\n",
    "        for image,_ in tqdm(dataloader):\n",
    "            image_width = image.shape[3]\n",
    "            condition = image[:, :, :, :image_width // 2]\n",
    "            condition = nn.functional.interpolate(condition, size=target_shape)\n",
    "            real = image[:, :, :, image_width // 2:]\n",
    "            real = nn.functional.interpolate(real, size=target_shape)\n",
    "            cur_batch_size = len(real)\n",
    "            condition = condition.to(device)\n",
    "            real = real.to(device)\n",
    "\n",
    "            disc_opt.zero_grad()\n",
    "            with torch.no_grad():\n",
    "                fake = gen(condition)\n",
    "            disc_fake_pred = disc(fake.detach(),condition)\n",
    "            disc_fake_loss = adv_criterion(disc_fake_pred,torch.zeros_like(disc_fake_pred))\n",
    "            disc_real_pred = disc(real,condition)\n",
    "            disc_real_loss = adv_criterion(disc_real_pred,torch.ones_like(disc_real_pred))\n",
    "            disc_loss = (disc_fake_loss + disc_real_loss) / 2\n",
    "            disc_loss.backward(retain_graph= True)\n",
    "            disc_opt.step()\n",
    "\n",
    "            gen.zero_grad()\n",
    "            gen_loss = get_gen_loss(gen,disc,real,condition,adv_criterion,recon_criterion,lambda_recon)\n",
    "            gen_loss.backward()\n",
    "            gen_opt.step()\n",
    "\n",
    "            mean_gen_loss += gen_loss.item() / display_step\n",
    "            mean_disc_loss += disc_loss.item() / display_step\n",
    "\n",
    "            if cur_step % display_step == 0:\n",
    "                print(f\"Epoch: {epoch} | Step: {cur_step} | Gen-Loss: {mean_gen_loss} | Disc-loss: {mean_disc_loss}\")\n",
    "                show_tensor_images(condition, size=(input_dim, target_shape, target_shape))\n",
    "                show_tensor_images(real, size=(real_dim, target_shape, target_shape))\n",
    "                show_tensor_images(fake, size=(real_dim, target_shape, target_shape))\n",
    "                mean_gen_loss = 0\n",
    "                mean_disc_loss = 0\n",
    "\n",
    "                torch.save({\n",
    "                    \"gen\":gen,\n",
    "                    \"disc\":disc,\n",
    "                    \"gen_opt\": gen_opt,\n",
    "                    \"disc_opt\": disc_opt\n",
    "                },f\"checkpoints/Pix2Pix_Epoch{epoch}.pth\")\n",
    "            cur_step += 1\n",
    "\n",
    "train()\n",
    "\n",
    "            "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}