{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5a8f05e",
   "metadata": {},
   "source": [
    "## Question 4 (10 marks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ddab76f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from ca_utils import ResNet, BasicBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b52a47c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ResNet(block=BasicBlock, layers=[1, 1, 1], num_classes=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a25e00",
   "metadata": {},
   "source": [
    "### Load the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "af4de07b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(16, 32, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(32, 64, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=1)\n",
       "  (fc): Linear(in_features=64, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint = torch.load(\"data/weights_resnet.pth\", map_location=torch.device('cpu'))\n",
    "\n",
    "model.load_state_dict(checkpoint)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98c03b18",
   "metadata": {},
   "source": [
    "## Question 5 (15 marks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3024a7a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "357891cd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "image_transform = transforms.Compose(\n",
    "    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "\n",
    "test_data = torchvision.datasets.ImageFolder('data/EXCV10/val/', transform=image_transform)\n",
    "test_loader = DataLoader(test_data, batch_size=64, shuffle=False, num_workers=4, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb07988",
   "metadata": {},
   "source": [
    "## Method 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b9d88037",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6a30d88c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def m1_test_cnn(model, test_loader):\n",
    "\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    all_predicted_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "\n",
    "            # Make predictions\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "\n",
    "            # Save results\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            \n",
    "            all_predicted_labels.append(predicted.cpu().numpy())\n",
    "\n",
    "    accuracy = 100 * correct / total\n",
    "    all_predicted_labels = np.concatenate(all_predicted_labels)\n",
    "\n",
    "    return all_predicted_labels, accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3fcb0a3a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 75.8%\n"
     ]
    }
   ],
   "source": [
    "m1_predicted_labels, m1_test_accuracy = m1_test_cnn(model, test_loader)\n",
    "print(f'Test Accuracy: {m1_test_accuracy}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75272741",
   "metadata": {},
   "source": [
    "### Put Students' implementations here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "47bd1202",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_cnn(model, test_loader):\n",
    "    \"\"\"\n",
    "    Test the trained ResNet model on the test dataset.\n",
    "\n",
    "    Args:\n",
    "        model (nn.Module): The trained ResNet model.\n",
    "        test_loader (DataLoader): Data loader for the test data.\n",
    "    \n",
    "    Returns:\n",
    "        float: Test accuracy.\n",
    "        list: Predicted labels.\n",
    "        list: True labels.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    predicted_labels = []\n",
    "    true_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            predicted_labels.extend(predicted.tolist())\n",
    "            true_labels.extend(labels.tolist())\n",
    "\n",
    "    accuracy = correct / total\n",
    "\n",
    "    return predicted_labels, accuracy*100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1b0db3bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 75.8%\n"
     ]
    }
   ],
   "source": [
    "predicted_labels, test_accuracy = test_cnn(model, test_loader)\n",
    "print(f'Test Accuracy: {test_accuracy}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17bc1948",
   "metadata": {},
   "source": [
    "### Test (Should output ALL PASS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "853c4db3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy:  75.8\n",
      "Score 100%: 15.0\n",
      "ALL PASS\n"
     ]
    }
   ],
   "source": [
    "assert np.allclose(predicted_labels, m1_predicted_labels)\n",
    "assert np.allclose(test_accuracy, m1_test_accuracy)\n",
    "\n",
    "print(\"Test accuracy: \", test_accuracy)\n",
    "\n",
    "if (test_accuracy >= 75):\n",
    "    print(\"Score 100%:\", 15 * 1.0)\n",
    "elif (test_accuracy >= 70):\n",
    "    print(\"Score 90%:\", 15 * 0.90)\n",
    "elif (test_accuracy >= 65):\n",
    "    print(\"Score 80%:\", 15 * 0.80)\n",
    "elif (test_accuracy >= 60):\n",
    "    print(\"Score 70%:\", 15 * 0.70)\n",
    "elif (test_accuracy >= 55):\n",
    "    print(\"Score 60%:\", 15 * 0.60)\n",
    "elif (test_accuracy >= 50):\n",
    "    print(\"Score 50%:\", 15 * 0.50)\n",
    "else:\n",
    "    print(\"Accuracy less than 50%\")\n",
    "print(\"ALL PASS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cef7dc17",
   "metadata": {},
   "source": [
    "## Question 6 (6 marks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0990f3b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "true_labels = []\n",
    "\n",
    "for images, labels in test_loader:\n",
    "    images, labels = images.to(device), labels.to(device)\n",
    "    true_labels.extend(labels.cpu().numpy())\n",
    "    \n",
    "true_labels = np.array(true_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8da35032",
   "metadata": {},
   "outputs": [],
   "source": [
    "def m1_compute_confusion_matrix(true, predictions):\n",
    "    unique_labels = np.unique(np.concatenate((true, predictions)))\n",
    "\n",
    "    confusion_mat = np.zeros((len(unique_labels), len(unique_labels)), dtype=np.int64)\n",
    "\n",
    "    label_to_index = {label: index for index,\n",
    "                      label in enumerate(unique_labels)}\n",
    "\n",
    "    for t, p in zip(true, predictions):\n",
    "        t_index = label_to_index[t]\n",
    "        p_index = label_to_index[p]\n",
    "        confusion_mat[t_index][p_index] += 1\n",
    "\n",
    "    return confusion_mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "16b6f9e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1_confusion_matrix = m1_compute_confusion_matrix(true_labels, m1_predicted_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "608265af",
   "metadata": {},
   "source": [
    "### Put Students' implementations here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1dce952c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_confusion_matrix(true_labels, predicted_labels):\n",
    "\n",
    "  # Ensure inputs are NumPy arrays\n",
    "  true_labels = np.array(true_labels)\n",
    "  predicted_labels = np.array(predicted_labels)\n",
    "\n",
    "  # Determine the number of classes\n",
    "  num_classes = len(np.unique(true_labels))\n",
    "\n",
    "  # Initialize the confusion matrix with zeros\n",
    "  cm = np.zeros((num_classes, num_classes))\n",
    "\n",
    "  # Count occurrences of true-predicted label pairs\n",
    "  for i in range(len(true_labels)):\n",
    "    cm[true_labels[i]][predicted_labels[i]] += 1\n",
    "\n",
    "  return cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "21917014",
   "metadata": {},
   "outputs": [],
   "source": [
    "confusion_matrix = m1_compute_confusion_matrix(true_labels, m1_predicted_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "935956b7",
   "metadata": {},
   "source": [
    "### Test (Should output ALL PASS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b77da2e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ALL PASS\n"
     ]
    }
   ],
   "source": [
    "assert np.allclose(confusion_matrix, m1_confusion_matrix)\n",
    "\n",
    "print(\"ALL PASS\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc0a7c7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "what",
   "language": "python",
   "name": "what"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}