{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03c2b72b-ce0e-4d2f-b283-80c5b0856078",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grad_cam(\n",
    "    model,\n",
    "    module,\n",
    "    img:np.ndarray,\n",
    "    class_id,\n",
    "    device = 'cuda',\n",
    "):\n",
    "    \"\"\"       다른 import 필요 없이 이 함수 하나만 쓰면 됩니다.\n",
    "    args :\n",
    "        model :\n",
    "            pytorch 모델\n",
    "        module :\n",
    "            model 파라미터의 일부분. 여기에 hook을 걸어야 하기 때문에 필요함\n",
    "        img : np.ndarray\n",
    "            model에 넣을 이미지. 8UC3 타입의 이미지\n",
    "        class_id : int\n",
    "            img 파라미터의 클래스    \n",
    "    \"\"\"\n",
    "    preprocess_transforms =  transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]),                 \n",
    "        transforms.Resize((224, 224))\n",
    "    ])\n",
    "    \n",
    "    acts = [0]\n",
    "    grads = [0]\n",
    "    \n",
    "    def f_hook(self, input, output):\n",
    "        acts[0] = output\n",
    "\n",
    "    def b_hook(self, grad_in, grad_out):\n",
    "        grads[0] = grad_out\n",
    "    \n",
    "    h1 = module.register_forward_hook(f_hook)\n",
    "    h2 = module.register_backward_hook(b_hook)\n",
    "\n",
    "    img_t = preprocess_transforms(img).unsqueeze(0).to(device)\n",
    "    model.to(device)\n",
    "    outs = model(img_t)\n",
    "    h1.remove()\n",
    "    h2.remove()\n",
    "    outs[0, class_id].backward()\n",
    "    \n",
    "    gap = torch.mean(grads[0][0].view(grads[0][0].size(0), grads[0][0].size(1), -1), dim = 2)\n",
    "    acts = acts[0][0]\n",
    "    gradcam = torch.nn.ReLU()(torch.sum(gap[0].reshape((gap.size()[1], 1, 1)) * acts, dim = 0))\n",
    "    arr = transforms.Resize(\n",
    "        (img.shape[0], img.shape[1]),\n",
    "    )(gradcam.unsqueeze(0))\n",
    "    gradcam_t = arr.detach().cpu().permute((1, 2, 0)).squeeze(-1)\n",
    "    \n",
    "    gradcam_img = gradcam_t * (255 / gradcam_t.max())\n",
    "    gradcam_img = gradcam_img.numpy().astype(np.uint8)\n",
    "    \n",
    "    # gradcam_img 는 1채널 np.uint8 이미지 입니다.\n",
    "    return gradcam_img\n",
    "\n",
    "\n",
    "\n",
    "# ----------------------------------------------------------\n",
    "ALPHA = 0.6   # gradcam 이미지랑 이미지를 섞을 때 gradcam_img의 세기\n",
    "# ----------------------------------------------------------\n",
    "\n",
    "img = None\n",
    "model = None                            # timm 모델\n",
    "target_layers = model.model.layer4[-2]  # model 의 레이어 중 하나\n",
    "\n",
    "\n",
    "\n",
    "gradcam_img = grad_cam(model.model, target_layers, img, 2)\n",
    "\n",
    "# uint8 타입의 gradcam_img를 컬러로 바꿈. 색을 칠하면서 1채널에서 3채널로 바뀜.\n",
    "gradcam_color = cv2.applyColorMap(\n",
    "    gradcam_img,\n",
    "    cv2.COLORMAP_PARULA\n",
    ")\n",
    "\n",
    "overlayed = (\n",
    "    img.astype(np.float32) * (1-ALPTHA) + \n",
    "    gradcam_color.astype(np.float32) * ALPHA\n",
    ").astype(np.uint8)\n",
    "\n",
    "\n",
    "plt.imshow(cv2.cvtColor( # opencv 에서는 이미지를 기본적으로 BGR 포맷으로 다룸\n",
    "    overlayed,\n",
    "    cv2.COLOR_BGR2RGB\n",
    "))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yt38",
   "language": "python",
   "name": "yt38"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}