{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lr=0.0001, L1=None\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import torch\n",
    "import cv2\n",
    "import numpy as np\n",
    "import Models as M\n",
    "\n",
    "import torchvision.transforms as T\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "def _removeBlackBorder(image):\n",
    "    image = np.array(image)\n",
    "    \n",
    "    copyImg = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)\n",
    "    h = copyImg[:,:,0]\n",
    "    mask = np.ones(h.shape, dtype=np.uint8) * 255\n",
    "    th = (25, 175)\n",
    "    mask[(h > th[0]) & (h < th[1])] = 0\n",
    "    copyImg = cv2.cvtColor(copyImg, cv2.COLOR_HSV2BGR)\n",
    "    resROI = cv2.bitwise_and(copyImg, copyImg, mask=mask)\n",
    "        \n",
    "    image_gray = cv2.cvtColor(resROI, cv2.COLOR_BGR2GRAY)\n",
    "    _, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))\n",
    "    morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)\n",
    "    contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    contours = contours[0] if len(contours) == 2 else contours[1]\n",
    "    bigCont = max(contours, key=cv2.contourArea)\n",
    "    x, y, w, h = cv2.boundingRect(bigCont)\n",
    "    crop = image[y : y + h, x : x + w]\n",
    "    return crop\n",
    "\n",
    "    \n",
    "\n",
    "#RN50ModelToEval = M.RARP_NVB_ResNet50_VAN.load_from_checkpoint(\"./log_X12_VAN_Review/lightning_logs/version_17/checkpoints/RARP-epoch=29.ckpt\")\n",
    "RN50ModelToEval = M.RARP_NVB_DINO_MultiTask.load_from_checkpoint(\"./log_X13_van_DINO/lightning_logs/version_33/checkpoints/RARP-epoch=32.ckpt\")\n",
    "th = 0.6299"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean, std = ([30.38144216, 42.03988769, 97.8896116], [40.63141752, 44.26910074, 50.29294373])\n",
    "transforms = T.Compose([\n",
    "    T.Resize((256,256), antialias=True, interpolation=T.InterpolationMode.BICUBIC),\n",
    "    T.CenterCrop(224),\n",
    "    T.Normalize(mean, std)\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.8702], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "pathImg = Path(\"D:/Users/user/Documents/postata/RARP/Clasification/DataSet_big/NVB/517.tiff\")\n",
    "frameToFind = cv2.imread(str(pathImg), cv2.IMREAD_COLOR)\n",
    "frameToFind = _removeBlackBorder(frameToFind)\n",
    "frameToFind = torch.Tensor(frameToFind)\n",
    "frameToFind = frameToFind.permute(2, 0, 1).float()\n",
    "\n",
    "originalSize = frameToFind.shape[-2:]\n",
    "originalImg = frameToFind\n",
    "\n",
    "frameToFind = transforms(frameToFind)\n",
    "\n",
    "frameToFind = frameToFind.repeat(1, 1, 1, 1)\n",
    "\n",
    "torch.set_float32_matmul_precision('high')\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "RN50ModelToEval.to(device)\n",
    "RN50ModelToEval.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    frameToFind = frameToFind.to(device)\n",
    "    \n",
    "    Doutput, _, _ = RN50ModelToEval(frameToFind)\n",
    "    \n",
    "    Doutput = Doutput.flatten()\n",
    "    \n",
    "    pred = torch.sigmoid(Doutput)\n",
    "    \n",
    "    print (pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_grad_cam import GradCAM, GradCAMPlusPlus, EigenGradCAM\n",
    "from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image\n",
    "\n",
    "targetL = [RN50ModelToEval.student.backbone.block4[-1], RN50ModelToEval.teacher_Features.backbone.block4[-1]]\n",
    "\n",
    "for param in RN50ModelToEval.teacher_Features.backbone.parameters():\n",
    "    param.requires_grad = True\n",
    "\n",
    "#CAM = EigenGradCAM(model=RN50ModelToEval, target_layers=targetL)\n",
    "CAM = GradCAMPlusPlus(model=RN50ModelToEval, target_layers=targetL) \n",
    "#CAM = GradCAM(model=RN50ModelToEval, target_layers=targetL)\n",
    "tar = [ClassifierOutputTarget(0)]\n",
    "\n",
    "\n",
    "gi = CAM(frameToFind, targets=tar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Denorlalize (img:torch.Tensor, std, mean):\n",
    "    ImgNumpy = img.numpy().transpose((1, 2, 0))\n",
    "    ImgNumpy = np.clip((std * ImgNumpy + mean) / 255, 0, 1)\n",
    "    ImgNumpy = ImgNumpy[...,::-1].copy()\n",
    "    \n",
    "    return ImgNumpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1972204 characters]",
      "text/plain": [
       "<Figure size 2000x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "\n",
    "rgb_img = Denorlalize(frameToFind.cpu()[0], std, mean)\n",
    "oriImg = Denorlalize(originalImg.cpu(), [1, 1, 1], [0, 0, 0])\n",
    "\n",
    "grayscale_cam = gi[0]\n",
    "\n",
    "visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)\n",
    "\n",
    "smallH, smallW, _ = rgb_img.shape\n",
    "x, y = (16, 16) #(256 - 224) // 2\n",
    "over = np.ones((256, 256), dtype=np.uint8)\n",
    "layer2 = np.zeros((224, 224), dtype=np.uint8)\n",
    "over[y:y + smallH, x:x + smallW] = layer2\n",
    "over = cv2.resize(over, (originalSize[1], originalSize[0]))\n",
    "over = cv2.bitwise_and(oriImg, oriImg, mask=over)\n",
    "\n",
    "bkMask = cv2.inRange(over, (0,0,0), (0,0,0))\n",
    "contours, _ = cv2.findContours(bkMask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "xBK, yBK, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea)) \n",
    "\n",
    "layer2 = cv2.resize(visualization / 255, (w, h))\n",
    "\n",
    "over[yBK:yBK + h, xBK:xBK + w] = layer2\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(20, 8))\n",
    "\n",
    "ax[0].imshow(oriImg)\n",
    "ax[0].set_title(f\"Original Image [{pathImg.name}]\")\n",
    "ax[0].axis(\"off\")\n",
    "\n",
    "ax[1].imshow(over)\n",
    "ax[1].set_title(f\"CAM Prediction: {'NVB' if pred.item() > th else 'NO_NVB'}; Conf.: {pred.item():.4f}\")\n",
    "ax[1].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyRARP",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
