{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a910e8f8-64ff-428d-bf86-98fb25885057",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<contextlib.ExitStack at 0x7f71ac3210f0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "from torch.utils.data import random_split\n",
    "from torch.utils.data.dataset import Dataset\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",
    "from time import time\n",
    "import os\n",
    "from PIL import Image\n",
    "from tempfile import TemporaryDirectory\n",
    "from copy import copy\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "# cudnn.benchmark = True\n",
    "plt.ion()   # interactive mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d508ea55-65cd-49c2-916b-ef8481af2cfe",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Computing mean and std..\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01bfd441eb874a4bb9228cf169651900",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1518 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.3227, 0.2076, 0.1405]) tensor([0.0816, 0.0715, 0.0502])\n",
      "time elapsed:  6.0041279792785645\n"
     ]
    }
   ],
   "source": [
    "# Determine specific mean and std for normalization\n",
    "\n",
    "N_CHANNELS = 3\n",
    "\n",
    "data_dir = 'img/GT_square'\n",
    "dataset = datasets.ImageFolder(data_dir, transform=transforms.ToTensor())\n",
    "full_loader = torch.utils.data.DataLoader(dataset, shuffle=False, num_workers=4)\n",
    "\n",
    "before = time()\n",
    "mean = torch.zeros(N_CHANNELS)\n",
    "std = torch.zeros(N_CHANNELS)\n",
    "print('==> Computing mean and std..')\n",
    "for inputs, _labels in tqdm(full_loader):\n",
    "    for i in range(N_CHANNELS):\n",
    "        mean[i] += inputs[:,i,:,:].mean()\n",
    "        std[i] += inputs[:,i,:,:].std()\n",
    "mean.div_(len(dataset))\n",
    "std.div_(len(dataset))\n",
    "print(mean, std)\n",
    "m_np = mean.numpy()\n",
    "std_np = std.numpy()\n",
    "\n",
    "print(\"time elapsed: \", time()-before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "edf0594c-e6ba-43f8-b460-d77b186800d3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1518\n",
      "1214\n",
      "303\n",
      "1517\n"
     ]
    }
   ],
   "source": [
    "print(len(dataset))\n",
    "print(int(len(dataset)*0.8))\n",
    "print(int(len(dataset)*0.2))\n",
    "print(int(len(dataset)*0.8) + int(len(dataset)*0.2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7040db69-fb4b-485c-bb88-20699cfd3063",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Data augmentation and normalization for training\n",
    "# Just normalization for validation\n",
    "data_transforms = {\n",
    "    'train': transforms.Compose([\n",
    "        # transforms.RandomResizedCrop(224),\n",
    "        # transforms.RandomHorizontalFlip(),\n",
    "        transforms.AutoAugment(),\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",
    "        # transforms.Normalize([0.2138, 0.0986, 0.0467], [0.0856, 0.0697, 0.0405])\n",
    "        transforms.Normalize(m_np, std_np)\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",
    "        # transforms.Normalize([0.2138, 0.0986, 0.0467], [0.0856, 0.0697, 0.0405])\n",
    "        transforms.Normalize(m_np, std_np)\n",
    "    ]),\n",
    "}\n",
    "\n",
    "data_dir = 'img/GT_square'\n",
    "image_dataset = datasets.ImageFolder(data_dir)\n",
    "class_names = image_dataset.classes\n",
    "\n",
    "train_data, val_data = random_split(image_dataset, [1214, 304])\n",
    "train_data.dataset = copy(image_dataset)\n",
    "\n",
    "train_data.dataset.transform = data_transforms['train']\n",
    "val_data.dataset.transform = data_transforms['val']\n",
    "\n",
    "image_datasets = {'train': train_data, 'val': val_data}\n",
    "\n",
    "dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,\n",
    "                                             shuffle=x=='train', num_workers=4)\n",
    "              for x in ['train', 'val']}\n",
    "dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "594fc164-68cc-48ab-a583-ffc8a1f48a1b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2415164 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 32 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# def imshow(inp, title=None):\n",
    "#     \"\"\"Display image for Tensor.\"\"\"\n",
    "#     inp = inp.numpy().transpose((1, 2, 0))\n",
    "#     # mean = np.array([0.485, 0.456, 0.406])\n",
    "#     # mean = np.array([0.2138, 0.0986, 0.0467])\n",
    "#     mean = np.array([0.3227, 0.2076, 0.1405])\n",
    "    \n",
    "#     # std = np.array([0.229, 0.224, 0.225])\n",
    "#     # std = np.array([0.0856, 0.0697, 0.0405])\n",
    "#     std = np.array([0.0816, 0.0715, 0.0502])\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\n",
    "\n",
    "def imshow(inp, title=None):\n",
    "    \"\"\"Display image for Tensor.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = m_np\n",
    "    std = std_np\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    plt.imshow(inp)\n",
    "    \n",
    "    \n",
    "def show_dataset(dataloader, num_images=16):\n",
    "    images_so_far = 0\n",
    "    cols = 8\n",
    "    rows = int(np.ceil(num_images/cols))\n",
    "    fig = plt.figure(figsize=(int(cols*2.5),int(rows*2.5)))\n",
    "    \n",
    "    for batch_idx, sample in enumerate(dataloader):\n",
    "        images, labels = sample\n",
    "        for i, image in enumerate(images):\n",
    "            \n",
    "            ax = plt.subplot(rows, cols, images_so_far + 1)\n",
    "            ax.axis('off')\n",
    "            ax.set_title(f\"GT: {class_names[labels[i]]}\")\n",
    "            imshow(image)\n",
    "            images_so_far += 1\n",
    "        \n",
    "            if images_so_far == num_images:\n",
    "                return\n",
    "\n",
    "show_dataset(dataloaders['train'], num_images=32)\n",
    "plt.savefig('img/documentation/training_sample.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "51d68c1f-e0da-446c-b332-c8b21e83dcf7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_model(model, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time()\n",
    "\n",
    "    # Create a temporary directory to save training checkpoints\n",
    "    with TemporaryDirectory() as tempdir:\n",
    "        best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')\n",
    "\n",
    "        torch.save(model.state_dict(), best_model_params_path)\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 tqdm(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",
    "                    torch.save(model.state_dict(), best_model_params_path)\n",
    "\n",
    "            print()\n",
    "\n",
    "        time_elapsed = 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(torch.load(best_model_params_path))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "11df6762-cc4b-4b76-977b-14f0e44acc12",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def visualize_model(model, num_images=8, data=dataloaders['val']):\n",
    "    was_training = model.training\n",
    "    model.eval()\n",
    "    images_so_far = 0\n",
    "    cols = 8\n",
    "    rows = int(np.ceil(num_images/cols))\n",
    "    fig = plt.figure(figsize=(int(cols*2.5),int(rows*2.5)))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i, (inputs, labels) in enumerate(data):\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(rows, cols, images_so_far)\n",
    "                ax.axis('off')\n",
    "                ax.set_title(f'GT: {class_names[labels[j]]}, 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": 8,
   "id": "beb80475-a343-45f7-a229-e7b974001a86",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_ft = models.resnext50_32x4d(weights='DEFAULT')\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, len(class_names))\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2d25a8c0-bff0-4693-b72c-e18fcc38d186",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7be28b8d28f341f0b4261f2b6a9dc5f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.5536 Acc: 0.8064\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "475404e51e1945afba40add9d9800d0e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.1915 Acc: 0.9572\n",
      "\n",
      "Epoch 1/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c746f4f61db4755833429fe90a3d5ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.1480 Acc: 0.9498\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87e183b36b2148f38495c5bca404b14b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.1209 Acc: 0.9737\n",
      "\n",
      "Epoch 2/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4829f5393a844e9a8aeaa9e5fefff395",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0972 Acc: 0.9695\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15964298294348fe9dce0cf905da1ce7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0549 Acc: 0.9967\n",
      "\n",
      "Epoch 3/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "579e6a3f2d6c46ed8ff972beb35811f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0588 Acc: 0.9835\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3503b28ee00848f9b6357a35ab5d879e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0407 Acc: 0.9967\n",
      "\n",
      "Epoch 4/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15ed40b48e314cb1be700f17a9d84499",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0527 Acc: 0.9835\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ce4a865b038485a917d6af8e6ef4e88",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0362 Acc: 0.9967\n",
      "\n",
      "Epoch 5/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a5b190f3b1b4c8e8064ad421c3cc9fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0346 Acc: 0.9918\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "49d46dfd39834ffeafe6c025dc06ab09",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0280 Acc: 1.0000\n",
      "\n",
      "Epoch 6/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc7ee96f68e84483852d8f577b798187",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0340 Acc: 0.9918\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b04e47ddd6e4df3be2f43e2b3727552",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0121 Acc: 1.0000\n",
      "\n",
      "Epoch 7/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "264b12325e2b42be8e4622f7fd0089cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0278 Acc: 0.9942\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "547843697504431986823faf59c831ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0360 Acc: 1.0000\n",
      "\n",
      "Epoch 8/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d2a11475d154bef9d7b298882e80229",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0161 Acc: 0.9992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7df764a98e4c45b8b10ffcb84e445be7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0235 Acc: 1.0000\n",
      "\n",
      "Epoch 9/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18b70b80deb04ad5b99be2a6bd35ba44",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0144 Acc: 0.9942\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "132f795a21da4bf9814592c396095de9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0269 Acc: 1.0000\n",
      "\n",
      "Epoch 10/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d2cccb6ae294aecbbcd12406ed4a82b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0186 Acc: 0.9967\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4298d409de794044958b18d5d79d0169",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0184 Acc: 1.0000\n",
      "\n",
      "Epoch 11/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00ce0d54b19343e894fc7310225c10c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0269 Acc: 0.9942\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7863c8eca24407f861fcf40b4fcaca9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0294 Acc: 1.0000\n",
      "\n",
      "Epoch 12/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dccc17bc9642445c8b38fc2d29512691",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0152 Acc: 0.9975\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "873cb88e5eee4057a9d04a6f02176641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0225 Acc: 1.0000\n",
      "\n",
      "Epoch 13/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38c36de72ca24b6f9b6d43ec85952ca4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0178 Acc: 0.9951\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6ad6c9b03f140aa80c25059abf62483",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0329 Acc: 1.0000\n",
      "\n",
      "Epoch 14/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbb2e1ef898c468b87a5510998e14352",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0136 Acc: 0.9975\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2a8aa707de14fd49a453025b1df0b46",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0251 Acc: 1.0000\n",
      "\n",
      "Epoch 15/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a418b86110dc4952b0d567019e074b71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0096 Acc: 0.9984\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0b4417a5efa4ae3a4d9e611c9577117",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0347 Acc: 1.0000\n",
      "\n",
      "Epoch 16/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bf060197c3d42a9ac3e2b1d996f1dd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0172 Acc: 0.9942\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e27a307b1733485585a54792504fea90",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0182 Acc: 1.0000\n",
      "\n",
      "Epoch 17/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65b4f4f11cd2407689abca5db1010ecf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0142 Acc: 0.9959\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71b4cc7955b84f419267a95324c4f6a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0354 Acc: 0.9967\n",
      "\n",
      "Epoch 18/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92962392ced84d76b1f1dd0837e87c4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0087 Acc: 0.9992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5e1f70e8e744359a2723dc5f3e3f94c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0336 Acc: 1.0000\n",
      "\n",
      "Epoch 19/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "868dede501ba4f31891567269c9927ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0180 Acc: 0.9992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5aad4345da754fa99b532891d0a5b1ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0182 Acc: 1.0000\n",
      "\n",
      "Epoch 20/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2537e500f17a4323a9ecfc94aea99bdd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0176 Acc: 0.9959\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a657263ae684745b3f4f4d3acbfa414",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0168 Acc: 1.0000\n",
      "\n",
      "Epoch 21/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea6d35f5dd0c48689768cf6e560b8e8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0098 Acc: 0.9992\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3e5a16eb99c4ad8913550f70fff8557",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0138 Acc: 1.0000\n",
      "\n",
      "Epoch 22/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11b9ee1cb6cb4d23a6dfbf5a7c8b4365",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0124 Acc: 0.9967\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b5014ea2448b45cfbff99887da48abda",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0231 Acc: 1.0000\n",
      "\n",
      "Epoch 23/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf67aa19fad44e7abcbdbaee97d9676b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0111 Acc: 0.9975\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b2e4c9faa3242e69095c869558c02c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0151 Acc: 1.0000\n",
      "\n",
      "Epoch 24/24\n",
      "----------\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21d0e0c156d445a8a916e3fb97d85e83",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/304 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0142 Acc: 0.9967\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e07aaa45621445359e2892045f898712",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/76 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0172 Acc: 1.0000\n",
      "\n",
      "Training complete in 6m 49s\n",
      "Best val Acc: 1.000000\n"
     ]
    }
   ],
   "source": [
    "model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,\n",
    "                       num_epochs=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2abf728f-c5ef-4951-ae72-4f0d81b65772",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.save(model_ft.state_dict(), 'models/final_GT_resnext50_32x4d_finetuned.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "eee8c387-6c4a-4ce6-99e8-fa261980f380",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model_ft.load_state_dict(torch.load('models/final_GT_resnext50_32x4d_finetuned.pth'))\n",
    "model_ft.eval()\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e686626d-2d9b-4eae-8120-7722b32da333",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1832084 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 32 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_model(model_ft, num_images=32)\n",
    "plt.savefig('img/documentation/validation_sample.jpg', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "73c126c0-555e-4428-b7f0-190d2b20aa56",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11\n"
     ]
    }
   ],
   "source": [
    "# data_dir = 'img/GT_cutout'\n",
    "# image_dataset = datasets.ImageFolder(data_dir, data_transforms['val'])\n",
    "# class_names = image_dataset.classes\n",
    "\n",
    "add_dir = 'img/BED_testdata'\n",
    "BED_add = datasets.ImageFolder(add_dir, data_transforms['val'])\n",
    "BED_dl = torch.utils.data.DataLoader(BED_add, batch_size=4, num_workers=1)\n",
    "print(len(BED_add))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "91d10134-9590-4beb-bb70-7842794b65c6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 712624 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 11 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_model(model_ft, num_images=len(BED_add), data=BED_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b19a651d-7444-496e-ad16-7e464e7f10d5",
   "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
