super_res_NN_simplified.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vnabZn8zu0QI"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torchvision.transforms import ToTensor, Resize\n",
        "from PIL import Image\n",
        "import os\n",
        "import glob\n",
        "\n",
        "# Residual Block Definition\n",
        "class ResidualBlock(nn.Module):\n",
        "    def __init__(self, in_channels):\n",
        "        super(ResidualBlock, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)\n",
        "        self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        identity = x\n",
        "        out = self.relu(self.conv1(x))\n",
        "        out = self.conv2(out)\n",
        "        out += identity\n",
        "        return self.relu(out)\n",
        "\n",
        "# Super Resolution Network with Residual Block\n",
        "class SuperResolutionNet(nn.Module):\n",
        "    def __init__(self, upscale_factor):\n",
        "        super(SuperResolutionNet, self).__init__()\n",
        "\n",
        "        self.upscale_factor = upscale_factor\n",
        "\n",
        "        # Initial convolution layer\n",
        "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=5, stride=1, padding=2)\n",
        "\n",
        "        # Residual block\n",
        "        self.residual_block = ResidualBlock(64)\n",
        "\n",
        "        # Second convolution layer\n",
        "        self.conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        # PixelShuffle layer\n",
        "        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)\n",
        "\n",
        "        # Output convolution layer\n",
        "        self.conv3 = nn.Conv2d(in_channels=32 // (upscale_factor ** 2), out_channels=3, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Apply initial convolution layer\n",
        "        x = F.relu(self.conv1(x))\n",
        "\n",
        "        # Apply residual block\n",
        "        x = self.residual_block(x)\n",
        "\n",
        "        # Apply second convolution layer\n",
        "        x = F.relu(self.conv2(x))\n",
        "\n",
        "        # Apply PixelShuffle layer\n",
        "        x = self.pixel_shuffle(x)\n",
        "\n",
        "        # Apply final convolution layer\n",
        "        x = self.conv3(x)\n",
        "        return x\n",
        "\n",
        "# Dataset creation by lowering resolution\n",
        "def create_low_res_dataset(image_dir, low_res_dir, scale_factor):\n",
        "    os.makedirs(low_res_dir, exist_ok=True)\n",
        "    transform = ToTensor()\n",
        "\n",
        "    for img_path in glob.glob(os.path.join(image_dir, '*.jpg')):\n",
        "        img = Image.open(img_path).convert('RGB')\n",
        "        high_res_img = transform(img)\n",
        "        low_res_img = Resize((high_res_img.shape[1] // scale_factor, high_res_img.shape[2] // scale_factor))(img)\n",
        "        low_res_img.save(os.path.join(low_res_dir, os.path.basename(img_path)))\n",
        "\n",
        "# Loss function: Mean Squared Error\n",
        "def loss_function(output, target):\n",
        "    return F.mse_loss(output, target)\n",
        "\n",
        "# Example usage\n",
        "if __name__ == \"__main__\":\n",
        "    # Define the network with an upscale factor of 2\n",
        "    net = SuperResolutionNet(upscale_factor=2)\n",
        "\n",
        "    # Create a random input tensor with shape (batch_size, channels, height, width)\n",
        "    input_tensor = torch.randn(1, 3, 24, 24)  # Example input size: (1, 3, 24, 24)\n",
        "\n",
        "    # Forward pass\n",
        "    output_tensor = net(input_tensor)\n",
        "\n",
        "    print(f\"Input tensor shape: {input_tensor.shape}\")\n",
        "    print(f\"Output tensor shape: {output_tensor.shape}\")\n",
        "\n",
        "    # Example dataset creation\n",
        "    image_dir = 'path/to/high_res_images'\n",
        "    low_res_dir = 'path/to/low_res_images'\n",
        "    scale_factor = 2\n",
        "    create_low_res_dataset(image_dir, low_res_dir, scale_factor)\n",
        "\n",
        "    # Example loss calculation\n",
        "    high_res_tensor = torch.randn(1, 3, 48, 48)  # Example high-resolution target\n",
        "    loss = loss_function(output_tensor, high_res_tensor)\n",
        "    print(f\"Loss: {loss.item()}\")\n"
      ]
    }
  ]
}