psnr_metric.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from itertools import chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_name = \"resnet\"\n",
    "# model_dir_name = \"resnet\"\n",
    "# model_dir_name = \"restnet_no_cycle_consistency\"\n",
    "\n",
    "model_name = \"fdcgan\"\n",
    "model_dir_name = \"fdcgan\"\n",
    "# model_dir_name = \"fdcgan_no_cycle_consistency\"\n",
    "\n",
    "dataset_name = \"apple2orange64\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Extract data from local files\n",
    "\n",
    "class ImageClass(Dataset):\n",
    "    def __init__(self, root, transformations, mode=\"train\", image_class=\"A\"):\n",
    "        self.transform = transforms.Compose(transformations)\n",
    "\n",
    "        dir_path = os.path.join(root, mode, image_class)\n",
    "        if not os.path.exists(dir_path):\n",
    "            raise FileNotFoundError(f\"Directory not found : {dir_path}\")\n",
    "        self.files = sorted(glob.glob(os.path.join(dir_path, \"*.*\")))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.transform(Image.open(self.files[index]))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.files)\n",
    "\n",
    "data_path = os.path.join(\".\", \"datasets\", dataset_name)\n",
    "img_size = 64\n",
    "img_shape = (3, img_size, img_size)\n",
    "\n",
    "transformations = [\n",
    "    transforms.Resize(int(img_size * 1.12), Image.BICUBIC),\n",
    "    transforms.RandomCrop((img_size, img_size)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "]\n",
    "\n",
    "dataset_a = ImageClass(data_path, transformations, mode=\"validation\", image_class=\"A\")\n",
    "dataset_b = ImageClass(data_path, transformations, mode=\"validation\", image_class=\"B\")\n",
    "\n",
    "def tensor_to_image(tensor_image):\n",
    "    image = tensor_image.detach().to('cpu').numpy()\n",
    "    image = np.moveaxis(image, [0, 1, 2], [2, 0, 1])\n",
    "    image = (image + 1) / 2\n",
    "    image[image < 0] = 0\n",
    "    image[image > 1] = 1\n",
    "    return image\n",
    "\n",
    "def show_sample(dataset):\n",
    "    fig, axs = plt.subplots(2, 5, figsize=(15, 6))\n",
    "    for i, j in enumerate(random.sample(range(len(dataset)), 10)):\n",
    "        img = dataset[j]\n",
    "        ax = axs[i // 5, i % 5]\n",
    "        ax.imshow(tensor_to_image(img))\n",
    "        ax.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "print(\"Samples of dataset A\")\n",
    "show_sample(dataset_a)\n",
    "\n",
    "print(\"Samples of dataset B\")\n",
    "show_sample(dataset_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Load the pretrained CycleGan generator\n",
    "\n",
    "model_save_dir = os.path.join(\".\", \"saved_models\", \"cycle-gan\", model_dir_name, dataset_name)\n",
    "\n",
    "if model_name == \"fdcgan\":\n",
    "    from fdcgan import Generator\n",
    "    g_ab = Generator(img_shape)\n",
    "    g_ba = Generator(img_shape)\n",
    "\n",
    "elif model_name == \"resnet\":\n",
    "    from resnet import Generator\n",
    "    g_ab = Generator(img_shape, num_residual_blocks=9)\n",
    "    g_ba = Generator(img_shape, num_residual_blocks=9)\n",
    "\n",
    "else:\n",
    "    raise ValueError(\"Invalid model name\")\n",
    "\n",
    "g_ab_save_path = os.path.join(model_save_dir, \"G_AB_199.pth\")\n",
    "g_ba_save_path = os.path.join(model_save_dir, \"G_BA_199.pth\")\n",
    "\n",
    "if not os.path.exists(g_ab_save_path):\n",
    "    raise FileNotFoundError(f\"Could not find the pretrained model: {g_ab_save_path}\")\n",
    "elif not os.path.exists(g_ba_save_path):\n",
    "    raise FileNotFoundError(f\"Could not find the pretrained model: {g_ba_save_path}\")\n",
    "else:\n",
    "    print(f\"Loading pretrained models:\\nA->B: {g_ab_save_path}\\nB->A: {g_ba_save_path}\")\n",
    "\n",
    "g_ab.load_state_dict(torch.load(g_ab_save_path))\n",
    "g_ba.load_state_dict(torch.load(g_ba_save_path))\n",
    "\n",
    "g_ab.eval()\n",
    "g_ba.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Translate all the images of the test dataset into images of the other class\n",
    "\n",
    "with torch.no_grad():\n",
    "    original_a = next(iter(DataLoader(dataset_a, batch_size=len(dataset_a), shuffle=False)))\n",
    "    generated_b = g_ab(original_a)\n",
    "    recovered_a = g_ba(generated_b)\n",
    "\n",
    "    original_b = next(iter(DataLoader(dataset_b, batch_size=len(dataset_b), shuffle=False)))\n",
    "    generated_a = g_ba(original_b)\n",
    "    recovered_b = g_ab(generated_a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Show samples of the images and their recovering\n",
    "\n",
    "for i in random.sample(range(len(original_a)), 3):\n",
    "    origin = original_a[i]\n",
    "    gener = generated_b[i]\n",
    "    recov = recovered_a[i]\n",
    "\n",
    "    fig, axis = plt.subplots(1, 3, figsize=(8, 2))\n",
    "    axis[0].set_title(\"original a\")\n",
    "    axis[0].imshow(tensor_to_image(origin))\n",
    "    axis[1].set_title(\"generated b\")\n",
    "    axis[1].imshow(tensor_to_image(gener))\n",
    "    axis[2].set_title(\"recovered a\")\n",
    "    axis[2].imshow(tensor_to_image(recov))\n",
    "\n",
    "for i in random.sample(range(len(original_b)), 3):\n",
    "    origin = original_b[i]\n",
    "    gener = generated_a[i]\n",
    "    recov = recovered_b[i]\n",
    "\n",
    "    fig, axis = plt.subplots(1, 3, figsize=(8, 2))\n",
    "    axis[0].set_title(\"original b\")\n",
    "    axis[0].imshow(tensor_to_image(origin))\n",
    "    axis[1].set_title(\"generated a\")\n",
    "    axis[1].imshow(tensor_to_image(gener))\n",
    "    axis[2].set_title(\"recovered b\")\n",
    "    axis[2].imshow(tensor_to_image(recov))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Compute the PSNR value\n",
    "\n",
    "def compute_psnr(original_images, recoverd_images):\n",
    "    assert len(original_images) == len(recoverd_images)\n",
    "    total_psnr = 0.\n",
    "    for origin, recov in zip(original_images, recoverd_images):\n",
    "        total_psnr += 10. * torch.log10(\n",
    "            (img_size * img_size) / ((origin - recov).norm(2) ** 2)\n",
    "        )\n",
    "    return total_psnr / len(original_images)\n",
    "\n",
    "with torch.no_grad():\n",
    "    psnr_a = compute_psnr(original_a, recovered_a)\n",
    "    psnr_b = compute_psnr(original_b, recovered_b)\n",
    "\n",
    "print(psnr_a)\n",
    "print(psnr_b)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "machine-learning",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}