{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "abda67c113f89d41",
   "metadata": {},
   "source": [
    "# `pyvisim` Introduction\n",
    "\n",
    "Welcome to this library! `pyvisim` hosts a collection of traditional and deep learning-based image similarity metrics and follow the Object-Oriented design. In this Tutorial, following metrics will be covered, through which you can get a general impression of this library:\n",
    "\n",
    "- `Structural Similarity Index (SSIM)`: compares two aligned images through local luminance, contrast and structure statistics. The score ranges from `0` to `1`, where only identical images score `1`.\n",
    "- `Peak Signal-to-Noise Ratio (PSNR)`: originiating from Signal Processing, it is commonly used to measure the \"lossiness\" of an image after compression. It ranges from `0` to `inf`, where only identical images score `inf`.\n",
    "- `Clip Embedder`: a neural network trained on a paradigm called `Contrastive Language-Image Pretraining (CLIP)` to embed images and texts into a shared latent space, where similarity can be measured by the cosine similarity of their embeddings. See [This Paper](https://arxiv.org/abs/2103.00020) for more details.\n",
    "\n",
    "## 1. Setting Up\n",
    "\n",
    "### Required Imports and Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "193123c05eb6c5c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "from itertools import islice\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "from pyvisim.datasets import OxfordFlowerDataset\n",
    "from pyvisim.dense.pixelwise import PSNR\n",
    "from pyvisim.dense.structural import SSIM\n",
    "from pyvisim.neural_networks import ClipEmbedder\n",
    "from pyvisim.typing import UInt8NumpyArray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eee539a",
   "metadata": {},
   "source": [
    "### Set the strength of the image degradations\n",
    "\n",
    "`JPEG_QUALITY` is the JPEG quality (from 1 to 95) of the compressed image, and `NOISE_SIGMA` is the standard deviation of the Gaussian noise added to the noised image. Lower the quality or raise the noise to degrade the image further.\n",
    "\n",
    "We will use these to see how `SSIM` and `PSNR` respond to image degradations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5a82a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "JPEG_QUALITY = 10\n",
    "NOISE_SIGMA = 25"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fff3325",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6382bc4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_image(image: UInt8NumpyArray | torch.Tensor, title: str = \"Image\") -> None:\n",
    "    \"\"\"\n",
    "    Plot a single image.\n",
    "\n",
    "    :param image: Image as a NumPy array (H, W, C) or torch tensor (C, H, W)\n",
    "    :param title: Title of the plot\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    if isinstance(image, torch.Tensor):\n",
    "        image = image.detach().cpu()\n",
    "        if image.ndim == 3:\n",
    "            image = image.permute(1, 2, 0)\n",
    "        image = image.numpy()\n",
    "    plt.imshow(image)\n",
    "    plt.axis(\"off\")\n",
    "    plt.title(title)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10e85464caa08385",
   "metadata": {},
   "source": [
    "## 2. Initializing the CLIP Embedder\n",
    "\n",
    "The `ClipEmbedder` downloads the pre-trained weights of the requested CLIP variant from the Hugging Face Hub on first use and caches them. Here, we use the `ViT-B-32` variant with the original `openai` weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6485084faa3b7593",
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_embedder = ClipEmbedder(variant=\"ViT-B-32\", pretrained=\"openai\")\n",
    "print(\"CLIP embedder:\", clip_embedder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d75f0dd",
   "metadata": {},
   "source": [
    "## 3. Embedding Images and Computing Similarities\n",
    "\n",
    "First, we load some sample images from the Oxford Flower dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a0cf21efec70271",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = OxfordFlowerDataset()\n",
    "images = [\n",
    "    image for i, (image, *_) in enumerate(islice(dataset, 200))\n",
    "]  # Set '200' to len(dataset) for full dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a2b9f9c",
   "metadata": {},
   "source": [
    "### Example: Embedding Images\n",
    "\n",
    "We will convert all images from the dataset batch above to embeddings (which is basically the numerical representation of the image using the algorithm of the embedder)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c545ba6e98369d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_embeddings_clip = clip_embedder.embed(images)\n",
    "print(\"Shape of CLIP embeddings:\", image_embeddings_clip.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee04d9b40251baf",
   "metadata": {},
   "source": [
    "\n",
    "### Example: Computing Similarity Between Two Images\n",
    "\n",
    "First, we choose and plot the images to compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c1bf60166fd98a",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_1, image_2 = images[0], images[1]\n",
    "plot_image(image_1, title=\"Image 1\")\n",
    "plot_image(image_2, title=\"Image 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885590c47bacfc3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute similarity using the CLIP embedder\n",
    "similarity = clip_embedder.similarity_score(image_1, image_2)\n",
    "\n",
    "print(\"Similarity score:\", similarity)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa895951",
   "metadata": {},
   "source": [
    "## 4. Comparing an Image with its Compressed and Noised Versions\n",
    "\n",
    "> [!NOTE]\n",
    "> `SSIM` and `PSNR` compare two aligned images pixel by pixel, so both images must have the same shape:\n",
    "\n",
    "### Example: Creating a Compressed and a Noised Image\n",
    "\n",
    "We create a JPEG-compressed and a noised version of `image_1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc2fef7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compress_jpeg(image, quality):\n",
    "    \"\"\"\n",
    "    Compress an image with JPEG and decode it again.\n",
    "\n",
    "    :param image: RGB image as a NumPy array (H, W, C)\n",
    "    :param quality: JPEG quality, from 1 to 95\n",
    "    :return: The decoded JPEG image as a NumPy array (H, W, C)\n",
    "    \"\"\"\n",
    "    buffer = io.BytesIO()\n",
    "    Image.fromarray(image).save(buffer, format=\"JPEG\", quality=quality)\n",
    "    with Image.open(buffer) as compressed:\n",
    "        return np.asarray(compressed.convert(\"RGB\"))\n",
    "\n",
    "\n",
    "def add_gaussian_noise(image, sigma, seed=42):\n",
    "    \"\"\"\n",
    "    Add Gaussian noise to an image.\n",
    "\n",
    "    :param image: RGB image as a NumPy array (H, W, C)\n",
    "    :param sigma: Standard deviation of the noise, in pixel values\n",
    "    :param seed: Seed of the random number generator\n",
    "    :return: The noised image as a NumPy array (H, W, C)\n",
    "    \"\"\"\n",
    "    rng = np.random.default_rng(seed)\n",
    "    noise = rng.normal(0.0, sigma, size=image.shape)\n",
    "    return np.clip(image + noise, 0, 255).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e34f6869",
   "metadata": {},
   "outputs": [],
   "source": [
    "compressed_image = compress_jpeg(image_1, JPEG_QUALITY)\n",
    "noised_image = add_gaussian_noise(image_1, NOISE_SIGMA)\n",
    "plot_image(compressed_image, title=f\"Image 1, JPEG quality {JPEG_QUALITY}\")\n",
    "plot_image(noised_image, title=f\"Image 1, Gaussian noise with sigma {NOISE_SIGMA}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3fc7e21",
   "metadata": {},
   "source": [
    "### Example: Computing SSIM and PSNR\n",
    "\n",
    "Now, we compare `image_1` with its compressed and its noised version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f41e9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute similarity using SSIM\n",
    "ssim = SSIM()\n",
    "\n",
    "print(\"SSIM, compressed image:\", ssim.similarity_score(image_1, compressed_image))\n",
    "print(\"SSIM, noised image:\", ssim.similarity_score(image_1, noised_image))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda90bc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute similarity using PSNR\n",
    "psnr = PSNR()\n",
    "\n",
    "print(\"PSNR, compressed image:\", psnr.similarity_score(image_1, compressed_image))\n",
    "print(\"PSNR, noised image:\", psnr.similarity_score(image_1, noised_image))"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}