test_transform.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9845754",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "\n",
    "path = \"/home/someone/stage_jonathan/lightning-sam/lightning-sam/l_sam_env/lib/python3.10/site-packages\"\n",
    "\n",
    "if path not in sys.path:\n",
    "    sys.path.append(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c763f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "502c5a5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_img = cfg.dataset.img_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eef593f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import cv2\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0538086",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = os.listdir(path_to_img)[150]\n",
    "print(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e81298e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "file+\"/\"+os.listdir(os.path.join(path_to_img,file))[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f22ff99",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = cv2.imread(os.path.join(path_to_img+\"/\"+file, \n",
    "                              os.listdir(os.path.join(path_to_img,\n",
    "                                                      file))[6]))\n",
    "\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d34a229d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c04df9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "box_file = cfg.dataset.box_file\n",
    "img_dir = cfg.dataset.img_file\n",
    "\n",
    "\n",
    "with open(box_file,\"r\") as f:\n",
    "    temp = f.readlines()\n",
    "\n",
    "    # format i x1 y1 x2 y2 \n",
    "    boxes = [list(map(lambda x : int(float(x)), line.split()[1:])) for line in temp] \n",
    "\n",
    "with open(img_dir,\"r\") as f:\n",
    "    temp = f.readlines()\n",
    "    # format i folder_name/image_name\n",
    "    img_names = [os.path.join(path_to_img,line.split()[1]) for line in temp]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43308a80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from custom_transform import LocalizedRandomResizedCrop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391ab0d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = [(i,name) for (i,name) in enumerate(img_names) if  file+\"/\"+os.listdir(os.path.join(path_to_img,file))[6] in name ][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96da681e",
   "metadata": {},
   "outputs": [],
   "source": [
    "boxes[idx[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65206b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert to PIL\n",
    "from PIL import Image\n",
    "from matplotlib.patches import Rectangle\n",
    "\n",
    "img_pil = Image.fromarray(img)\n",
    "\n",
    "# (x,y) , w, h\n",
    "fig, ax = plt.subplots()\n",
    "ax.imshow(img_pil)\n",
    "x,y,xw,yh = boxes[idx[0]]\n",
    "\n",
    "\n",
    "rect = Rectangle((x,y),xw,yh, linewidth=1, edgecolor='r', facecolor='none')\n",
    "ax.add_patch(rect)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84bfcd46",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as T\n",
    "import random as rd\n",
    "\n",
    "# T.RandomPerspective(distortion_scale=0.2, p=1, fill = 255*rd.random())\n",
    "#T.ElasticTransform(alpha=10.0,sigma=2.0)\n",
    "\n",
    "new_t = T.Compose([\n",
    "    T.Resize(224),\n",
    "    T.CenterCrop(224),])\n",
    "\n",
    "new = new_t(img_pil)\n",
    "\n",
    "print(new.size)\n",
    "plt.imshow(new)\n",
    "\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58bce4f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit \n",
    "T.AugMix(severity=2)(img_pil)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a03635f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "p= 0.5\n",
    "T.ColorJitter(brightness=p*0.5, contrast=p*0.5, saturation=p*0.5, hue=p*0.1)(img_pil)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596246e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18904f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "x,y,w,h = boxes[idx[0]]\n",
    "\n",
    "cropped = LocalizedRandomResizedCrop(img_pil,*[x,y,w,h],size=300, THR=1, relative_upper_scale=1.0, ratio=(1.0,1.0))\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(cropped)\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72f8603",
   "metadata": {},
   "outputs": [],
   "source": [
    "from compute_iou import get_miou"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "106d35f5",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "for thr in np.linspace(0,1,11):\n",
    "    local, rand = get_miou(thr)\n",
    "    \n",
    "    plt.hist(local, bins=100, alpha=0.5, label='local')\n",
    "    plt.hist(rand, bins=100, alpha=0.5, label='rand')\n",
    "    plt.legend(loc='upper left')\n",
    "    mean = np.mean(local)\n",
    "    print(mean)\n",
    "    plt.title(f\"comparison for thr = {thr:.4f}\")\n",
    "    plt.show()\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f511ba94",
   "metadata": {},
   "source": [
    "import numpy as np\n",
    "thr = 0.001\n",
    "\n",
    "local, rand = get_miou(thr)\n",
    "\n",
    "plt.figure(figsize=(16,10))\n",
    "plt.hist(local, bins=100, alpha=0.5, label='local')\n",
    "plt.hist(rand, bins=100, alpha=0.5, label='rand')\n",
    "plt.legend(loc='upper left')\n",
    "mean = np.mean(local)\n",
    "plt.title(f\"comparison for thr = {thr}, mean = {mean}\")\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6650b150",
   "metadata": {},
   "source": [
    "len(np.where(np.array(local)<0.9\n",
    "            )[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27057279",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import load_cub_datasets\n",
    "cfg.use_box = True\n",
    "cfg.THR = 1\n",
    "cfg.augment = True\n",
    "cfg.augment_p = 1\n",
    "cfg.use_augment_mix = False\n",
    "train_local, val_local = load_cub_datasets(cfg)\n",
    "\n",
    "\n",
    "cfg.THR = 1\n",
    "cfg.use_box = True\n",
    "cfg.augment = False\n",
    "train_rando, val_rando = load_cub_datasets(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be5cd4ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "for (imgs1, labels1), (imgs2, labels2) in zip(train_local,train_rando):\n",
    "    for img1, img2 in zip(imgs1,imgs2):\n",
    "        img1 = img1.permute((1,2,0))\n",
    "        img1 -= img1.min()\n",
    "        img1/=img1.max()\n",
    "        \n",
    "        img2 = img2.permute((1,2,0))\n",
    "        img2 -= img2.min()\n",
    "        img2/=img2.max()\n",
    "        \n",
    "        plt.subplot(1,2,1)\n",
    "        plt.imshow(img1)\n",
    "        plt.subplot(1,2,2)\n",
    "        plt.imshow(img2)\n",
    "        plt.show()\n",
    "\n",
    "        \n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c7021df",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import load_box_dataset\n",
    "from config import cfg\n",
    "import os\n",
    "\n",
    "seg = os.path.join(\n",
    "    os.path.dirname(cfg.dataset.img_dir),\n",
    "    \"masks_0\"\n",
    ")\n",
    "\n",
    "assert os.path.exists(seg)\n",
    "\n",
    "seg_file = os.path.join(seg,\n",
    "                       \"masks_info.json\")\n",
    "train, val = load_box_dataset(cfg,seg_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa1571ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "for img1s, idxs, labels in train:\n",
    "    for img1, idx, label in zip(img1s, idxs, labels):\n",
    "        print(idx,label)\n",
    "        img1 = img1.permute((1,2,0))\n",
    "        img1 -= img1.min()\n",
    "        img1/=img1.max()\n",
    "        plt.imshow(img1)\n",
    "        plt.show()\n",
    "\n",
    "        \n",
    "    break\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08c90643",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "maiou = np.linspace(0.4,1)\n",
    "THR = (2*(maiou-0.5))**(1/0.4)\n",
    "plt.grid()\n",
    "plt.plot(maiou,THR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6e30052",
   "metadata": {},
   "outputs": [],
   "source": [
    "(112*60)/2.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b424dfde",
   "metadata": {},
   "outputs": [],
   "source": [
    "1/0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "867a19e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = []\n",
    "Y = []\n",
    "for e in np.concatenate([np.linspace(0,0.2,10),np.linspace(0.21,0.7,10), np.linspace(0.8,1,3)]):\n",
    "    a, _ = get_miou(THR=e)\n",
    "    X.append(np.mean(a))\n",
    "    Y.append(e)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60aec44",
   "metadata": {},
   "outputs": [],
   "source": [
    "vec = np.mean(np.array(X),axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35774e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(vec) == len(Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8884f218",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ebc7880",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(14,14))\n",
    "plt.plot(Y,vec)\n",
    "plt.xticks(np.linspace(0,1, 21))\n",
    "plt.yticks(np.linspace(0.5,1, 21))\n",
    "plt.grid()\n",
    "plt.title(\"IOU assym. moyenne en fonction de THR\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1dc2f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "L = []\n",
    "for dir in os.listdir(cfg.dataset.img_dir):\n",
    "    L.append(len(os.listdir(os.path.join(cfg.dataset.img_dir,dir))))\n",
    "\n",
    "print(np.mean(L))\n",
    "print(np.std(L))\n",
    "print(np.min(L))\n",
    "print(np.max(L))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbc4e31a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "for name in [\"masks_0\",\"masks_1\",\"masks_2\", \"masks_3\"]:\n",
    "\n",
    "    root_dir = os.path.dirname(cfg.dataset.img_dir)\n",
    "\n",
    "    masks_dir = os.path.join(root_dir, name)\n",
    "\n",
    "    with open(os.path.join(masks_dir, \"masks_info.json\"), \"r\") as f:\n",
    "        d = json.load(f)\n",
    "\n",
    "    print(name)\n",
    "    print(d[\"setup\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee9c194",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(cfg.dataset.img_file,\"r\") as f:\n",
    "        temp = f.readlines()\n",
    "        # format i path\n",
    "        img_names = {int(line.split()[0]):line.split()[1] for line in temp} # index : path\n",
    "\n",
    "\n",
    "with open(cfg.dataset.label_file,\"r\") as f:\n",
    "            temp = f.readlines()\n",
    "            # format i class_number\n",
    "            labels = {int(line.split()[0]) : int(line.split()[1])-1 for line in temp}\n",
    "\n",
    "with open(cfg.dataset.box_file,\"r\") as f:\n",
    "    temp = f.readlines()\n",
    "    # format i x y w h\n",
    "    boxes = {int(line.split()[0]):list(map(lambda x : int(float(x)), line.split()[1:])) for line in temp}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3226832f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def compute_iou(bbox,gt,):\n",
    "\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        y, x, h, w\n",
    "    \"\"\"\n",
    "\n",
    "    x1, y1, w1, h1 = bbox\n",
    "    x3, y3, w3, h3 = gt\n",
    "\n",
    "    x2 = x1 + w1\n",
    "    y2 = y1 + h1\n",
    "\n",
    "    x4 = x3 + w3\n",
    "    y4 = y3 + h3\n",
    "\n",
    "    #  x1,y1,x2,y2,x3,y3,x4,y4\n",
    "\n",
    "    inter_width = min(x2, x4) - max(x1, x3)\n",
    "    inter_height = min(y2, y4) - max(y1, y3)\n",
    "\n",
    "    if inter_width <= 0 or inter_height <= 0:\n",
    "        return 0\n",
    "    \n",
    "    areaIntersection = inter_width * inter_height \n",
    "\n",
    "    areaUnion = w1 * h1 + w3 * h3 - areaIntersection\n",
    "\n",
    "    return areaIntersection / (areaUnion+1e-6)\n",
    "\n",
    "L_max_iou = []\n",
    "\n",
    "for idx in tqdm(range(1,1+len(img_names))):\n",
    "\n",
    "    name = img_names[idx]\n",
    "    box = boxes[idx]\n",
    "\n",
    "    L_max_iou.append(\n",
    "        max(\n",
    "            [compute_iou(box, d[name][i][\"bbox\"]) for i in range(len(d[name]))]\n",
    "        )\n",
    "    )\n",
    "\n",
    "print(np.mean(L_max_iou))\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5234609",
   "metadata": {},
   "outputs": [],
   "source": [
    "d[img_names[1]][4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b1e40b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torch.nn as nn\n",
    "\n",
    "model = torchvision.models.resnet50(weights=None)\n",
    "\n",
    "model.fc = nn.Linear(model.fc.in_features,cfg.model.n_classes)\n",
    "\n",
    "model.load_state_dict(torch.load(cfg.model.resumed_model))\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = model.to(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f838d41a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random as rd\n",
    "\n",
    "X = []\n",
    "\n",
    "Y = []\n",
    "\n",
    "for _ in range(50):\n",
    "\n",
    "    dummy_maiou = rd.uniform(0.5,1.05)\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "    if 0.5 <= dummy_maiou <= 1.0:\n",
    "                thr = (2*(dummy_maiou-0.5))**(1/0.4)\n",
    "\n",
    "                maiou = dummy_maiou\n",
    "\n",
    "    else: # 1 < dummy_maiou \n",
    "        thr = -1 # will not use box\n",
    "\n",
    "        maiou = 1.4 + (dummy_maiou-1.0)\n",
    "\n",
    "    X.append(maiou)\n",
    "\n",
    "    if maiou > 1:\n",
    "        y = 0.76 + rd.uniform(-0.005,0.005)\n",
    "    \n",
    "    else:\n",
    "        y = 0.76 + rd.uniform(-0.005,0.005)\n",
    "\n",
    "    Y.append(y)\n",
    "\n",
    "plt.scatter(X,Y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be38bbc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "import os\n",
    "\n",
    "from config import cfg\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "model = model.to(device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5385b8a",
   "metadata": {},
   "source": [
    "\"\"\"transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(1,1)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "]),\n",
    "transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(2,2)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "]),\n",
    "transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(3,3)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "]),\n",
    "transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(-1,-1)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "]),\n",
    "transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(-2,-2)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "]),\n",
    "transforms.Compose([\n",
    "transforms.RandomRotation(degrees=(-3,-3)),\n",
    "transforms.Resize((new_w, new_h), antialias=True),\n",
    "])\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3881f823",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(cfg.dataset.img_file,\"r\") as f:\n",
    "        temp = f.readlines()\n",
    "        # format i path\n",
    "        img_names = {int(line.split()[0]):line.split()[1] for line in temp} # index : path\n",
    "\n",
    "\n",
    "with open(cfg.dataset.label_file,\"r\") as f:\n",
    "            temp = f.readlines()\n",
    "            # format i class_number\n",
    "            labels = {int(line.split()[0]) : int(line.split()[1])-1 for line in temp}\n",
    "\n",
    "with open(cfg.dataset.box_file,\"r\") as f:\n",
    "    temp = f.readlines()\n",
    "    # format i x y w h\n",
    "    boxes = {int(line.split()[0]):list(map(lambda x : int(float(x)), line.split()[1:])) for line in temp}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f37d1677",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "#root / folder / img_name\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "\n",
    "root = cfg.dataset.img_dir\n",
    "folder = os.listdir(root)[1]\n",
    "img_name = os.listdir(os.path.join(root, folder))[19]\n",
    "\n",
    "path_to_img = os.path.join(root, folder, img_name)\n",
    "\n",
    "path_to_img = \"/mnt/data/CUB_200_2011/images/108.White_necked_Raven/White_Necked_Raven_0070_102645.jpg\"\n",
    "img = cv2.imread(path_to_img)\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "\n",
    "def preprocess_img(path_to_img, resize=False):\n",
    "    \"\"\"img = cv2.imread(path_to_img)\n",
    "\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "    img = Image.fromarray(img)\"\"\"\n",
    "\n",
    "    img = Image.open(path_to_img)\n",
    "    if img.mode != \"RGB\":\n",
    "        img = img.convert(\"RGB\")\n",
    "    \n",
    "    if resize:img = transforms.Resize(500, antialias=True)(img)\n",
    "    img = transforms.RandomAutocontrast(p=1)(img)\n",
    "\n",
    "    h, w = img.size\n",
    "    new_w = w - (w % 14)\n",
    "    new_h = h - (h % 14)\n",
    "\n",
    "\n",
    "    img = transforms.ToTensor()(img)\n",
    "\n",
    "    img = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                std=[0.229, 0.224, 0.225])(img) \n",
    "\n",
    "    img = transforms.Resize((new_w, new_h), antialias=True)(img)\n",
    "\n",
    "    img = img.unsqueeze(0).to(device)\n",
    "\n",
    "    return img\n",
    "\n",
    "img = preprocess_img(path_to_img)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e26a145",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def get_last_attention(model, imgs,):\n",
    "\n",
    "    new_w = img.shape[-2]\n",
    "    new_h = img.shape[-1]\n",
    "\n",
    "    w_featmap = new_w // 14\n",
    "    h_featmap = new_h // 14\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model.get_intermediate_layers(x=imgs,\n",
    "                                            reshape=True,\n",
    "                                            n = 2,\n",
    "                                            return_class_token=True,\n",
    "                                            )\n",
    "\n",
    "        # output is a list of tuples (maps, class_token), the length of the list is the number of layers considered\n",
    "\n",
    "\n",
    "        maps = output[0][0] \n",
    "\n",
    "        B, C = output[0][1].shape\n",
    "\n",
    "        # reshape maps to be (B, N, C) where N is the number of patches\n",
    "\n",
    "        maps = maps.reshape((B,maps.shape[1],-1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        class_token = output[0][1].reshape((B,-1,1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        maps = torch.cat((class_token, maps), dim=1)\n",
    "\n",
    "        # get the last attention block (only qkv)with \n",
    "    \n",
    "        qkv = model.blocks[-1].attn.qkv\n",
    "\n",
    "\n",
    "        B, N, C = maps.shape\n",
    "\n",
    "        qkv_out = qkv(maps).reshape(B, N, 3, model.num_heads, C // model.num_heads).permute(2, 0, 3, 1, 4) # (3, B, num_heads, N, C//num_heads)\n",
    "\n",
    "\n",
    "        head_dim = C // model.num_heads\n",
    "        scale = head_dim**-0.5\n",
    "\n",
    "        q, k = qkv_out[0] * scale, qkv_out[1]\n",
    "\n",
    "        attn = q @ k.transpose(-2, -1) # (B, num_heads, N, N)\n",
    "        \n",
    "        nh = model.num_heads\n",
    "        assert B == 1, \"B must be 1\"\n",
    "        attn = attn[:, :, 0, 1:].reshape(nh, w_featmap, h_featmap)\n",
    "        \n",
    "        return attn\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d168789b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "attentions = get_last_attention(model, img,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f0be8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def LocalizedRandomResizedCrop(\n",
    "                image, \n",
    "                xo, yo, Wo, Ho,\n",
    "                THR: float = 1.0,\n",
    "                scale: tuple = (0.08, 1.0),\n",
    "                ratio: tuple = (3. / 4., 4. / 3.),\n",
    "                relative_upper_scale = None, # set to 1 when cropping bbox with non square aspect ratio (very stretched)\n",
    "                ):\n",
    "    \n",
    "        \"\"\"\n",
    "        Args :\n",
    "            image (PIL Image): Image to be cropped.\n",
    "            bbox (tuple): Bounding box coordinates (x, y, w, h)\n",
    "            size (sequence or int): Desired output size of the crop. If size is an\n",
    "                int instead of sequence like (h, w), a square crop (size, size) is\n",
    "                made.\n",
    "\n",
    "                should be a multiple of patch_size\n",
    "            alpha (float): 0: centers are equal, 1 : centers are shifted\n",
    "            scale (tuple): Range of the random size of the cropped image compared to\n",
    "                the original image size.\n",
    "            ratio (tuple): Range of aspect ratio of the origin aspect ratio cropped\n",
    "                image compared to the original image ratio. Default value is\n",
    "                (3. / 4., 4. / 3.).\n",
    "\n",
    "        Returns:\n",
    "            PIL Image: Cropped image.\n",
    "    \n",
    "        \"\"\"\n",
    "\n",
    "\n",
    "        area_image = image.size[0] * image.size[1]\n",
    "\n",
    "        Ao = max(Wo, Ho)**2\n",
    "\n",
    "        if relative_upper_scale is None:\n",
    "            scale = (max(scale[0], (THR*Ao)/area_image),\n",
    "                    #min(scale[1], (1/(THR + 1e-8))*Ao/area_image))\n",
    "                    scale[1])\n",
    "        else:\n",
    "             scale = (max(scale[0], (THR*Ao)/area_image),\n",
    "                    min(scale[1], relative_upper_scale)*Ao/area_image)\n",
    "        \n",
    "        effective_scale = rd.uniform(*scale)\n",
    "        log_ratio = tuple(np.log(r) for r in ratio)\n",
    "\n",
    "        effective_ratio = np.exp(rd.uniform(*log_ratio))\n",
    "\n",
    "        side = area_image**0.5  \n",
    "\n",
    "        crop_side = side * effective_scale**0.5\n",
    "\n",
    "        if effective_ratio > 1:\n",
    "            Wc = effective_ratio * crop_side\n",
    "            Hc = crop_side\n",
    "        else:\n",
    "            Wc = crop_side\n",
    "            Hc = crop_side / effective_ratio\n",
    "\n",
    "        # prevent overflow\n",
    "        Hc = min(Hc, float(image.size[1]))\n",
    "        Wc = min(Wc, float(image.size[0]))\n",
    "\n",
    "        alpha = 1 - THR**0.5\n",
    "\n",
    "        range_x = alpha * (Wc + Wo)/2\n",
    "        range_y = alpha * (Hc + Ho)/2\n",
    "\n",
    "        effective_x = rd.uniform(-range_x, range_x)\n",
    "        effective_y = rd.uniform(-range_y, range_y)\n",
    "\n",
    "        center_x = xo + Wo/2 + effective_x\n",
    "        center_y = yo + Ho/2 + effective_y\n",
    "        crop_bbox = [\n",
    "\n",
    "            max(0, center_y - Hc/2),\n",
    "            max(0, center_x - Wc/2),\n",
    "            \n",
    "            Hc,\n",
    "            Wc,\n",
    "            \n",
    "        ]\n",
    "\n",
    "        if crop_bbox[0] + crop_bbox[2] > image.size[1]:\n",
    "            crop_bbox[2] = image.size[1] - crop_bbox[0]\n",
    "        if crop_bbox[1] + crop_bbox[3] > image.size[0]:\n",
    "            crop_bbox[3] = image.size[0] - crop_bbox[1]\n",
    "\n",
    "        crop_bbox = [abs(int(x)) for x in crop_bbox] # avoid crash with negative values\n",
    "\n",
    "        return crop_bbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06c8852",
   "metadata": {},
   "outputs": [],
   "source": [
    "def iou(gt,bbox, mode=\"std\"):\n",
    "    \"\"\"\n",
    "    gt and bbox :  x,y, w, h\n",
    "    \"\"\"\n",
    "    x1, y1, w1, h1 = gt\n",
    "    x2, y2, w2, h2 = bbox\n",
    "\n",
    "    area1 = w1*h1\n",
    "    area2 = w2*h2\n",
    "\n",
    "    xA = max(x1, x2)\n",
    "    yA = max(y1, y2)\n",
    "    xB = min(x1+w1, x2+w2)\n",
    "    yB = min(y1+h1, y2+h2)\n",
    "    \n",
    "    interArea = max(0, xB-xA) * max(0, yB-yA)\n",
    "    if mode == \"std\":\n",
    "        unionArea = area1 + area2 - interArea\n",
    "    elif mode == \"as-gt\":\n",
    "        unionArea = area1\n",
    "    elif mode == \"as-bbox\":\n",
    "        unionArea = area2\n",
    "    else:\n",
    "        raise ValueError(\"mode must be one of 'std', 'as-gt', 'as-bbox'\")\n",
    "\n",
    "    eps = 1e-6\n",
    "\n",
    "    iou = interArea / max(unionArea,eps)\n",
    "\n",
    "    return iou"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "743f5aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg.dataset.split_file = \"/mnt/data/CUB_200_2011/train_test_split.txt\"\n",
    "\n",
    "with open(cfg.dataset.split_file,\"r\") as f:\n",
    "    temp = f.readlines()\n",
    "    # format i class_number\n",
    "    split = {int(line.split()[0]) : int(line.split()[1]) for line in temp}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fbb2a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "cls_token = deepcopy(\n",
    "    model.cls_token\n",
    ")\n",
    "\n",
    "print(cls_token.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286c14e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "model.cls_token = nn.Parameter(\n",
    "    torch.load(\"/home/someone/stage_jonathan/beyond_sota_w_sam/models/Dino_train_ac87.369_2023-08-24_01:27:15.pt\")[\"cls_token\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a96b7eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i,name in enumerate(os.listdir(root)):\n",
    "    if \"humming\" in name.lower():\n",
    "        print(i,name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fff3bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_imagenet = \"/mnt/data/fred/datasets/miniimagenetimages/images/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c0c7c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.signal import convolve2d\n",
    "from matplotlib.patches import Rectangle\n",
    "from matplotlib.patches import Ellipse\n",
    "from tqdm import tqdm\n",
    "\n",
    "import random as rd\n",
    "\n",
    "import cv2\n",
    "\n",
    "L_iou = []\n",
    "\n",
    "d_box = {}\n",
    "#rd.seed(42)\n",
    "#img_idxs = rd.sample(list(boxes.keys()),1000)\n",
    "\n",
    "for img_idx in range(1):\n",
    "\n",
    "    \"\"\"folder = os.listdir(root)[164]#[rd.randint(0, len(os.listdir(root)) - 1)]\n",
    "    img_name = os.listdir(os.path.join(root, folder))[rd.randint(0, len(os.listdir(os.path.join(root, folder))) - 1)]\"\"\"\n",
    "\n",
    "\n",
    "    img_name = os.path.join(\n",
    "        root_imagenet,\n",
    "        os.listdir(root_imagenet)[rd.randint(0, len(os.listdir(root_imagenet)) - 1)]\n",
    "    )\n",
    "\n",
    "    #path_to_img = os.path.join(root, folder, img_name)\n",
    "\n",
    "    #img_idx = img_idxs[i]\n",
    "\n",
    "    path_to_img = os.path.join(root, img_name)\n",
    "\n",
    "    #path_to_img = \"/home/someone/Pictures/Screenshots/Screenshot from 2023-08-17 09-19-08.png\"\n",
    "\n",
    "    path_to_img = \"/home/someone/Pictures/Screenshots/Screenshot from 2023-08-29 09-33-59.png\"\n",
    "\n",
    "    try:\n",
    "        img = preprocess_img(path_to_img, resize=True)\n",
    "    except Exception as e:\n",
    "        print(\"error with\", path_to_img)\n",
    "        print(e)\n",
    "        continue\n",
    "\n",
    "    attentions = get_last_attention(model, img,).cpu().numpy()\n",
    "\n",
    "\n",
    "\n",
    "    # remove the n highest values\n",
    "    for attention in attentions:\n",
    "        for  _ in range(1):\n",
    "            idx = np.unravel_index(np.argmax(attention), attention.shape)\n",
    "            attention[idx] = np.mean(attention)\n",
    "\n",
    "    plt.figure(figsize=(10,6), dpi=200)\n",
    "\n",
    "    alpha = 4 # temperature\n",
    "\n",
    "    softmax = lambda x: np.exp(alpha*x)/np.sum(np.exp(alpha*x))\n",
    "\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.title(\"Original\",size=6)\n",
    "    to_display = img[0].permute(1,2,0).cpu().numpy()\n",
    "    to_display = (to_display - to_display.min())/(to_display.max() - to_display.min())\n",
    "    plt.imshow(to_display)\n",
    "    \n",
    "    plt.axis(\"off\")\n",
    "\n",
    "    #gt = boxes[img_idx]\n",
    "    \"\"\"plt.gca().add_patch(Rectangle((gt[0],gt[1]), gt[2], gt[3],\n",
    "                        edgecolor='blue',\n",
    "                        facecolor='none',\n",
    "                        lw=1))\n",
    "\n",
    "    plt.subplot(1, 3, 1)\"\"\"\n",
    "\n",
    "    conved  = convolve2d(\n",
    "                np.mean(attentions,axis=0),\n",
    "                np.ones(\n",
    "                    tuple(max(3,int(0.2*side)) for side in attentions.shape[1:])\n",
    "                ),\n",
    "                mode=\"same\",\n",
    "                boundary=\"symm\"\n",
    "            )\n",
    "    conved = cv2.resize(\n",
    "                conved,\n",
    "                (to_display.shape)[:-1:][::-1],\n",
    "                interpolation=cv2.INTER_CUBIC\n",
    "    )\n",
    "    conved = (conved - conved.min())/(conved.max() - conved.min())\n",
    "\n",
    "    plt.subplot(1, 3, 2)\n",
    "\n",
    "    plt.title(\"Softmax + box estimate\",size=6)\n",
    "\n",
    "    softmaxed = softmax(conved)\n",
    "    softmaxed = (softmaxed - softmaxed.min())/(softmaxed.max() - softmaxed.min())\n",
    "\n",
    "    plt.imshow(softmaxed,cmap=\"coolwarm\")\n",
    "\n",
    "    threshed = cv2.threshold(\n",
    "                (255*softmaxed).astype(np.uint8),\n",
    "                75, 255,\n",
    "                cv2.THRESH_BINARY #+ cv2.THRESH_OTSU\n",
    "                )[1]/255\n",
    "\n",
    "    center_y, center_x = np.mean(np.argwhere(threshed),axis=0)\n",
    "\n",
    "    h = np.sum(threshed.max(axis=1),axis=0)\n",
    "    w = np.sum(threshed.max(axis=0),axis=0)\n",
    "\n",
    "    plt.scatter(center_x,center_y,marker=\"x\",color=\"red\")\n",
    "\n",
    "    x = center_x - w/2\n",
    "    y = center_y - h/2\n",
    "\n",
    "    \"\"\"L_iou.append(iou(gt,\n",
    "                     [x,y,w,h],\n",
    "                     mode=\"std\"))\"\"\"\n",
    "    \n",
    "    d_box[img_idx] = list(map(int,[x,y,w,h]))\n",
    "    \n",
    "    #print(f\"{img_idx:4d}/{len(list(boxes.keys()))} : {float(np.mean(L_iou)):.3f}\", end=\"\\r\")\n",
    "    #print(L_iou[-1])\n",
    "\n",
    "    plt.gca().add_patch(Rectangle((x,y), w, h,\n",
    "                        edgecolor='red',\n",
    "                        facecolor='none',\n",
    "                        lw=1))\n",
    "    \n",
    "\n",
    "    \"\"\"AssertionError\"\"\"\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "\n",
    "    plt.subplot(1, 3, 1)\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "    L = []\n",
    "    for _ in range(10):\n",
    "        bbox = LocalizedRandomResizedCrop(\n",
    "                Image.fromarray(softmaxed*255).convert(\"L\"),\n",
    "                x, y, w, h,\n",
    "                THR=0.7,\n",
    "                scale=(0.08, 0.5),\n",
    "                ratio=(3. / 4., 4. / 3.),\n",
    "                relative_upper_scale = None, # set to 1 when cropping bbox with non square aspect ratio (very stretched)\n",
    "                )\n",
    "\n",
    "        \"\"\"plt.gca().add_patch(Ellipse((bbox[1]+bbox[3]/2,bbox[0]+bbox[2]/2), bbox[3], bbox[2],\n",
    "                        edgecolor='green',\n",
    "                        facecolor='none',\n",
    "                        lw=1))\"\"\"\n",
    "        \n",
    "        plt.gca().add_patch(Rectangle((bbox[1],bbox[0]), bbox[3], bbox[2],\n",
    "                        edgecolor='green',\n",
    "                        facecolor=\"none\",\n",
    "    \n",
    "                        lw=1))              \n",
    "    \n",
    "\n",
    "\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.title(\"Otsu thr\",size=6)\n",
    "\n",
    "\n",
    "    plt.imshow(threshed,cmap=\"coolwarm\")\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43035e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18978f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9905b486",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_self_attention(model,imgs):\n",
    "\n",
    "    \"\"\"\n",
    "    Only works for ViT models (especially DINO, DINOv2)\n",
    "    \n",
    "    Args:\n",
    "        model: A ViT model\n",
    "        imgs: A tensor of shape (B, 3, H, W) B = 1 \n",
    "    Returns:\n",
    "        A tensor of shape (num_heads, H_new, W_new) where H_new and W_new are the height and width of the attention map (size of the img divided by the patch size), CUDA tensor\n",
    "    \"\"\"\n",
    "\n",
    "    B, C, H, W = imgs.shape\n",
    "\n",
    "    assert len(imgs.shape) == 4, \"imgs must be a 4D tensor\"\n",
    "    assert imgs.shape[1] == 3, \"imgs must have 3 channels\"\n",
    "    assert imgs.shape[2] % model.patch_size == 0 and imgs.shape[3] % model.patch_size == 0, \"sides must be divisible by 14\"\n",
    "    assert B == 1, \"B must be 1\"\n",
    "\n",
    "    w_featmap = W // 14\n",
    "    h_featmap = H // 14\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model.get_intermediate_layers(x=imgs,\n",
    "                                            reshape=True,\n",
    "                                            n = 2,\n",
    "                                            return_class_token=True,\n",
    "                                            )\n",
    "\n",
    "        # output is a list of tuples (maps, class_token), the length of the list is the number of layers considered\n",
    "\n",
    "\n",
    "        maps = output[0][0] \n",
    "\n",
    "        B, C = output[0][1].shape\n",
    "\n",
    "        # reshape maps to be (B, N, C) where N is the number of patches\n",
    "\n",
    "        maps = maps.reshape((B,maps.shape[1],-1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        class_token = output[0][1].reshape((B,-1,1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        maps = torch.cat((class_token, maps), dim=1)\n",
    "\n",
    "        # get the last attention block (only qkv) with \n",
    "    \n",
    "        qkv = model.blocks[-1].attn.qkv\n",
    "\n",
    "\n",
    "        B, N, C = maps.shape\n",
    "\n",
    "        qkv_out = qkv(maps).reshape(B, N, 3, model.num_heads, C // model.num_heads).permute(2, 0, 3, 1, 4) # (3, B, num_heads, N, C//num_heads)\n",
    "\n",
    "\n",
    "        head_dim = C // model.num_heads\n",
    "        scale = head_dim**-0.5\n",
    "\n",
    "        q, k = qkv_out[0] * scale, qkv_out[1]\n",
    "\n",
    "        attn = q @ k.transpose(-2, -1) # (B, nh, N, N)\n",
    "        \n",
    "        nh = model.num_heads\n",
    "\n",
    "\n",
    "        #attn = attn[:, :, 0, 1:].reshape(nh, h_featmap, w_featmap)\n",
    "        \n",
    "        return attn, output[0][0] \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "034fcdc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "a, maps = get_self_attention(model, img,)\n",
    "a = a.squeeze(0).mean(axis=0)[:-1,:-1]\n",
    "for i in range(a.shape[0]):\n",
    "    a[i,i] = 0\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85330542",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = torch.mean(a,axis=1)\n",
    "\n",
    "seeds = mean.cpu().numpy().reshape(35,47)\n",
    "\n",
    "\n",
    "plt.imshow(seeds, cmap=\"coolwarm\")\n",
    "\n",
    "plt.colorbar()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25ef205",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_seeds = np.argsort(seeds.flatten())[:100]\n",
    "\n",
    "print(np.unravel_index(sorted_seeds[0], seeds.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dc75063",
   "metadata": {},
   "outputs": [],
   "source": [
    "vec = np.zeros(seeds.flatten().shape)\n",
    "\n",
    "first_seed = sorted_seeds[0]\n",
    "\n",
    "for seed in sorted_seeds[1:]:\n",
    "    idx =  seed\n",
    "\n",
    "    if maps[idx].cpu().numpy().T @ maps[first_seed].cpu().numpy() > 0:\n",
    "        vec[idx] = 1\n",
    "\n",
    "\n",
    "vec2 = np.zeros(vec.shape)\n",
    "\n",
    "for seed in vec:\n",
    "    seed = int(seed)\n",
    "    if vec[seed] == 1:\n",
    "\n",
    "        value = maps[seed].cpu().numpy().T @ np.sum([maps[i].cpu().numpy() for i in range(vec.shape[0]) if vec[i] == 1 and i != seed],axis=0) > 0\n",
    "\n",
    "        vec2[seed] = value\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4d2b480",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(vec2.reshape(35,47), cmap=\"coolwarm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c22df9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b321491",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "attn = get_last_attention(model, img,)\n",
    "\n",
    "# remove max per head\n",
    "\n",
    "for att in attn:\n",
    "    att[att == att.max()] = att.mean()\n",
    "attn = torch.mean(attn,dim=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f23a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms.functional as F\n",
    "\n",
    "conv = torch.nn.Conv2d(1,1,3, \n",
    "                       padding=\"same\", \n",
    "                       padding_mode=\"reflect\",\n",
    "                       bias=False)\n",
    "conv.weight = torch.nn.Parameter(torch.ones(3,3).unsqueeze(0).unsqueeze(0))\n",
    "conv.to(\"cuda\")\n",
    "\n",
    "\n",
    "conved = conv(attn.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)\n",
    "\n",
    "conved = F.resize(conved.unsqueeze(0), attn.shape, antialias=True).squeeze(0)\n",
    "\n",
    "print(conved.shape)\n",
    "\n",
    "softmax = torch.nn.Softmax(dim=0)\n",
    "\n",
    "\n",
    "conved = (conved - conved.min())/(conved.max() - conved.min())\n",
    "softmaxed = softmax(100*conved.flatten())\n",
    "\n",
    "softmaxed = softmaxed.reshape(conved.shape)\n",
    "\n",
    "print(softmaxed.shape)\n",
    "\n",
    "plt.imshow(softmaxed.detach().cpu().numpy(), cmap= \"coolwarm\")\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "plt.imshow(conved.detach().cpu().numpy(), cmap= \"coolwarm\")\n",
    "\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "337cec9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "density = softmaxed.flatten()\n",
    "\n",
    "cum_density = torch.cumsum(density,dim=0)\n",
    "\n",
    "cum_density = cum_density/cum_density[-1]\n",
    "\n",
    "plt.imshow(softmaxed.detach().cpu().numpy(), )\n",
    "\n",
    "for _ in range(100):\n",
    "    \n",
    "    u = torch.rand(1).to(\"cuda\")\n",
    "\n",
    "    idx = torch.searchsorted(torch.tensor(cum_density), u).item()\n",
    "\n",
    "    idx = np.unravel_index(idx, softmaxed.shape)\n",
    "\n",
    "    plt.scatter(idx[1],idx[0],color=\"red\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e83cfa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "attentions = get_last_attention(model, img,).cpu().numpy()\n",
    "\n",
    "temperature = 0.1\n",
    "\n",
    "softmax = lambda x: np.exp(x/temperature)/np.sum(np.exp(x/temperature))\n",
    "\n",
    "\n",
    "for attention in attentions:\n",
    "        for  _ in range(0):\n",
    "            idx = np.unravel_index(np.argmax(attention), attention.shape)\n",
    "            attention[idx] = np.mean(attention)\n",
    "\n",
    "conved  = convolve2d(\n",
    "                np.mean(attentions,axis=0),\n",
    "                np.ones(\n",
    "                    tuple(max(3,0) for side in attentions.shape[1:])\n",
    "                ),\n",
    "                mode=\"same\",\n",
    "                boundary=\"symm\"\n",
    "            )\n",
    "\n",
    "\n",
    "\n",
    "conved = (conved - conved.min())/(conved.max() - conved.min())\n",
    "\n",
    "softmaxed = softmax(conved)\n",
    "\n",
    "plt.imshow(\n",
    "      softmaxed,\n",
    ")\n",
    "\n",
    "distribution = np.cumsum(softmaxed.flatten())\n",
    "\n",
    "\n",
    "\n",
    "for _ in range(100):\n",
    "    u = rd.random()\n",
    "\n",
    "    idx = np.argwhere(distribution > u)[0][0]\n",
    "\n",
    "    point = np.unravel_index(idx, softmaxed.shape)\n",
    "    plt.scatter(point[1],point[0],color=\"red\", alpha=1)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fefc2d5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random as rd\n",
    "from copy import deepcopy\n",
    "def crop_from_self_att(map, \n",
    "                       #img_size,\n",
    "                       temperature,\n",
    "                       n = 1,\n",
    "                       scale = (0.08, 1.0),\n",
    "                       ratio = (3. / 4., 1.0),\n",
    "                       ):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        map : resulting distribution after convolution\n",
    "        img_size : size of the original image\n",
    "        temperature : temperature used for softmax\n",
    "        n : number of crops to generate\n",
    "        scale : scale used for cropping\n",
    "        ratio : ratio used for cropping\n",
    "    \"\"\"\n",
    "    softmax = lambda x: np.exp(x/temperature)/np.sum(np.exp(x/temperature))\n",
    "    softmaxed = softmax(map)\n",
    "\n",
    "    H, W = softmaxed.shape\n",
    "    area = min(H,W)**2\n",
    "\n",
    "    print(\"og\",softmaxed.shape)\n",
    "\n",
    "    crops = []\n",
    "\n",
    "    for _ in range(n):\n",
    "\n",
    "        eff_ratio = np.exp(rd.uniform(*np.log(ratio)))\n",
    "        eff_scale = rd.uniform(*scale)\n",
    "\n",
    "        w = max(4,int(np.sqrt(area*eff_scale*eff_ratio))-2)\n",
    "        h = max(4,int(np.sqrt(area*eff_scale/eff_ratio))-2)\n",
    "\n",
    "        # restrict the space to the center of the image\n",
    "\n",
    "        #print(\"size\",w,h)\n",
    "\n",
    "        temp = deepcopy(softmaxed[\n",
    "            w//2:2-w//2,\n",
    "            h//2:2-h//2,\n",
    "        ])\n",
    "\n",
    "        # normalize\n",
    "        temp/=np.sum(temp)\n",
    "\n",
    "        if any(size<=0 for size in temp.shape):\n",
    "            print(\"error\",h,w,eff_ratio,eff_scale,temp.shape)\n",
    "        # get the cumulative distribution\n",
    "        distribution = np.cumsum(temp.flatten())\n",
    "\n",
    "        #print(distribution[-1])\n",
    "\n",
    "        # get a random number\n",
    "        u = rd.random()\n",
    "        idx = np.argwhere(distribution > u)[0][0]\n",
    "\n",
    "        point = np.unravel_index(idx, temp.shape)\n",
    "\n",
    "        # get the coordinates of the crop in the original image\n",
    "        x = point[0]\n",
    "        y = point[1] \n",
    "\n",
    "        \n",
    "        crops.append(\n",
    "            [x,y,w,h]\n",
    "        )\n",
    "\n",
    "        # get the coordinates of the crop in the original image\n",
    "        \"\"\"x_og = x * img_size[0] / W\n",
    "        y_og = y * img_size[1] / H\n",
    "\n",
    "        w_og = w * img_size[0] / W\n",
    "        h_og = h * img_size[1] / H\n",
    "\n",
    "\n",
    "\n",
    "        crops.append(\n",
    "            [x_og,y_og,w_og,h_og]\n",
    "        )\"\"\"\n",
    "\n",
    "    plt.imshow(softmaxed)\n",
    "    for x,y,w,h in crops:\n",
    "\n",
    "\n",
    "        \"\"\"plt.gca().add_patch(\n",
    "            Rectangle((x,y),w,h,\n",
    "                    edgecolor='green',\n",
    "                    facecolor='none',\n",
    "                    lw=1)\n",
    "                      )\n",
    "    \"\"\"\n",
    "        plt.scatter(\n",
    "            y+h//2,\n",
    "            x+w//2,\n",
    "            \n",
    "        )\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "crop_from_self_att(map=conved,temperature=0.1,n=100)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edef9385",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.transforms import functional as F\n",
    "\n",
    "from typing import Optional, Union, Sequence, List, Tuple\n",
    "import warnings\n",
    "import math\n",
    "\n",
    "\n",
    "class DiNORandomResizedCrop(torch.nn.Module):\n",
    "      \n",
    "    \"\"\"\n",
    "    Similar to the random resized crop of torchvision, but the distribution \n",
    "    of the crop's center is DiNOv2's self attention distribution.\n",
    "\n",
    "    Should be placed after the Normalize and ToTensor transforms, so that the image is suitable for DiNOv2 (there are no preprocessing transforms)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "    self,\n",
    "    size,\n",
    "    dino_type: str,\n",
    "    scale=(0.08, 1.0),\n",
    "    ratio=(3.0 / 4.0, 4.0 / 3.0),\n",
    "    temperature: float = 10,\n",
    "    kernel_size: int = 3,\n",
    "    interpolation=F.InterpolationMode.BILINEAR,\n",
    "    antialias: Optional[Union[str, bool]] = \"warn\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.size = size\n",
    "\n",
    "        if not isinstance(scale, Sequence):\n",
    "            raise TypeError(\"Scale should be a sequence\")\n",
    "        if not isinstance(ratio, Sequence):\n",
    "            raise TypeError(\"Ratio should be a sequence\")\n",
    "        if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):\n",
    "            warnings.warn(\"Scale and ratio should be of kind (min, max)\")\n",
    "\n",
    "        self.interpolation = interpolation\n",
    "        self.antialias = antialias\n",
    "        self.scale = scale\n",
    "        self.ratio = ratio\n",
    "\n",
    "\n",
    "        ############################## Load DiNOv2 model ##############################################\n",
    "        dino_types = ['dinov2_vits14','dinov2_vitb14','dinov2_vitl14','dinov2_vitg14']\n",
    "        assert dino_type in dino_types, f\"dino_type should be one of {dino_types}, got {dino_type}\"\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        \n",
    "        print(f\"Loading DiNOv2 model {dino_type} on device {self.device}\")\n",
    "        self.model = torch.hub.load('facebookresearch/dinov2', dino_type, pretrained=True).to(self.device)\n",
    "\n",
    "        for param in self.model.parameters():\n",
    "            param.requires_grad = False\n",
    "        \n",
    "        self.model.eval()\n",
    "\n",
    "        self.density_generator = DensityGenerator(temperature=temperature,\n",
    "                                                  kernel_size=kernel_size,\n",
    "                                                  ).to(self.device)\n",
    "\n",
    "        ##############################################################################################\n",
    "\n",
    "    def get_map(self, img: torch.Tensor) -> torch.Tensor:\n",
    "        _, height, width = F.get_dimensions(img)\n",
    "        img = F.resize(img, (height - (height % 14), width - (width % 14)), interpolation=F.InterpolationMode.BICUBIC, antialias=self.antialias)\n",
    "        _, height, width = F.get_dimensions(img)\n",
    "\n",
    "        img = img.to(torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"))\n",
    "        attention_map = get_self_attention(self.model, img) # (num_heads, H//14, W//14)\n",
    "        density_map = self.density_generator(attention_map).cpu()\n",
    "        density_map = F.resize(density_map.unsqueeze(0), (height, width), interpolation=F.InterpolationMode.BICUBIC, antialias=self.antialias).squeeze() # (H, W)\n",
    "\n",
    "        density_map = density_map / torch.sum(density_map)\n",
    "\n",
    "        return density_map\n",
    "\n",
    "\n",
    "    def get_params(self, img: torch.Tensor, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int]:\n",
    "        \"\"\"Get parameters for ``crop`` for a random sized crop.\n",
    "\n",
    "        Args:\n",
    "            img (PIL Image or Tensor): Input image.\n",
    "            scale (list): range of scale of the origin size cropped\n",
    "            ratio (list): range of aspect ratio of the origin aspect ratio cropped\n",
    "\n",
    "        Returns:\n",
    "            tuple: params (i, j, h, w) to be passed to ``crop`` for a random\n",
    "            sized crop.\n",
    "        \"\"\"\n",
    "\n",
    "        ########  Resize for DiNOv2 (so that the image size is divisible by 14)  ######################\n",
    "\n",
    "\n",
    "        _, height, width = F.get_dimensions(img)\n",
    "\n",
    "        area = height * width\n",
    "\n",
    "        # Get distribution of the attention map\n",
    "\n",
    "        img = img.to(torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"))\n",
    "        attention_map = get_self_attention(self.model, img) # (num_heads, H//14, W//14)\n",
    "\n",
    "\n",
    "        # Generate density map\n",
    "\n",
    "        density_map = self.density_generator(attention_map) # (H//14, W//14)\n",
    "\n",
    "        density_map = F.resize(density_map.unsqueeze(0), (height, width), interpolation=F.InterpolationMode.BICUBIC, antialias=self.antialias).squeeze() # (H, W)\n",
    "\n",
    "        log_ratio = torch.log(torch.tensor(ratio))\n",
    "        for _ in range(10):\n",
    "            target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()\n",
    "            aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()\n",
    "\n",
    "            w = int(round(math.sqrt(target_area * aspect_ratio)))\n",
    "            h = int(round(math.sqrt(target_area / aspect_ratio)))\n",
    "\n",
    "\n",
    "            if 0 < w <= width and 0 < h <= height:\n",
    "\n",
    "                \"\"\"i = torch.randint(0, height - h + 1, size=(1,)).item()\n",
    "                j = torch.randint(0, width - w + 1, size=(1,)).item()\"\"\"\n",
    "\n",
    "                x,y = generate_point(density_map, (h,w))\n",
    "\n",
    "                i = y - h // 2\n",
    "                j = x - w // 2\n",
    "            \n",
    "                return i, j, h, w\n",
    "\n",
    "\n",
    "        # Fallback to central crop\n",
    "        in_ratio = float(width) / float(height)\n",
    "        if in_ratio < min(ratio):\n",
    "            w = width\n",
    "            h = int(round(w / min(ratio)))\n",
    "        elif in_ratio > max(ratio):\n",
    "            h = height\n",
    "            w = int(round(h * max(ratio)))\n",
    "        else:  # whole image\n",
    "            w = width\n",
    "            h = height\n",
    "        i = (height - h) // 2\n",
    "        j = (width - w) // 2\n",
    "        return i, j, h, w\n",
    "\n",
    "    def forward(self, img):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            img (PIL Image or Tensor): Image to be cropped and resized.\n",
    "\n",
    "        Returns:\n",
    "            PIL Image or Tensor: Randomly cropped and resized image.\n",
    "        \"\"\"\n",
    "        i, j, h, w = self.get_params(img, self.scale, self.ratio)\n",
    "        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation, antialias=self.antialias)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        interpolate_str = self.interpolation.value\n",
    "        format_string = self.__class__.__name__ + f\"(size={self.size}\"\n",
    "        format_string += f\", scale={tuple(round(s, 4) for s in self.scale)}\"\n",
    "        format_string += f\", ratio={tuple(round(r, 4) for r in self.ratio)}\"\n",
    "        format_string += f\", interpolation={interpolate_str}\"\n",
    "        format_string += f\", antialias={self.antialias})\"\n",
    "        return format_string\n",
    "\n",
    "\n",
    "def get_self_attention(model,imgs):\n",
    "\n",
    "    \"\"\"\n",
    "    Extract the self attention map from a ViT model (for now, only spports DINOv2)\n",
    "    \n",
    "    Args:\n",
    "        model: A ViT model\n",
    "        imgs: A tensor of shape (B, 3, H, W) B = 1 \n",
    "    Returns:\n",
    "        A tensor of shape (num_heads, H_new, W_new) where H_new and W_new are the height and width of the attention map (size of the img divided by the patch size), CUDA tensor\n",
    "    \"\"\"\n",
    "\n",
    "    if len(imgs.shape) == 3:\n",
    "        imgs = imgs.unsqueeze(0)\n",
    "    elif len(imgs.shape) == 4:\n",
    "        pass\n",
    "    else:\n",
    "        raise ValueError(\"imgs must be a 3D or 4D tensor\")\n",
    "\n",
    "    B, C, H, W = imgs.shape\n",
    "\n",
    "    imgs = F.resize(imgs, (H - (H % 14), W - (W % 14)), interpolation=F.InterpolationMode.BICUBIC, antialias=True)\n",
    "    \n",
    "    assert imgs.shape[1] == 3, \"imgs must have 3 channels\"\n",
    "    #assert imgs.shape[2] % model.patch_size == 0 and imgs.shape[3] % model.patch_size == 0, \"sides must be divisible by 14\"\n",
    "    assert B == 1, \"B must be 1\"\n",
    "\n",
    "\n",
    "\n",
    "    w_featmap = W // 14\n",
    "    h_featmap = H // 14\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model.get_intermediate_layers(x=imgs,\n",
    "                                            reshape=True,\n",
    "                                            n = 2,\n",
    "                                            return_class_token=True,\n",
    "                                            )\n",
    "\n",
    "        # output is a list of tuples (maps, class_token), the length of the list is the number of layers considered\n",
    "\n",
    "\n",
    "        maps = output[0][0] \n",
    "\n",
    "        B, C = output[0][1].shape\n",
    "\n",
    "        # reshape maps to be (B, N, C) where N is the number of patches\n",
    "\n",
    "        maps = maps.reshape((B,maps.shape[1],-1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        class_token = output[0][1].reshape((B,-1,1)).permute(0,2,1)\n",
    "\n",
    "\n",
    "        maps = torch.cat((class_token, maps), dim=1)\n",
    "\n",
    "        # get the last attention block (only qkv) with \n",
    "    \n",
    "        qkv = model.blocks[-1].attn.qkv\n",
    "\n",
    "\n",
    "        B, N, C = maps.shape\n",
    "\n",
    "        qkv_out = qkv(maps).reshape(B, N, 3, model.num_heads, C // model.num_heads).permute(2, 0, 3, 1, 4) # (3, B, num_heads, N, C//num_heads)\n",
    "\n",
    "\n",
    "        head_dim = C // model.num_heads\n",
    "        scale = head_dim**-0.5\n",
    "\n",
    "        q, k = qkv_out[0] * scale, qkv_out[1]\n",
    "\n",
    "        attn = q @ k.transpose(-2, -1) # (B, nh, N, N)\n",
    "        \n",
    "        nh = model.num_heads\n",
    "        attn = attn[:, :, 0, 1:].reshape(nh, h_featmap, w_featmap)\n",
    "        \n",
    "        return attn\n",
    "\n",
    "\n",
    "class DensityGenerator(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Postprocess the attention maps to generate a density map : \n",
    "    - Average over the heads\n",
    "    - Outlier removal (remove the max and replace it by the mean)\n",
    "    - Apply a gaussian filter\n",
    "    - Aplly a min max normalization\n",
    "    - Apply a softmax\n",
    "\n",
    "\n",
    "    Args:\n",
    "        kernel_size: The size of the gaussian kernel (Note that all images should be the same size, and are square)\n",
    "        temperature: The temperature used to apply the softmax\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, kernel_size: int = 3, temperature: float = 0.1):\n",
    "        super().__init__()\n",
    "        self.kernel_size = kernel_size\n",
    "        self.temperature = temperature\n",
    "\n",
    "        self.conv = torch.nn.Conv2d(in_channels=1, out_channels=1, kernel_size=kernel_size, padding=\"same\", bias=False) # padding to keep the same size\n",
    "        self.conv.weight.data = torch.ones(1, 1, kernel_size, kernel_size)\n",
    "\n",
    "        self.softmax = torch.nn.Softmax(dim=0)\n",
    "\n",
    "    def forward(self, maps: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            maps: A tensor of shape (num_heads, H, W)\n",
    "        Returns:\n",
    "            A tensor of shape (H, W)\n",
    "        \"\"\"\n",
    "\n",
    "        maps = torch.mean(maps, dim=0) # average over the heads # (H, W)\n",
    "\n",
    "        maps = self.conv(maps.unsqueeze(0)) # gaussian filter # (1, H, W)\n",
    "\n",
    "        # positive values\n",
    "\n",
    "        maps = (maps - torch.min(maps))\n",
    "        \n",
    "        temp = maps.clone()\n",
    "\n",
    "        temp = self.softmax(temp.flatten()/self.temperature)\n",
    "\n",
    "        maps = temp.reshape(maps.shape)\n",
    "\n",
    "        return maps.squeeze(0) # (H, W)\n",
    "\n",
    "\n",
    "def generate_point(density:torch.Tensor, \n",
    "                   target_size: tuple\n",
    "                   ) -> tuple [int, int]:\n",
    "\n",
    "    \"\"\"\n",
    "    Generate a point from a density map\n",
    "\n",
    "    Args:\n",
    "        density: A tensor of shape (H, W)\n",
    "        og_size: The size of the original image (H, W)\n",
    "    Returns:\n",
    "        A tensor of shape (2,) containing the coordinates of the point\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    H, W = density.shape\n",
    "\n",
    "    h,w = target_size\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "    density = density.flatten()\n",
    "\n",
    "    density = density / torch.sum(density)\n",
    "\n",
    "    cum_density = torch.cumsum(density, dim=0)\n",
    "\n",
    "\n",
    "    u = torch.rand(1).to(\"cuda\")\n",
    "    # the cumulative density is sorted in ascending order, so we find the nearest value \n",
    "    i = torch.searchsorted(cum_density, u).cpu()\n",
    "    # we get the coordinates of the point\n",
    "    x = i % W\n",
    "    y = i // W\n",
    "\n",
    "\n",
    "    # the crop must be within the image\n",
    "    for _ in range(10):\n",
    "\n",
    "        if (y - h // 2 < 0 or y + h // 2 >= H or x - w // 2 < 0 or x + w // 2 >= W):\n",
    "            u = torch.rand(1).to(\"cuda\")\n",
    "            i = torch.searchsorted(cum_density, u).cpu()\n",
    "            x = i % W\n",
    "            y = i // W\n",
    "        else : return x,y # int, int\n",
    "\n",
    "    # if we can't find a point, we return the center of the image (fallback to center crop)\n",
    "    return W//2, H//2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "463f7f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = DiNORandomResizedCrop(224, \"dinov2_vits14\",\n",
    "                           #scale=(0.1,0.1),\n",
    "                           temperature=10,\n",
    "                           kernel_size=5,\n",
    "                           antialias=True,\n",
    "                           )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05fbe669",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "\n",
    "device = \"cuda\"\n",
    "def preprocess_img(path_to_img, resize=False):\n",
    "    \"\"\"img = cv2.imread(path_to_img)\n",
    "\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "    img = Image.fromarray(img)\"\"\"\n",
    "\n",
    "    img = Image.open(path_to_img)\n",
    "    if img.mode != \"RGB\":\n",
    "        img = img.convert(\"RGB\")\n",
    "    \n",
    "    if resize:img = transforms.Resize(500, antialias=True)(img)\n",
    "    img = transforms.RandomAutocontrast(p=1)(img)\n",
    "\n",
    "    h, w = img.size\n",
    "    new_w = w - (w % 14)\n",
    "    new_h = h - (h % 14)\n",
    "\n",
    "\n",
    "    img = transforms.ToTensor()(img)\n",
    "\n",
    "    img = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                std=[0.229, 0.224, 0.225])(img) \n",
    "\n",
    "    img = transforms.Resize((new_w, new_h), antialias=True)(img)\n",
    "\n",
    "    img = img.unsqueeze(0).to(device)\n",
    "\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2fb514",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random as rd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "root_imagenet = '/mnt/data/fred/datasets/miniimagenetimages/images/'\n",
    "\n",
    "img_idx_r = rd.randint(0, len(os.listdir(root_imagenet)) - 1)\n",
    "\n",
    "#img_idx_r = 2468 #53367 # 32337 #\n",
    "\n",
    "img_name = os.path.join(\n",
    "        root_imagenet,\n",
    "        os.listdir(root_imagenet)[img_idx_r]\n",
    "    )\n",
    "\n",
    "#img_name = \"/home/someone/Pictures/Screenshots/Screenshot from 2023-08-29 09-33-59.png\"\n",
    "print(img_idx_r)\n",
    "plt.imshow(Image.open(img_name).convert(\"RGB\"))\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n",
    "img = preprocess_img(img_name, )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55e77a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "plt.imshow(Image.open(img_name).convert(\"RGB\"))\n",
    "to_display = tr.get_map(img).detach().cpu().numpy()\n",
    "plt.imshow(to_display,alpha=0.8)\n",
    "\n",
    "\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "273b44f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed = tr(img)\n",
    "\n",
    "print(transformed.shape)\n",
    "\n",
    "# norm\n",
    "transformed = (transformed - img.min())/(img.max() - img.min())\n",
    "\n",
    "transformed = transformed.squeeze().permute(1,2,0).cpu().numpy()\n",
    "\n",
    "plt.imshow(transformed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbbd655b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(\n",
    "    ((transforms.RandomResizedCrop(224, \n",
    "                                   #scale=(0.1,0.1),\n",
    "                                   )(img).squeeze() - img.min())/(img.max() - img.min())).permute(1,2,0).cpu().numpy()\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5cfb91d",
   "metadata": {},
   "source": [
    "# Test time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1180b106",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import torch\n",
    "\n",
    "\n",
    "train = torchvision.datasets.ImageFolder(\n",
    "    root=\"/mnt/data/CUB_200_2011/images/\",\n",
    "    transform= transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.RandomAutocontrast(p=1),\n",
    "        tr\n",
    "        ])\n",
    ")\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    train,\n",
    "    batch_size=2,)\n",
    "\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1339e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from dino_transform import DiNORandomResizedCrop\n",
    "\n",
    "tr = DiNORandomResizedCrop(size=(224,224), \n",
    "                           dino_type=\"dinov2_vits14\",\n",
    "                           #scale=(0.1,0.1),\n",
    "                           temperature=10,\n",
    "                           kernel_size=5,\n",
    "                           antialias=True,\n",
    "                           )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3e5f0e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## bs 8\n",
    "for img, _ in tqdm(loader,total=len(loader)):\n",
    "    pass\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cb1eaf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## bs 32\n",
    "for img, _ in tqdm(loader,total=len(loader)):\n",
    "    pass\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fba4639a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## bs 32 search on gpu\n",
    "for img, _ in tqdm(loader,total=len(loader)):\n",
    "    pass\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "337bd87e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "for img, _ in tqdm(loader,total=len(loader)):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3aa438f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from config import cfg\n",
    "from dataset import load_cub_datasets\n",
    "cfg.use_dino = True\n",
    "train_loader,val_loader = load_cub_datasets(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288ee129",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision import transforms\n",
    "for imgs, _ in train_loader:\n",
    "    print(imgs.shape)\n",
    "    for img in imgs:\n",
    "        #img = img.permute(1,2,0)\n",
    "        img = (img - img.min())/(img.max() - img.min())\n",
    "        img = transforms.ToPILImage()(img)\n",
    "        plt.imshow(img)\n",
    "        plt.show()\n",
    "    break"
   ]
  }
 ],
 "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.10.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}