{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "179f6bb1-ef6f-4328-b857-d85740b69b9b",
   "metadata": {},
   "source": [
    "# Compare 2 videos with SSIM, PSNR and VMAF metrics\n",
    "\n",
    "The classic metrics **PSNR** (Peak Signal to Noise Ratio) and **SSIM** (Structural SIMilarity) are implemented in C for improved performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "be964f70-cf5f-4fd5-871e-188f4955e2d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import pathlib\n",
    "import timeit\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "import tqdm\n",
    "\n",
    "from cutcutcodec.core.io import read\n",
    "from cutcutcodec.core.opti.parallel import starmap\n",
    "from cutcutcodec.core.analysis.video.metric import psnr as compute_psnr, ssim as compute_ssim, vmaf as compute_vmaf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5c31cef-e0e1-4ced-8aa1-8e5e6c3c9b26",
   "metadata": {},
   "source": [
    "## Preparing videos to be compared\n",
    "Here we'll compare the original video with a noisy version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c2dab0a-8e62-4d65-b86f-bd15fed6426c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/rrichard/cutcutcodec_git/media/video/intro.webm at 30000/1001 fps, of size 1280x720\n"
     ]
    }
   ],
   "source": [
    "from cutcutcodec.core.analysis.stream.rate_video import optimal_rate_video\n",
    "from cutcutcodec.core.analysis.stream.shape import optimal_shape_video\n",
    "from cutcutcodec.utils import get_project_root\n",
    "\n",
    "file = get_project_root().parent / \"media\" / \"video\" / \"intro.webm\"\n",
    "stream_ref = read(file).out_select(\"video\")[0]\n",
    "rate = optimal_rate_video(stream_ref)\n",
    "shape = optimal_shape_video(stream_ref)\n",
    "print(f\"{file} at {rate} fps, of size {shape[1]}x{shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "85eacf9c-3256-4437-9e27-84bc94b8cb91",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cutcutcodec.core.generation.video.noise import GeneratorVideoNoise\n",
    "\n",
    "stream_comp = (stream_ref | GeneratorVideoNoise()).apply_video_equation(\"(9*r0+r1)/10\", \"(9*g0+g1)/10\", \"(9*b0+b1)/10\").out_streams[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a45ac754-6dc3-42be-90a4-a76f7a554750",
   "metadata": {},
   "source": [
    "## Compare each frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a33c4311-3cab-4291-9fe9-949c47e6e5b4",
   "metadata": {},
   "source": [
    "### Convert all streams into YUV in the same colorspace\n",
    "Image comparisons are always made in yuv, not rgb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e7939ef-ab7e-4343-8b14-ec971c74169c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cutcutcodec.core.filter.video.colorspace import FilterVideoColorspace\n",
    "from cutcutcodec.config import Config  # to get the default used colorspace\n",
    "\n",
    "dst_color = f\"y'pbpr_{Config().target_trc}_{Config().target_prim}\"\n",
    "ref_yuv = FilterVideoColorspace([stream_ref], dst=dst_color).out_streams[0]\n",
    "comp_yuv = FilterVideoColorspace([stream_comp], dst=dst_color).out_streams[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "88798d3c-8fcf-4c13-8ac3-2bdc4a58ac63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare(stream1, stream2, time, shape):\n",
    "    \"\"\"Compare 2 images.\"\"\"\n",
    "    frame1, frame2 = stream1.snapshot(time, shape), stream2.snapshot(time, shape)\n",
    "    psnr = compute_psnr(frame1, frame2, weights=(6, 1, 1), threads=1)  # factor 6 on y\n",
    "    ssim = compute_ssim(frame1, frame2, weights=(6, 1, 1), data_range=1.0, threads=1)\n",
    "    vmaf = compute_vmaf(frame1, frame2, threads=1)\n",
    "    return psnr, ssim, vmaf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edd46ea0-7962-4813-8c60-b9ab204ff3e5",
   "metadata": {},
   "source": [
    "## Compute metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d715ff1e-3c0d-4586-9e88-162310007efc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██████████████████████████████████████▋                                                                                                           | 78/294 [01:53<01:22,  2.61it/s]"
     ]
    }
   ],
   "source": [
    "times = np.arange(0, ref_yuv.duration, 1/rate).tolist()\n",
    "all_psnr, all_ssim, all_vmaf = zip(*starmap(compare, ((ref_yuv, comp_yuv, t, shape) for t in tqdm.tqdm(times))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbf03e40-a89e-44b9-87a1-dd562214a845",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(times, all_psnr)\n",
    "plt.xlabel(\"time (s)\")\n",
    "plt.ylabel(\"psnr (db)\")\n",
    "plt.show()\n",
    "\n",
    "plt.plot(times, all_ssim)\n",
    "plt.xlabel(\"time (s)\")\n",
    "plt.ylabel(\"ssim\")\n",
    "plt.show()\n",
    "\n",
    "plt.plot(times, all_vmaf)\n",
    "plt.xlabel(\"time (s)\")\n",
    "plt.ylabel(\"vmaf (db)\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}