{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from PIL import Image\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class FeaturesDataset(Dataset):\n",
    "    def __init__(self, features_json_path, mode='train', side_size=None, image_processor=None):\n",
    "        with open(features_json_path, 'r') as map_file:\n",
    "            self.map = list(json.load(map_file).items())\n",
    "        self.mode = mode\n",
    "\n",
    "        self.side_size = side_size\n",
    "        self.image_processor = image_processor\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.map)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_path, feature_path = self.map[idx]\n",
    "        feature = torch.load(feature_path, map_location='cpu').to(torch.float32)\n",
    "        \n",
    "        image = Image.open(image_path).convert('RGB')\n",
    "        \n",
    "        processed_image = self.image_processor(image)['pixel_values'][0]\n",
    "        processed_image = torch.tensor(processed_image)\n",
    "        if processed_image.shape[-1] != self.side_size:\n",
    "            interpolated_image = F.interpolate(processed_image[None], size=(self.side_size, self.side_size), mode='bilinear', align_corners=False)[0]\n",
    "        else:\n",
    "            interpolated_image = processed_image\n",
    "        \n",
    "        if self.mode == 'train':\n",
    "            return interpolated_image, feature.permute(2, 0, 1) # tensor[3 x 224 x 224], tensor[768 x 14 x 14]\n",
    "        else:\n",
    "            return np.asarray(image), interpolated_image, feature.permute(2, 0, 1) # array(427, 640, 3), tensor[3 x 224 x 224], tensor[768 x 14 x 14]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jovyan/.mlspace/envs/nips-interpretation/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from transformers import SiglipImageProcessor\n",
    "\n",
    "side_size=256\n",
    "vision_model_name = 'google/siglip2-base-patch16-512'\n",
    "features_json_path = os.path.join(os.getcwd(), f'../generated_datasets/{\"-\".join(vision_model_name.split(\"/\"))}/map_train.json')\n",
    "weights_dir = os.path.join(os.getcwd(), '../feature_extractor_weights')\n",
    "\n",
    "image_processor = SiglipImageProcessor.from_pretrained(vision_model_name, cache_dir=weights_dir)\n",
    "val_dataset = FeaturesDataset(features_json_path, 'val', side_size, image_processor)\n",
    "print(len(val_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1397476 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 50 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "im_mean = torch.as_tensor(image_processor.image_mean)\n",
    "im_std = torch.as_tensor(image_processor.image_std)\n",
    "\n",
    "nrows, ncols, scale = 5, 10, 3\n",
    "fig, axes = plt.subplots(ncols=ncols, nrows=nrows, figsize=(ncols * scale, nrows * scale))\n",
    "\n",
    "for i in range(nrows):\n",
    "    im, inter_im, t = val_dataset[i]\n",
    "    inter_im = inter_im.permute(1, 2, 0)\n",
    "    inter_im = inter_im * im_std[None, None] + im_mean[None, None]\n",
    "\n",
    "    axes[i, 0].imshow(im)\n",
    "    axes[i, 1].imshow(inter_im)\n",
    "    for j in range(2, ncols):\n",
    "        axes[i, j].imshow(t[t.shape[0] // (ncols - 2) * (j - 2)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
