{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before starting this notebook make sure that:\n",
    "1. You have generated datasets for all models using the ../../dataset_generation/generation.sh script\n",
    "2. You have trained your weights or taken already trained weights and placed them in the ../precalculated_weights folder\n",
    "3. You have calculated the clip similarity for all models using the script ./calculate_similarity.sh\n",
    "\n",
    "List of models:\n",
    "    \n",
    "    * google-siglip-base-patch16-224\n",
    "\n",
    "    * google-siglip-base-patch16-256\n",
    "    \n",
    "    * google-siglip-base-patch16-384\n",
    "    \n",
    "    * google-siglip-base-patch16-512\n",
    "    \n",
    "    * google-siglip2-base-patch16-224\n",
    "    \n",
    "    * google-siglip2-base-patch16-256\n",
    "    \n",
    "    * google-siglip2-base-patch16-384\n",
    "    \n",
    "    * google-siglip2-base-patch16-512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "siglips = [\n",
    "    'google-siglip-base-patch16-224',\n",
    "    'google-siglip-base-patch16-256',\n",
    "    'google-siglip-base-patch16-384',\n",
    "    'google-siglip-base-patch16-512'\n",
    "]\n",
    "\n",
    "siglips2 = [\n",
    "    'google-siglip2-base-patch16-224',\n",
    "    'google-siglip2-base-patch16-256',\n",
    "    'google-siglip2-base-patch16-384',\n",
    "    'google-siglip2-base-patch16-512',\n",
    "]\n",
    "\n",
    "path = './results'\n",
    "\n",
    "first_10 = 'first_10.png'\n",
    "metrics_siglip_json = 'metrics_clip-10__google-siglip2-large-patch16-256.json'\n",
    "metrics_clip_json = 'metrics_clip-10__openai-clip-vit-large-patch14.json'\n",
    "\n",
    "model_names = os.listdir(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "\n",
    "siglips_siglip_pathes = [os.path.join(path, name, metrics_siglip_json) for name in siglips if '.ipynb' not in name]\n",
    "siglips_siglip_sim = {}\n",
    "for p in siglips_siglip_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        stat = json.load(f)['clip_sim']\n",
    "        mean = np.mean(stat)\n",
    "        std = np.std(stat)\n",
    "        siglips_siglip_sim[p.split('/')[-2]] = (mean, std)\n",
    "\n",
    "siglips_clip_pathes = [os.path.join(path, name, metrics_clip_json) for name in siglips if '.ipynb' not in name]\n",
    "siglips_clip_sim = {}\n",
    "for p in siglips_clip_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        stat = json.load(f)['clip_sim']\n",
    "        mean = np.mean(stat)\n",
    "        std = np.std(stat)\n",
    "        siglips_clip_sim[p.split('/')[-2]] = (mean, std)\n",
    "\n",
    "\n",
    "\n",
    "siglips2_siglip_pathes = [os.path.join(path, name, metrics_siglip_json) for name in siglips2 if '.ipynb' not in name]\n",
    "siglips2_siglip_sim = {}\n",
    "for p in siglips2_siglip_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        stat = json.load(f)['clip_sim']\n",
    "        mean = np.mean(stat)\n",
    "        std = np.std(stat)\n",
    "        siglips2_siglip_sim[p.split('/')[-2]] = (mean, std)\n",
    "\n",
    "siglips2_clip_pathes = [os.path.join(path, name, metrics_clip_json) for name in siglips2 if '.ipynb' not in name]\n",
    "siglips2_clip_sim = {}\n",
    "for p in siglips2_clip_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        stat = json.load(f)['clip_sim']\n",
    "        mean = np.mean(stat)\n",
    "        std = np.std(stat)\n",
    "        siglips2_clip_sim[p.split('/')[-2]] = (mean, std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "siglip_image_pathes = [os.path.join(path, n, first_10) for n in siglips if '.ipynb' not in n]\n",
    "siglip_images = {}\n",
    "for p in siglip_image_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        siglip_images[p.split('/')[-2]] = Image.open(p).convert('RGB')\n",
    "\n",
    "siglip2_image_pathes = [os.path.join(path, n, first_10) for n in siglips2 if '.ipynb' not in n]\n",
    "siglip2_images = {}\n",
    "for p in siglip2_image_pathes:\n",
    "    with open(p, 'r') as f:\n",
    "        siglip2_images[p.split('/')[-2]] = Image.open(p).convert('RGB')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "siglip_np_images = [np.asarray(siglip_images[n])[600:-600, 90:-90] for n in siglips]\n",
    "siglip_recs = np.concatenate(siglip_np_images, axis=1)\n",
    "\n",
    "siglip2_np_images = [np.asarray(siglip2_images[n])[600:-600, 90:-90] for n in siglips2]\n",
    "siglip2_recs = np.concatenate(siglip2_np_images, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graph below may differ from the graph from the article because it is plotted using only 10 examples. \n",
    "However, even for the ten examples, we can see that:\n",
    "1. siglip2 beats siglip in all cases.\n",
    "2. As the resolution of the phyche tensor increases, the quality of the reconstruction increases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10718920 characters]",
      "text/plain": [
       "<Figure size 3600x4800 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "ncols, nrows, scale = 12, 16, 3\n",
    "fig = plt.figure(figsize=(ncols * scale, nrows * scale))\n",
    "gs_main = fig.add_gridspec(3, 1, height_ratios=[1, 1, 8])\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "ax_top = plt.subplot(gs_main[0])\n",
    "width=0.1\n",
    "ax_top.bar([n - width / 2 for n in range(len(siglips))],\n",
    "        [siglips_clip_sim[n][0] for n in siglips],\n",
    "        width=width, label='SigLIP base patch16')\n",
    "ax_top.bar([n + width / 2 for n in range(len(siglips2))],\n",
    "        [siglips2_clip_sim[n][0] for n in siglips2],\n",
    "        width=width, label='SigLIP2 base patch16')\n",
    "ax_top.set_yticks(np.arange(0.6, 0.8, 0.025).round(3), np.arange(0.6, 0.8, 0.025).round(3), fontsize=20)\n",
    "ax_top.set_xticks([n for n in range(len(siglips))], [n.split('-')[-1] for n in siglips], fontsize=20)\n",
    "\n",
    "ax_top.set_xlim(-0.5, 3.5)\n",
    "ax_top.set_ylim(0.6, 0.8)\n",
    "ax_top.set_xlabel('image size', fontsize=30)\n",
    "ax_top.set_ylabel('CLIP large\\nsimilarity', fontsize=30)\n",
    "ax_top.legend(fontsize=30)\n",
    "\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "ax_medium = plt.subplot(gs_main[1])\n",
    "width=0.1\n",
    "ax_medium.bar([n - width / 2 for n in range(len(siglips))],\n",
    "        [siglips_siglip_sim[n][0] for n in siglips],\n",
    "        width=width, label='SigLIP base patch16')\n",
    "ax_medium.bar([n + width / 2 for n in range(len(siglips2))],\n",
    "        [siglips2_siglip_sim[n][0] for n in siglips2],\n",
    "        width=width, label='SigLIP2 base patch16')\n",
    "ax_medium.set_yticks(np.arange(0.6, 0.8, 0.025).round(3), np.arange(0.6, 0.8, 0.025).round(3), fontsize=20)\n",
    "ax_medium.set_xticks([n for n in range(len(siglips))], [n.split('-')[-1] for n in siglips], fontsize=20)\n",
    "\n",
    "ax_medium.set_xlim(-0.5, 3.5)\n",
    "ax_medium.set_ylim(0.6, 0.8)\n",
    "ax_medium.set_xlabel('image size', fontsize=30)\n",
    "ax_medium.set_ylabel('SigLIP 2 large\\nsimilarity', fontsize=30)\n",
    "ax_medium.legend(fontsize=30)\n",
    "\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "#################################################################################\n",
    "d, p = 0.4, 0.2\n",
    "gs_bottom = gs_main[2].subgridspec(1, 13, width_ratios=[p, 1, 1, d, 1, 1, d, 1, 1, d, 1, 1, p], wspace=0.1)\n",
    "\n",
    "#################################################################################\n",
    "siglip_np_images = [np.asarray(siglip_images[n])[600:-600, 90:-90] for n in siglips]\n",
    "siglip_recs = np.concatenate(siglip_np_images, axis=1)\n",
    "for i, name, im in zip([0, 3, 6, 9], ['SigLIP: 224', 'SigLIP: 256', 'SigLIP: 384', 'SigLIP: 512'],siglip_np_images):\n",
    "    ax2 = fig.add_subplot(gs_bottom[0, i + 1])\n",
    "    ax2.imshow(im)\n",
    "    ax2.set_title(name, fontsize=20)\n",
    "    ax2.axis('off')\n",
    "\n",
    "#################################################################################\n",
    "siglip2_np_images = [np.asarray(siglip2_images[n])[600:-600, 90:-90] for n in siglips2]\n",
    "siglip2_recs = np.concatenate(siglip2_np_images, axis=1)\n",
    "for i, name, im in zip([1, 4, 7, 10], ['SigLIP 2: 224', 'SigLIP 2: 256', 'SigLIP 2: 384', 'SigLIP 2: 512'], siglip2_np_images):\n",
    "    ax2 = fig.add_subplot(gs_bottom[0, i + 1])\n",
    "    ax2.imshow(im)\n",
    "    ax2.set_title(name, fontsize=20)\n",
    "    ax2.axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nips-interpretation",
   "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.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
