{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before starting this notebook make sure that:\n",
    "1. You have generated datasets for google-siglip2-base-patch16-512 model 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda:0'\n",
    "vision_model_name = 'google/siglip2-base-patch16-512'\n",
    "feature_extractor_weights_dir = '../../feature_extractor_weights'\n",
    "reconstructor_weights_path = '../precalculated_weights/models--google--siglip2-base-patch16-512.pt'\n",
    "json_path = f'../../generated_datasets/{\"-\".join(vision_model_name.split(\"/\"))}/map_val.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's verify that the files and directories exist:\n",
    "\n",
    "The feature_extractor_weights_dir folder and the json_path file should appear after the dataset generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dir with feature extractor weights:  True\n",
      "file with weights of reconstructor R:  True\n",
      "file with pathes to data:  True\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "print('dir with feature extractor weights: ', os.path.isdir(feature_extractor_weights_dir))\n",
    "print('file with weights of reconstructor R: ', os.path.isfile(reconstructor_weights_path))\n",
    "print('file with pathes to data: ', os.path.isfile(json_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading the feature extractor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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",
      "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from transformers import SiglipVisionConfig, AutoImageProcessor, SiglipVisionModel\n",
    "\n",
    "class SigLipVisionTower(nn.Module):\n",
    "    def __init__(self, vision_model_name, weights_dir):\n",
    "        super().__init__()\n",
    "\n",
    "        self.is_loaded = False\n",
    "\n",
    "        self.vision_tower_name = vision_model_name\n",
    "        self.weights_dir = weights_dir\n",
    "        self.processor_path = 'facebook/dinov2-base'\n",
    "        self.select_layer = -2\n",
    "        self.cfg_only = SiglipVisionConfig.from_pretrained(self.vision_tower_name, cache_dir=self.weights_dir)\n",
    "\n",
    "    def load_model(self, device_map=None):\n",
    "        if self.is_loaded:\n",
    "            return\n",
    "\n",
    "        self.image_processor = AutoImageProcessor.from_pretrained(self.processor_path, \n",
    "                                                                  cache_dir=self.weights_dir, \n",
    "                                                                  crop_size={\"height\": self.cfg_only.image_size, \n",
    "                                                                             \"width\": self.cfg_only.image_size},\n",
    "                                                                  image_mean=[0.5, 0.5, 0.5], \n",
    "                                                                  image_std=[0.5, 0.5, 0.5], \n",
    "                                                                  size={\"shortest_edge\": self.cfg_only.image_size})\n",
    "        self.vision_tower = SiglipVisionModel.from_pretrained(self.vision_tower_name, \n",
    "                                                              cache_dir=self.weights_dir, \n",
    "                                                              device_map=device_map)\n",
    "\n",
    "        self.vision_tower.vision_model.head = nn.Identity()\n",
    "        self.vision_tower.requires_grad_(False)\n",
    "        self.eval()\n",
    "\n",
    "        self.is_loaded = True\n",
    "\n",
    "    def forward(self, images):\n",
    "        if type(images) is list:\n",
    "            image_features = []\n",
    "            for image in images:\n",
    "                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)\n",
    "                image_feature = image_forward_out.hidden_states[self.select_layer].to(image.dtype)\n",
    "                image_features.append(image_feature)\n",
    "        return image_features\n",
    "\n",
    "    @property\n",
    "    def dtype(self):\n",
    "        for p in self.vision_tower.parameters():\n",
    "            return p.dtype\n",
    "\n",
    "    @property\n",
    "    def device(self):\n",
    "        for p in self.vision_tower.parameters():\n",
    "            return p.device\n",
    "\n",
    "    @property\n",
    "    def config(self):\n",
    "        if self.is_loaded:\n",
    "            return self.vision_tower.config\n",
    "        else:\n",
    "            return self.cfg_only\n",
    "\n",
    "    @property\n",
    "    def hidden_size(self):\n",
    "        return self.config.hidden_size\n",
    "\n",
    "    @property\n",
    "    def num_patches(self):\n",
    "        return (self.config.image_size // self.config.patch_size) ** 2\n",
    "\n",
    "    @property\n",
    "    def num_patches_per_side(self):\n",
    "        return self.config.image_size // self.config.patch_size\n",
    "\n",
    "    @property\n",
    "    def image_size(self):\n",
    "        return self.config.image_size\n",
    "    \n",
    "vision_tower = SigLipVisionTower(vision_model_name, feature_extractor_weights_dir)\n",
    "vision_tower.load_model(device_map=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For each image, I create an analog of it by swapping the r and b channels in places.\n",
    "\n",
    "On five pictures I will train and on another five I will check the result. Turns out five pictures are enough to learn a turn!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:00<00:00, 166.52it/s]\n",
      "100%|██████████| 10/10 [00:00<00:00, 192.98it/s]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def rgb_to_bgr(im):\n",
    "    gbr_im = np.array(im)[:, :, ::-1]\n",
    "    gbr_im = Image.fromarray(gbr_im).convert('RGB')\n",
    "    return gbr_im\n",
    "\n",
    "\n",
    "def load_and_process_images(json_path, processor, number):\n",
    "    with open(json_path, 'r') as json_file:\n",
    "        json_dict = json.load(json_file)\n",
    "\n",
    "    ###################### images_list #########################\n",
    "    images_pathes_list = []\n",
    "    for i, (im_path, _) in enumerate(json_dict.items()):\n",
    "        if i == number: break\n",
    "        images_pathes_list.append(im_path)\n",
    "    images_list = [Image.open(p).convert('RGB') for p in tqdm(images_pathes_list)]\n",
    "\n",
    "    ###################### gbr_images_list #########################\n",
    "    processed_images_list = []\n",
    "    for im in tqdm(images_list):\n",
    "        gbr_im = processor(im)\n",
    "        processed_images_list.append(gbr_im)\n",
    "    \n",
    "    return images_list, processed_images_list\n",
    "\n",
    "N_train, N_val = 5, 5\n",
    "all_images_list, all_processed_images_list = load_and_process_images(json_path, rgb_to_bgr, number=N_train + N_val)\n",
    "\n",
    "images_list, processed_images_list = all_images_list[-N_train:], all_processed_images_list[-N_train:]\n",
    "val_images_list, val_processed_images_list = all_images_list[:N_val], all_processed_images_list[:N_val]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Counting features for initial and modified images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]/home/jovyan/.mlspace/envs/nips-interpretation/lib/python3.10/site-packages/torch/nn/modules/conv.py:456: UserWarning: Plan failed with a cudnnException: CUDNN_BACKEND_EXECUTION_PLAN_DESCRIPTOR: cudnnFinalize Descriptor Failed cudnn_status: CUDNN_STATUS_NOT_SUPPORTED (Triggered internally at /home/conda/feedstock_root/build_artifacts/libtorch_1716578970623/work/aten/src/ATen/native/cudnn/Conv_v8.cpp:919.)\n",
      "  return F.conv2d(input, weight, bias, self.stride,\n",
      "100%|██████████| 5/5 [00:02<00:00,  1.81it/s]\n",
      "100%|██████████| 5/5 [00:00<00:00, 30.78it/s]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from torch.nn import functional as F\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def calc_feature(me):\n",
    "    inputs = vision_tower.image_processor(me, return_tensors=\"pt\")['pixel_values'].to(device)[0]\n",
    "    features = vision_tower([inputs])\n",
    "\n",
    "    features = features[0].reshape(32, 32, 768)\n",
    "    features = features.permute(2, 0, 1)\n",
    "    features = features / features.norm(dim=0, keepdim=True)\n",
    "    return inputs, features\n",
    "\n",
    "\n",
    "def calc_features_for_images_and_processed_images(images_list, processed_images_list):\n",
    "    inputs_list = []\n",
    "    feature_list = []\n",
    "    for im in tqdm(images_list):\n",
    "        inputs, features = calc_feature(im)\n",
    "        inputs_list.append(inputs.cpu())\n",
    "        feature_list.append(features.cpu())\n",
    "\n",
    "    processed_inputs_list = []\n",
    "    processed_feature_list = []\n",
    "    for im in tqdm(processed_images_list):\n",
    "        inputs, features = calc_feature(im)\n",
    "        processed_inputs_list.append(inputs.cpu())\n",
    "        processed_feature_list.append(features.cpu())\n",
    "    \n",
    "    return inputs_list, feature_list, processed_inputs_list, processed_feature_list\n",
    "\n",
    "inputs_list, feature_list, processed_inputs_list, processed_feature_list = calc_features_for_images_and_processed_images(images_list, processed_images_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = torch.stack(feature_list).to(device)\n",
    "processed_feature = torch.stack(processed_feature_list).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the orthogonal matrix on the training sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_Q(real_F, gbr_F): # [1024, 1024]\n",
    "    M = real_F.permute(1, 0) @ gbr_F\n",
    "    U, S, Vh = torch.linalg.svd(M)\n",
    "    Q = U @ Vh\n",
    "    return Q\n",
    "\n",
    "Q = calculate_Q(features.permute(2, 3, 0, 1).flatten(0, 2),\n",
    "                processed_feature.permute(2, 3, 0, 1).flatten(0, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply the transformation and restore the changed fiches on the validation sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_Q(Fe, Q, n=1):\n",
    "    gbr_Fe = torch.zeros_like(Fe)\n",
    "\n",
    "    Qn = Q\n",
    "    for _ in range(n - 1):\n",
    "        Qn = Qn @ Q\n",
    "\n",
    "    for i in range(Fe.shape[1]):\n",
    "        for j in range(Fe.shape[2]):\n",
    "            gbr_Fe[:, i, j] = Fe[:, i, j] @ Qn\n",
    "    return gbr_Fe\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def interpolate(im, Fe, model, device):\n",
    "    processed_image = torch.tensor(im).to(device)\n",
    "    processed_Fe = Fe.to(device)\n",
    "    reconstructed_image, processed_Fe, check_dict = model.forward(processed_Fe, processed_image)\n",
    "    return processed_image, reconstructed_image, processed_Fe, check_dict\n",
    "\n",
    "image_mean = np.array(vision_tower.image_processor.image_mean)\n",
    "image_std = np.array(vision_tower.image_processor.image_std)\n",
    "def from_1_to_255(image, image_mean, image_std):\n",
    "    image_unnormed = np.asarray(image) * image_std[:, None, None] + image_mean[:, None, None]\n",
    "    image_01 = image_unnormed.clip(0, 1)\n",
    "    image_255 = image_01 * 255\n",
    "    image_255_uint = np.array(image_255, dtype=np.uint8)\n",
    "    return image_255_uint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import R\n",
    "model = R(vision_model_name, feature_extractor_weights_dir).to(device)\n",
    "model.load_state_dict(torch.load(reconstructor_weights_path, weights_only=True), strict=True)\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:00<00:00, 10.74it/s]\n"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "from PIL import Image\n",
    "\n",
    "def Qed(images_list, Q, n=1):\n",
    "    inputs_list, features_list, Qed_features_list, processed_inputs_list, processed_features_list = [], [], [], [], []\n",
    "\n",
    "    for im in tqdm(images_list):\n",
    "        inputs, features = calc_feature(im)\n",
    "        processed_inputs, processed_features = calc_feature(rgb_to_bgr(im))\n",
    "        Qed_me_features = apply_Q(features, Q, n)\n",
    "        \n",
    "        inputs_list.append(inputs)\n",
    "        features_list.append(features)\n",
    "        processed_inputs_list.append(processed_inputs)\n",
    "        processed_features_list.append(processed_features)\n",
    "        Qed_features_list.append(Qed_me_features)\n",
    "\n",
    "    return inputs_list, features_list, Qed_features_list, processed_inputs_list, processed_features_list\n",
    "\n",
    "inputs_list, features_list, Qed_features_list, processed_inputs_list, processed_features_list = Qed(val_images_list[:10], Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5 [00:00<?, ?it/s]/tmp/ipykernel_1759902/3103030834.py:16: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  processed_image = torch.tensor(im).to(device)\n",
      "100%|██████████| 5/5 [00:01<00:00,  4.26it/s]\n"
     ]
    }
   ],
   "source": [
    "def interpolate_all_features(inputs_list, features_list, Qed_features_list, \n",
    "                             processed_inputs_list, processed_features_list, \n",
    "                             get_features_of_recs=False, get_features_of_processed_recs=False):\n",
    "    inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list = [], [], [], [], []\n",
    "    if get_features_of_recs: features_of_recs_list = []\n",
    "    if get_features_of_processed_recs: features_of_processed_recs_list = []\n",
    "    \n",
    "    z = zip(tqdm(inputs_list), features_list, Qed_features_list, processed_inputs_list, processed_features_list)\n",
    "    for inputs, features, Qed_features, processed_inputs, processed_features in z:\n",
    "        inter, rec, _, _ = interpolate(inputs[None], features[None], model, device)\n",
    "        _, Qed_rec, _, _ = interpolate(inputs[None], Qed_features[None], model, device)\n",
    "        processed_inter, processed_rec, _, _ = interpolate(processed_inputs[None], processed_features[None], model, device)\n",
    "\n",
    "        inter = from_1_to_255(inter[0].cpu(), image_mean, image_std)\n",
    "        rec = from_1_to_255(rec[0].cpu(), image_mean, image_std)\n",
    "        processed_inter = from_1_to_255(processed_inter[0].cpu(), image_mean, image_std)\n",
    "        processed_rec = from_1_to_255(processed_rec[0].cpu(), image_mean, image_std)\n",
    "        Qed_rec = from_1_to_255(Qed_rec[0].cpu(), image_mean, image_std)\n",
    "\n",
    "        inter_list.append(inter)\n",
    "        rec_list.append(rec)\n",
    "        processed_inter_list.append(processed_inter)\n",
    "        processed_rec_list.append(processed_rec)\n",
    "        Qed_rec_list.append(Qed_rec)\n",
    "\n",
    "        if get_features_of_recs:\n",
    "            feature_of_rec = calc_feature(Image.fromarray(Qed_rec.transpose(1, 2, 0)))\n",
    "            features_of_recs_list.append(feature_of_rec[-1].cpu())\n",
    "        if get_features_of_processed_recs:\n",
    "            features_of_processed_recs = calc_feature(Image.fromarray(processed_rec.transpose(1, 2, 0)))\n",
    "            features_of_processed_recs_list.append(features_of_processed_recs[-1].cpu())\n",
    "    \n",
    "    if get_features_of_recs and (not get_features_of_processed_recs):\n",
    "        return inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list, features_of_recs_list\n",
    "    if (not get_features_of_recs) and get_features_of_processed_recs:\n",
    "        return inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list, features_of_processed_recs_list\n",
    "    if get_features_of_recs and get_features_of_processed_recs:\n",
    "        return inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list, features_of_recs_list, features_of_processed_recs_list\n",
    "    \n",
    "    return inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list\n",
    "\n",
    "inter_list, rec_list, processed_inter_list, processed_rec_list, Qed_rec_list, features_of_recs_list, features_of_processed_recs_list = \\\n",
    "interpolate_all_features(inputs_list, features_list, Qed_features_list, \n",
    "                         processed_inputs_list, processed_features_list, \n",
    "                         get_features_of_recs=True, get_features_of_processed_recs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Draw the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2284892 characters]",
      "text/plain": [
       "<Figure size 1500x1800 with 25 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "def draw(inter_list, rec_list, \n",
    "         processed_inter_list, processed_rec_list, \n",
    "         Qed_rec_list, \n",
    "         Qed_features_list=None, processed_features_list=None, features_of_recs_list=None, features_of_processed_recs_list=None):\n",
    "    ncols, nrows, scale = 5, len(inter_list), 3\n",
    "    fig, axes = plt.subplots(ncols=ncols, nrows=nrows, figsize=(ncols * scale, (nrows + 1) * scale))\n",
    "\n",
    "    for i in range(nrows):\n",
    "        axes[i, 0].imshow(inter_list[i].transpose(1, 2, 0))\n",
    "        axes[i, 1].imshow(rec_list[i].transpose(1, 2, 0))\n",
    "        axes[i, 2].imshow(processed_inter_list[i].transpose(1, 2, 0))\n",
    "        axes[i, 3].imshow(processed_rec_list[i].transpose(1, 2, 0))\n",
    "        axes[i, 4].imshow(Qed_rec_list[i].transpose(1, 2, 0))\n",
    "\n",
    "        axes[i, 0].axis('off')\n",
    "        axes[i, 1].axis('off')\n",
    "        axes[i, 2].axis('off')\n",
    "        axes[i, 3].axis('off')\n",
    "        axes[i, 4].axis('off')\n",
    "\n",
    "        title_4 = ''\n",
    "        if i == 0:\n",
    "            axes[i, 0].set_title('image')\n",
    "            axes[i, 1].set_title('reconstruction from\\n'\n",
    "                                 'features of image')\n",
    "            axes[i, 2].set_title('image with swapped\\n'\n",
    "                                 'r and g channels')\n",
    "            axes[i, 3].set_title('reconstruction from\\n'\n",
    "                                 'image with swapped\\n'\n",
    "                                 'r and g channels')\n",
    "            axes[i, 4].set_title('reconstruction from\\n'+\\\n",
    "                                 'features of image,\\n'+\\\n",
    "                                 ' which was modified\\n'+\\\n",
    "                                 'by orthogonal transformation')\n",
    "        \n",
    "        if (Qed_features_list is not None) and (processed_features_list is not None):\n",
    "            x1 = Qed_features_list[i].cpu()\n",
    "            x2 = processed_features_list[i].cpu()\n",
    "            mean_cos = round((x1 * x2).sum(dim=0).mean().item(), 2)\n",
    "\n",
    "        if (Qed_features_list is not None) and (features_of_recs_list is not None):\n",
    "            x1 = Qed_features_list[i].cpu()\n",
    "            x2 = features_of_recs_list[i].cpu()\n",
    "            mean_cos = round((x1 * x2).sum(dim=0).mean().item(), 2)\n",
    "\n",
    "        if (Qed_features_list is not None) and (features_of_processed_recs_list is not None):\n",
    "            x1 = Qed_features_list[i].cpu()\n",
    "            x2 = features_of_processed_recs_list[i].cpu()\n",
    "            mean_cos = round((x1 * x2).sum(dim=0).mean().item(), 2)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "draw(inter_list, rec_list, \n",
    "     processed_inter_list, processed_rec_list, \n",
    "     Qed_rec_list, \n",
    "     Qed_features_list, processed_features_list, features_of_recs_list, features_of_processed_recs_list)\n"
   ]
  },
  {
   "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
}
