{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f57d4d5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import torch\n",
    "import cv2\n",
    "import numpy as np\n",
    "import Models as M\n",
    "import torchvision\n",
    "from torchvision import transforms, datasets\n",
    "import van\n",
    "import defs\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.metrics import silhouette_score, davies_bouldin_score, calinski_harabasz_score\n",
    "from sklearn.decomposition import PCA\n",
    "import numpy as np\n",
    "from Loaders import RARP_ChannelSwap\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class UnlabeledImageDataset(Dataset):\n",
    "    def __init__(self, root_dir, transform=None, exts=(\".png\",\".jpg\",\".jpeg\",\".bmp\", \".tiff\")):\n",
    "        self.paths = [p for p in Path(root_dir).rglob(\"*\")\n",
    "                      if p.suffix.lower() in exts]\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img = defs.load_Img_RBG_tensor_norm(self.paths[idx])\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "        return img\n",
    "\n",
    "Mean, Std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "\n",
    "def LoadModel(path_model:str, Device, b:int = 1):\n",
    "    model = van.van_b1(pretrained=False, num_classes = 0) if b == 1 else van.van_b2(pretrained=False, num_classes = 0)\n",
    "    model.load_state_dict(torch.load(path_model))\n",
    "    model.to(Device)\n",
    "    model.eval()\n",
    "    \n",
    "    return model\n",
    "\n",
    "def Denorlalize (img:torch.Tensor, std, mean):\n",
    "    ImgNumpy = img.numpy().transpose((1, 2, 0))\n",
    "    ImgNumpy = np.clip((std * ImgNumpy + mean), 0, 1)\n",
    "    \n",
    "    return ImgNumpy\n",
    "\n",
    "Transforms = transforms.Compose([\n",
    "    transforms.Resize(360, interpolation=transforms.InterpolationMode.BICUBIC),\n",
    "    transforms.CenterCrop(300),\n",
    "    transforms.Normalize(Mean, Std)\n",
    "])\n",
    "\n",
    "ds = UnlabeledImageDataset(\"./DataSet_Ando_All_no20/NVB/\", transform=Transforms)\n",
    "#ds = torchvision.datasets.DatasetFolder(\n",
    "#        \"./DataSet_Ando_All_no20/\",\n",
    "#        loader=defs.load_Img_RBG_tensor_norm,\n",
    "#        extensions=\"tiff\",\n",
    "#        transform=transforms\n",
    "#)\n",
    "\n",
    "dl = DataLoader(ds, batch_size=64, shuffle=False)\n",
    "\n",
    "torch.set_float32_matmul_precision('high')\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "device = torch.device(\"cuda:1\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model_s, model_t = LoadModel(\"./van_b1_student.pth\", device), LoadModel(\"./van_b1_teacher.pth\", device)\n",
    "model_MAE = LoadModel(\"./van_b1_MAE_MSE.pth\", device)\n",
    "\n",
    "model_IN = van.van_b1(pretrained=True, num_classes = 0)\n",
    "model_IN.to(device)\n",
    "model_IN.eval()\n",
    "\n",
    "model_Hybrid = [LoadModel(pth_model, device, 2) for pth_model in [\"./Hybrid T-S Model_S.pth\", \"./Hybrid T-S Model_T.pth\", ]] #\"./van_b2_student.pth\", \"./van_b2_teacher.pth\"\n",
    "models = [model_s, model_t, model_MAE, model_IN] + model_Hybrid\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a46a2e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_label_free_clustering(embs: np.ndarray,\n",
    "                                   n_clusters: int = 10,\n",
    "                                   random_state: int = 42):\n",
    "    \"\"\"\n",
    "    Fits k-means and returns the silhouette score.\n",
    "    \"\"\"\n",
    "    km = KMeans(n_clusters=n_clusters,\n",
    "                random_state=random_state,\n",
    "                n_init=10,\n",
    "                verbose=0)\n",
    "    labels = km.fit_predict(embs)             # (N,)\n",
    "    score  = silhouette_score(embs, labels)   # float in [-1, 1]\n",
    "    return score, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "01cafeba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_embeddings(model, dataloader, device):\n",
    "    embs = []\n",
    "    with torch.no_grad():\n",
    "        for imgs in dataloader:              # imgs is [B,3,224,224]\n",
    "            if isinstance(imgs, list):\n",
    "                imgs = imgs[0]\n",
    "            imgs = imgs.to(device)\n",
    "            z = model(imgs)           # [B, 512]\n",
    "            z = torch.nn.functional.normalize(z, dim=1)\n",
    "            embs.append(z.cpu())\n",
    "    return torch.cat(embs, dim=0).numpy()    # (N, 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "45bde1f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = defs.load_Img_RBG_tensor_norm(\"./DataSet_Ando_All_no20/NVB/429.tiff\")\n",
    "img = Transforms(img)\n",
    "img = img.repeat(1, 1, 1, 1).to(device)\n",
    "\n",
    "ref_img_Outside = defs.load_Img_RBG_tensor_norm(\"./Laparoscopic-Abdominal-Hysterectomy.jpg\")\n",
    "ref_img_Outside = Transforms(ref_img_Outside)\n",
    "ref_img_Outside = ref_img_Outside.repeat(1, 1, 1, 1).to(device)\n",
    "\n",
    "\n",
    "ref_img_inside = defs.load_Img_RBG_tensor_norm(\"./IMG-66451.jpg\")\n",
    "ref_img_inside = Transforms(ref_img_inside)\n",
    "ref_img_inside = ref_img_inside.repeat(1, 1, 1, 1).to(device)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d731eb14",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_byModel = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for m in models:\n",
    "        ref_byModel.append(m(ref_img_Outside))\n",
    "        \n",
    "ref_in_byModel = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for m in models:\n",
    "        ref_in_byModel.append(m(ref_img_inside))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "95e7ced0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_grad_cam import EigenGradCAM as G_CAM\n",
    "from pytorch_grad_cam import EigenCAM as F_CAM\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fdd2e326",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimilarityToConceptTarget:\n",
    "    def __init__(self, features):\n",
    "        self.features = features\n",
    "    \n",
    "    def __call__(self, model_output):\n",
    "        cosSim = torch.nn.CosineSimilarity(dim=0)\n",
    "        return cosSim(model_output, self.features)\n",
    "    \n",
    "class DifferenceFromConceptTarget:\n",
    "    def __init__(self, features):\n",
    "        self.features = features\n",
    "    \n",
    "    def __call__(self, model_output):\n",
    "        cos = torch.nn.CosineSimilarity(dim=0)\n",
    "        return 1 - cos(model_output, self.features)\n",
    "\n",
    "def CAM_by_model(m, input_img, std, mean, Similarity:bool = True, FeaturesToCompare:torch.Tensor = None, class_agnostic:bool = False):\n",
    "    target = [m.block4[-1]]\n",
    "    cam = G_CAM(model=m, target_layers=target) if not class_agnostic else F_CAM(model=m, target_layers=target)\n",
    "    \n",
    "    tar = None\n",
    "    if FeaturesToCompare is not None:\n",
    "        tar = [SimilarityToConceptTarget(FeaturesToCompare)] if Similarity else [DifferenceFromConceptTarget(FeaturesToCompare)]\n",
    "        \n",
    "    gi = cam(img, targets=tar)\n",
    "    \n",
    "    rgb_img = Denorlalize(input_img.cpu()[0], std, mean)\n",
    "    \n",
    "    visualization = show_cam_on_image(rgb_img, gi[0], use_rgb=True, colormap=cv2.COLORMAP_MAGMA if not class_agnostic else cv2.COLORMAP_JET)\n",
    "    \n",
    "    return visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3763f510",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5916072 characters]",
      "text/plain": [
       "<Figure size 1500x3000 with 18 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dis_info = {\n",
    "  \"cams_list\":[],\n",
    "  \"models_names\":[\"Student\", \"Teacher\", \"MAE\", \"ImageNet\", \"Hybrid T-S (Student)\",  \"Hybrid T-S (Teacher)\", \"Student_b2\", \"Teacher_b2\",]\n",
    "}\n",
    "\n",
    "for i, m in enumerate(models):\n",
    "    dis_info[\"cams_list\"].append(CAM_by_model(m, img, Std, Mean, FeaturesToCompare=ref_in_byModel[i][0, :]))\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(len(models), 3, figsize=(15, 30))\n",
    "\n",
    "for row, cam in enumerate(dis_info[\"cams_list\"]):\n",
    "    ax[row][0].imshow(Denorlalize(img.cpu()[0], Std, Mean))\n",
    "    ax[row][0].set_title(f\"Original Image\")\n",
    "    ax[row][0].axis(\"off\")\n",
    "    \n",
    "    ax[row][1].imshow(Denorlalize(ref_img_inside.cpu()[0], Std, Mean))\n",
    "    ax[row][1].set_title(f\"Sample Image\")\n",
    "    ax[row][1].axis(\"off\")\n",
    "\n",
    "    ax[row][2].imshow(cam)\n",
    "    ax[row][2].set_title(f\"CAM {dis_info['models_names'][row]}\")\n",
    "    ax[row][2].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a2e8bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dis_info = {\n",
    "  \"cams_list\":[],\n",
    "  \"models_names\":[\"Student\", \"Teacher\", \"MAE\", \"ImageNet\", \"Hybrid T-S (Student)\",  \"Hybrid T-S (Teacher)\", \"Student_b2\", \"Teacher_b2\",]\n",
    "}\n",
    "\n",
    "for i, m in enumerate(models):\n",
    "    dis_info[\"cams_list\"].append(CAM_by_model(m, img, Std, Mean, FeaturesToCompare=ref_byModel[i][0, :]))\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(len(models), 3, figsize=(15, 30))\n",
    "\n",
    "for row, cam in enumerate(dis_info[\"cams_list\"]):\n",
    "    ax[row][0].imshow(Denorlalize(img.cpu()[0], Std, Mean))\n",
    "    ax[row][0].set_title(f\"Original Image\")\n",
    "    ax[row][0].axis(\"off\")\n",
    "    \n",
    "    ax[row][1].imshow(Denorlalize(ref_img_Outside.cpu()[0], Std, Mean))\n",
    "    ax[row][1].set_title(f\"Sample Image\")\n",
    "    ax[row][1].axis(\"off\")\n",
    "\n",
    "    ax[row][2].imshow(cam)\n",
    "    ax[row][2].set_title(f\"CAM {dis_info['models_names'][row]}\")\n",
    "    ax[row][2].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "684664db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4692644 characters]",
      "text/plain": [
       "<Figure size 1500x3000 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dis_info = {\n",
    "  \"cams_list\":[],\n",
    "  \"models_names\":[\"Student\", \"Teacher\", \"MAE\", \"ImageNet\", \"Hybrid T-S (Student)\",  \"Hybrid T-S (Teacher)\", \"Student_b2\", \"Teacher_b2\",]\n",
    "}\n",
    "\n",
    "for i, m in enumerate(models):\n",
    "    dis_info[\"cams_list\"].append(CAM_by_model(m, img, Std, Mean, class_agnostic=True))\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(len(models), 2, figsize=(15, 30))\n",
    "\n",
    "for row, cam in enumerate(dis_info[\"cams_list\"]):\n",
    "    ax[row][0].imshow(Denorlalize(img.cpu()[0], Std, Mean))\n",
    "    ax[row][0].set_title(f\"Original Image\")\n",
    "    ax[row][0].axis(\"off\")\n",
    "  \n",
    "    ax[row][1].imshow(cam)\n",
    "    ax[row][1].set_title(f\"CAM {dis_info['models_names'][row]}\")\n",
    "    ax[row][1].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2bba0ff4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model 0 VAN\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k=10 → silhouette score = 0.0614\n",
      "model 1 VAN\n",
      "k=10 → silhouette score = 0.0434\n",
      "model 2 VAN\n",
      "k=10 → silhouette score = 0.1061\n",
      "model 3 VAN\n",
      "k=10 → silhouette score = 0.0296\n",
      "model 4 VAN\n",
      "k=10 → silhouette score = 0.0337\n",
      "model 5 VAN\n",
      "k=10 → silhouette score = 0.0011\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 151184 characters]",
      "text/plain": [
       "<Figure size 1500x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "k = 10\n",
    "\n",
    "fig, ax = plt.subplots(1, len(models), figsize=(15, 8))\n",
    "\n",
    "for i, m in enumerate(models):\n",
    "    print (f\"model {i} {type(m).__name__}\")\n",
    "    embeddings = extract_embeddings(m, dl, device)\n",
    "    \n",
    "    sil_score, l = evaluate_label_free_clustering(embeddings, n_clusters=k, random_state=0)\n",
    "    print(f\"k={k:2d} → silhouette score = {sil_score:.4f}\")\n",
    "    \n",
    "    pca_kmean = PCA(2, random_state=0)\n",
    "    emb_2d = pca_kmean.fit_transform(embeddings)\n",
    "    \n",
    "    ax[i].scatter(emb_2d[:,0], emb_2d[:,1], c=l, s=5, alpha=0.6)\n",
    "    \n",
    "    #for k_ in range(2, 20):\n",
    "    #    km = KMeans(n_clusters=k_, random_state=0).fit(embeddings)\n",
    "    #    labels = km.labels_\n",
    "    #    print(f\"k={k_:2d}  silh={silhouette_score(embeddings,labels):.3f}\"\n",
    "    #        f\"  db={davies_bouldin_score(embeddings,labels):.3f}\"\n",
    "    #        f\"  ch={calinski_harabasz_score(embeddings,labels):.0f}\")\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8c384b5",
   "metadata": {},
   "source": [
    "def setup_seed(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e37d7ba7",
   "metadata": {},
   "source": [
    "transforms_rand = transforms.Compose([\n",
    "    transforms.Resize(360, interpolation=transforms.InterpolationMode.BICUBIC),\n",
    "    transforms.CenterCrop(300),\n",
    "    transforms.RandomPerspective(0.7, p=1),\n",
    "    transforms.ElasticTransform(150.0),\n",
    "    transforms.RandomApply([\n",
    "        RARP_ChannelSwap()\n",
    "    ], p=1),\n",
    "    transforms.Normalize(Mean, Std)\n",
    "])\n",
    "\n",
    "img_1 = defs.load_Img_RBG_tensor_norm(\"./DataSet_Ando_All_no20/NVB/350.tiff\")\n",
    "\n",
    "img_1 = transforms_rand(img_1)\n",
    "img_1 = img_1.repeat(1, 1, 1, 1)\n",
    "img_1 = img_1.to(device)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(15, 8))\n",
    "\n",
    "ax[0].imshow(Denorlalize(img_1.cpu()[0], Std, Mean))\n",
    "ax[0].set_title(f\"Radn. Aug. Image\")\n",
    "ax[0].axis(\"off\")\n",
    "\n",
    "ax[1].imshow(Denorlalize(img.cpu()[0], Std, Mean))\n",
    "ax[1].set_title(f\"Original Image \")\n",
    "ax[1].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cb851f",
   "metadata": {},
   "source": [
    "with torch.no_grad():\n",
    "    for i, m in enumerate(models):\n",
    "        print(f\"Model {i} {type(m).__name__}\")\n",
    "        z_1, z = m(img_1), m(img)\n",
    "        \n",
    "        sim = torch.nn.functional.cosine_similarity(z_1, z, dim=1)\n",
    "        print (f\"Sim: {sim.item():.4f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
