{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b88ce481",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import torch.backends.cudnn as cudnn\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import os\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d0cc05b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cudnn.benchmark = True\n",
    "plt.ion()   # interactive mode\n",
    "\n",
    "data_transforms = {\n",
    "    'train': transforms.Compose([\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'val': transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "}\n",
    "\n",
    "data_dir = '/home/zenon/training/plantsdata'\n",
    "image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),\n",
    "                                          data_transforms[x])\n",
    "                  for x in ['train', 'val']}\n",
    "dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,\n",
    "                                             shuffle=True, num_workers=4)\n",
    "              for x in ['train', 'val']}\n",
    "dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}\n",
    "class_names = image_datasets['train'].classes\n",
    "\n",
    "device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a779e636",
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(inp, title=None):\n",
    "    \"\"\"Imshow for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.pause(0.001)  # pause a bit so that plots are updated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "25cc39cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f'Epoch {epoch}/{num_epochs - 1}')\n",
    "        print('-' * 10)\n",
    "\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # Set model to training mode\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "\n",
    "            # Iterate over data.\n",
    "            for inputs, labels in dataloaders[phase]:\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    loss = criterion(outputs, labels)\n",
    "\n",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / dataset_sizes[phase]\n",
    "            epoch_acc = running_corrects.double() / dataset_sizes[phase]\n",
    "\n",
    "            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')\n",
    "\n",
    "            # deep copy the model\n",
    "            if phase == 'val' and epoch_acc > best_acc:\n",
    "                best_acc = epoch_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')\n",
    "    print(f'Best val Acc: {best_acc:4f}')\n",
    "\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "03835995",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_model(model, num_images=6):\n",
    "    was_training = model.training\n",
    "    model.eval()\n",
    "    images_so_far = 0\n",
    "    fig = plt.figure()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i, (inputs, labels) in enumerate(dataloaders['val']):\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "\n",
    "            for j in range(inputs.size()[0]):\n",
    "                images_so_far += 1\n",
    "                ax = plt.subplot(num_images//2, 2, images_so_far)\n",
    "                ax.axis('off')\n",
    "                ax.set_title(f'predicted: {class_names[preds[j]]}')\n",
    "                imshow(inputs.cpu().data[j])\n",
    "\n",
    "                if images_so_far == num_images:\n",
    "                    model.train(mode=was_training)\n",
    "                    return\n",
    "        model.train(mode=was_training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "35821584",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/zenon/.local/share/miniconda3/lib/python3.7/site-packages/torchvision/models/_utils.py:209: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead.\n",
      "  f\"The parameter '{pretrained_param}' is deprecated since 0.13 and will be removed in 0.15, \"\n",
      "/home/zenon/.local/share/miniconda3/lib/python3.7/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing `weights=ResNet50_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet50_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/49\n",
      "----------\n",
      "train Loss: 0.6112 Acc: 0.7157\n",
      "val Loss: 0.2829 Acc: 0.8725\n",
      "\n",
      "Epoch 1/49\n",
      "----------\n",
      "train Loss: 0.6153 Acc: 0.7531\n",
      "val Loss: 0.3555 Acc: 0.8725\n",
      "\n",
      "Epoch 2/49\n",
      "----------\n",
      "train Loss: 0.5646 Acc: 0.7718\n",
      "val Loss: 0.6510 Acc: 0.8137\n",
      "\n",
      "Epoch 3/49\n",
      "----------\n",
      "train Loss: 0.5638 Acc: 0.7556\n",
      "val Loss: 0.6726 Acc: 0.7843\n",
      "\n",
      "Epoch 4/49\n",
      "----------\n",
      "train Loss: 0.5397 Acc: 0.7768\n",
      "val Loss: 0.3414 Acc: 0.8824\n",
      "\n",
      "Epoch 5/49\n",
      "----------\n",
      "train Loss: 0.6160 Acc: 0.7706\n",
      "val Loss: 0.4935 Acc: 0.8137\n",
      "\n",
      "Epoch 6/49\n",
      "----------\n",
      "train Loss: 0.5071 Acc: 0.7893\n",
      "val Loss: 0.3808 Acc: 0.8725\n",
      "\n",
      "Epoch 7/49\n",
      "----------\n",
      "train Loss: 0.3826 Acc: 0.8329\n",
      "val Loss: 0.3154 Acc: 0.8922\n",
      "\n",
      "Epoch 8/49\n",
      "----------\n",
      "train Loss: 0.3947 Acc: 0.8254\n",
      "val Loss: 0.2607 Acc: 0.9118\n",
      "\n",
      "Epoch 9/49\n",
      "----------\n",
      "train Loss: 0.3388 Acc: 0.8491\n",
      "val Loss: 0.2881 Acc: 0.9020\n",
      "\n",
      "Epoch 10/49\n",
      "----------\n",
      "train Loss: 0.3538 Acc: 0.8441\n",
      "val Loss: 0.2617 Acc: 0.8627\n",
      "\n",
      "Epoch 11/49\n",
      "----------\n",
      "train Loss: 0.3317 Acc: 0.8479\n",
      "val Loss: 0.3005 Acc: 0.8824\n",
      "\n",
      "Epoch 12/49\n",
      "----------\n",
      "train Loss: 0.3711 Acc: 0.8479\n",
      "val Loss: 0.2452 Acc: 0.9020\n",
      "\n",
      "Epoch 13/49\n",
      "----------\n",
      "train Loss: 0.3140 Acc: 0.8653\n",
      "val Loss: 0.2736 Acc: 0.9020\n",
      "\n",
      "Epoch 14/49\n",
      "----------\n",
      "train Loss: 0.3133 Acc: 0.8678\n",
      "val Loss: 0.2505 Acc: 0.9020\n",
      "\n",
      "Epoch 15/49\n",
      "----------\n",
      "train Loss: 0.3213 Acc: 0.8628\n",
      "val Loss: 0.3173 Acc: 0.8824\n",
      "\n",
      "Epoch 16/49\n",
      "----------\n",
      "train Loss: 0.3030 Acc: 0.8641\n",
      "val Loss: 0.2764 Acc: 0.9020\n",
      "\n",
      "Epoch 17/49\n",
      "----------\n",
      "train Loss: 0.2895 Acc: 0.8766\n",
      "val Loss: 0.2413 Acc: 0.8824\n",
      "\n",
      "Epoch 18/49\n",
      "----------\n",
      "train Loss: 0.3089 Acc: 0.8778\n",
      "val Loss: 0.2791 Acc: 0.8922\n",
      "\n",
      "Epoch 19/49\n",
      "----------\n",
      "train Loss: 0.2669 Acc: 0.8890\n",
      "val Loss: 0.2582 Acc: 0.8824\n",
      "\n",
      "Epoch 20/49\n",
      "----------\n",
      "train Loss: 0.2608 Acc: 0.8853\n",
      "val Loss: 0.2510 Acc: 0.8824\n",
      "\n",
      "Epoch 21/49\n",
      "----------\n",
      "train Loss: 0.2779 Acc: 0.8890\n",
      "val Loss: 0.3014 Acc: 0.8922\n",
      "\n",
      "Epoch 22/49\n",
      "----------\n",
      "train Loss: 0.2893 Acc: 0.8691\n",
      "val Loss: 0.2880 Acc: 0.8824\n",
      "\n",
      "Epoch 23/49\n",
      "----------\n",
      "train Loss: 0.3186 Acc: 0.8678\n",
      "val Loss: 0.2615 Acc: 0.9020\n",
      "\n",
      "Epoch 24/49\n",
      "----------\n",
      "train Loss: 0.2633 Acc: 0.8953\n",
      "val Loss: 0.3043 Acc: 0.8824\n",
      "\n",
      "Epoch 25/49\n",
      "----------\n",
      "train Loss: 0.2944 Acc: 0.8803\n",
      "val Loss: 0.2490 Acc: 0.9020\n",
      "\n",
      "Epoch 26/49\n",
      "----------\n",
      "train Loss: 0.3139 Acc: 0.8603\n",
      "val Loss: 0.2632 Acc: 0.9020\n",
      "\n",
      "Epoch 27/49\n",
      "----------\n",
      "train Loss: 0.3356 Acc: 0.8491\n",
      "val Loss: 0.2702 Acc: 0.9118\n",
      "\n",
      "Epoch 28/49\n",
      "----------\n",
      "train Loss: 0.2870 Acc: 0.8840\n",
      "val Loss: 0.2819 Acc: 0.8824\n",
      "\n",
      "Epoch 29/49\n",
      "----------\n",
      "train Loss: 0.2924 Acc: 0.8803\n",
      "val Loss: 0.2686 Acc: 0.9020\n",
      "\n",
      "Epoch 30/49\n",
      "----------\n",
      "train Loss: 0.3077 Acc: 0.8741\n",
      "val Loss: 0.3752 Acc: 0.8529\n",
      "\n",
      "Epoch 31/49\n",
      "----------\n",
      "train Loss: 0.3051 Acc: 0.8766\n",
      "val Loss: 0.2668 Acc: 0.8824\n",
      "\n",
      "Epoch 32/49\n",
      "----------\n",
      "train Loss: 0.2768 Acc: 0.8928\n",
      "val Loss: 0.2893 Acc: 0.8824\n",
      "\n",
      "Epoch 33/49\n",
      "----------\n",
      "train Loss: 0.2892 Acc: 0.8716\n",
      "val Loss: 0.2558 Acc: 0.9118\n",
      "\n",
      "Epoch 34/49\n",
      "----------\n",
      "train Loss: 0.3119 Acc: 0.8691\n",
      "val Loss: 0.2258 Acc: 0.9118\n",
      "\n",
      "Epoch 35/49\n",
      "----------\n",
      "train Loss: 0.2934 Acc: 0.8791\n",
      "val Loss: 0.2436 Acc: 0.9020\n",
      "\n",
      "Epoch 36/49\n",
      "----------\n",
      "train Loss: 0.2545 Acc: 0.8828\n",
      "val Loss: 0.2557 Acc: 0.8824\n",
      "\n",
      "Epoch 37/49\n",
      "----------\n",
      "train Loss: 0.3102 Acc: 0.8641\n",
      "val Loss: 0.2693 Acc: 0.9020\n",
      "\n",
      "Epoch 38/49\n",
      "----------\n",
      "train Loss: 0.2767 Acc: 0.8878\n",
      "val Loss: 0.2432 Acc: 0.8725\n",
      "\n",
      "Epoch 39/49\n",
      "----------\n",
      "train Loss: 0.3211 Acc: 0.8516\n",
      "val Loss: 0.2712 Acc: 0.8922\n",
      "\n",
      "Epoch 40/49\n",
      "----------\n",
      "train Loss: 0.2922 Acc: 0.8753\n",
      "val Loss: 0.2635 Acc: 0.8922\n",
      "\n",
      "Epoch 41/49\n",
      "----------\n",
      "train Loss: 0.2849 Acc: 0.8815\n",
      "val Loss: 0.2647 Acc: 0.8922\n",
      "\n",
      "Epoch 42/49\n",
      "----------\n",
      "train Loss: 0.2840 Acc: 0.8728\n",
      "val Loss: 0.2591 Acc: 0.8922\n",
      "\n",
      "Epoch 43/49\n",
      "----------\n",
      "train Loss: 0.2831 Acc: 0.8766\n",
      "val Loss: 0.2419 Acc: 0.8824\n",
      "\n",
      "Epoch 44/49\n",
      "----------\n",
      "train Loss: 0.2637 Acc: 0.8953\n",
      "val Loss: 0.2505 Acc: 0.9020\n",
      "\n",
      "Epoch 45/49\n",
      "----------\n",
      "train Loss: 0.2898 Acc: 0.8741\n",
      "val Loss: 0.2581 Acc: 0.8725\n",
      "\n",
      "Epoch 46/49\n",
      "----------\n",
      "train Loss: 0.2735 Acc: 0.8840\n",
      "val Loss: 0.2740 Acc: 0.8922\n",
      "\n",
      "Epoch 47/49\n",
      "----------\n",
      "train Loss: 0.2644 Acc: 0.8953\n",
      "val Loss: 0.2490 Acc: 0.9020\n",
      "\n",
      "Epoch 48/49\n",
      "----------\n",
      "train Loss: 0.3132 Acc: 0.8579\n",
      "val Loss: 0.2509 Acc: 0.8824\n",
      "\n",
      "Epoch 49/49\n",
      "----------\n",
      "train Loss: 0.2851 Acc: 0.8840\n",
      "val Loss: 0.2480 Acc: 0.8824\n",
      "\n",
      "Training complete in 82m 19s\n",
      "Best val Acc: 0.911765\n"
     ]
    }
   ],
   "source": [
    "model_ft = models.resnet50(pretrained=True)\n",
    "num_ftrs = model_ft.fc.in_features\n",
    "# Here the size of each output sample is set to 2.\n",
    "# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).\n",
    "model_ft.fc = nn.Linear(num_ftrs, 2)\n",
    "\n",
    "model_ft = model_ft.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Observe that all parameters are being optimized\n",
    "optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "# Decay LR by a factor of 0.1 every 7 epochs\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)\n",
    "\n",
    "model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,\n",
    "                       num_epochs=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "075eed83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize_model(model_ft)\n",
    "#torch.save(model_ft, \"resnet-model.pt\")\n",
    "model_ft = torch.load(\"resnet-model.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "256788f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205a012d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## img = PIL.Image.open(\"/home/zenon/test.webp\")\n",
    "rgb_img = np.array(PIL.Image.open(\"/home/zenon/test.webp\"))\n",
    "rgb_img = rgb_img.astype(np.float32)\n",
    "rgb_img /= 255\n",
    "\n",
    "data_transforms = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "\n",
    "img = data_transforms(img)\n",
    "out = model_ft(img.unsqueeze(0))\n",
    "torch.nn.functional.softmax(out, dim=1)[0] * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "2d65d453",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_grad_cam import *\n",
    "from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image\n",
    "from torchvision.models import resnet50\n",
    "\n",
    "target_layers = [model_ft.layer4[-1]]\n",
    "input_tensor = img.unsqueeze(0)# Create an input tensor image for your model..\n",
    "# Note: input_tensor can be a batch tensor with several images!\n",
    "\n",
    "# Construct the CAM object once, and then re-use it on many images:\n",
    "cam = GradCAM(model=model_ft, target_layers=target_layers)\n",
    "\n",
    "# You can also use it within a with statement, to make sure it is freed,\n",
    "# In case you need to re-create it inside an outer loop:\n",
    "# with GradCAM(model=model, target_layers=target_layers, use_cuda=args.use_cuda) as cam:\n",
    "#   ...\n",
    "\n",
    "# We have to specify the target we want to generate\n",
    "# the Class Activation Maps for.\n",
    "# If targets is None, the highest scoring category\n",
    "# will be used for every image in the batch.\n",
    "# Here we use ClassifierOutputTarget, but you can define your own custom targets\n",
    "# That are, for example, combinations of categories, or specific outputs in a non standard model.\n",
    "\n",
    "targets = [ClassifierOutputTarget(0)]\n",
    "\n",
    "# You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.\n",
    "grayscale_cam = cam(input_tensor=input_tensor, targets=targets)\n",
    "\n",
    "# In this example grayscale_cam has only one image in the batch:\n",
    "grayscale_cam = grayscale_cam[0, :]\n",
    "visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "775243a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2ee1b8be50>"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 581044 characters]\n",
      "text/plain": [
       "<Figure size 1200x1200 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, axarr = plt.subplots(1,2, figsize=(12,12))\n",
    "axarr[0].imshow(rgb_img)\n",
    "axarr[1].imshow(visualization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0470561c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
