Detecting COVID-19 with Chest X Ray.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detecting COVID-19 with Chest X Ray using PyTorch\n",
    "\n",
    "Image classification of Chest X Rays in one of three classes: Normal, Viral Pneumonia, COVID-19\n",
    "Dataset from [COVID-19 Radiography Dataset](https://www.kaggle.com/tawsifurrahman/covid19-radiography-database) on Kaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "import random\n",
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "print('Using PyTorch version', torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparing Training and Test Sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = ['normal', 'viral', 'covid']\n",
    "root_dir = 'COVID-19 Radiography Database'\n",
    "source_dirs = ['NORMAL', 'Viral Pneumonia', 'COVID-19']\n",
    "\n",
    "if os.path.isdir(os.path.join(root_dir, source_dirs[1])):\n",
    "    os.mkdir(os.path.join(root_dir, 'test'))\n",
    "\n",
    "    for i, d in enumerate(source_dirs):\n",
    "        os.rename(os.path.join(root_dir, d), os.path.join(root_dir, class_names[i]))\n",
    "\n",
    "    for c in class_names:\n",
    "        os.mkdir(os.path.join(root_dir, 'test', c))\n",
    "\n",
    "    for c in class_names:\n",
    "        images = [x for x in os.listdir(os.path.join(root_dir, c)) if x.lower().endswith('png')]\n",
    "        selected_images = random.sample(images, 30)\n",
    "        for image in selected_images:\n",
    "            source_path = os.path.join(root_dir, c, image)\n",
    "            target_path = os.path.join(root_dir, 'test', c, image)\n",
    "            shutil.move(source_path, target_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChestXRayDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, image_dirs, transform):\n",
    "        def get_images(class_name):\n",
    "            images = [x for x in os.listdir(image_dirs[class_name]) if x[-3:].lower().endswith('png')]\n",
    "            print(f'Found {len(images)} {class_name} examples')\n",
    "            return images\n",
    "        \n",
    "        self.images = {}\n",
    "        self.class_names = ['normal', 'viral', 'covid']\n",
    "        \n",
    "        for class_name in self.class_names:\n",
    "            self.images[class_name] = get_images(class_name)\n",
    "            \n",
    "        self.image_dirs = image_dirs\n",
    "        self.transform = transform\n",
    "        \n",
    "    \n",
    "    def __len__(self):\n",
    "        return sum([len(self.images[class_name]) for class_name in self.class_names])\n",
    "    \n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        class_name = random.choice(self.class_names)\n",
    "        index = index % len(self.images[class_name])\n",
    "        image_name = self.images[class_name][index]\n",
    "        image_path = os.path.join(self.image_dirs[class_name], image_name)\n",
    "        image = Image.open(image_path).convert('RGB')\n",
    "        return self.transform(image), self.class_names.index(class_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(size=(224, 224)),\n",
    "    torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "test_transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(size=(224, 224)),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dirs = {\n",
    "    'normal': 'COVID-19 Radiography Database/normal',\n",
    "    'viral': 'COVID-19 Radiography Database/viral',\n",
    "    'covid': 'COVID-19 Radiography Database/covid'\n",
    "}\n",
    "\n",
    "train_dataset = ChestXRayDataset(train_dirs, train_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dirs = {\n",
    "    'normal': 'COVID-19 Radiography Database/test/normal',\n",
    "    'viral': 'COVID-19 Radiography Database/test/viral',\n",
    "    'covid': 'COVID-19 Radiography Database/test/covid'\n",
    "}\n",
    "\n",
    "test_dataset = ChestXRayDataset(test_dirs, test_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 6\n",
    "\n",
    "dl_train = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "dl_test = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "print('Number of training batches', len(dl_train))\n",
    "print('Number of test batches', len(dl_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = train_dataset.class_names\n",
    "\n",
    "\n",
    "def show_images(images, labels, preds):\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    for i, image in enumerate(images):\n",
    "        plt.subplot(1, 6, i + 1, xticks=[], yticks=[])\n",
    "        image = image.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",
    "        image = image * std + mean\n",
    "        image = np.clip(image, 0., 1.)\n",
    "        plt.imshow(image)\n",
    "        col = 'green'\n",
    "        if preds[i] != labels[i]:\n",
    "            col = 'red'\n",
    "            \n",
    "        plt.xlabel(f'{class_names[int(labels[i].numpy())]}')\n",
    "        plt.ylabel(f'{class_names[int(preds[i].numpy())]}', color=col)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images, labels = next(iter(dl_train))\n",
    "show_images(images, labels, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images, labels = next(iter(dl_test))\n",
    "show_images(images, labels, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet18 = torchvision.models.resnet18(pretrained=True)\n",
    "\n",
    "print(resnet18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet18.fc = torch.nn.Linear(in_features=512, out_features=3)\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(resnet18.parameters(), lr=3e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_preds():\n",
    "    resnet18.eval()\n",
    "    images, labels = next(iter(dl_test))\n",
    "    outputs = resnet18(images)\n",
    "    _, preds = torch.max(outputs, 1)\n",
    "    show_images(images, labels, preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_preds()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epochs):\n",
    "    print('Starting training..')\n",
    "    for e in range(0, epochs):\n",
    "        print('='*20)\n",
    "        print(f'Starting epoch {e + 1}/{epochs}')\n",
    "        print('='*20)\n",
    "\n",
    "        train_loss = 0.\n",
    "        val_loss = 0.\n",
    "\n",
    "        resnet18.train() # set model to training phase\n",
    "\n",
    "        for train_step, (images, labels) in enumerate(dl_train):\n",
    "            optimizer.zero_grad()\n",
    "            outputs = resnet18(images)\n",
    "            loss = loss_fn(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            train_loss += loss.item()\n",
    "            if train_step % 20 == 0:\n",
    "                print('Evaluating at step', train_step)\n",
    "\n",
    "                accuracy = 0\n",
    "\n",
    "                resnet18.eval() # set model to eval phase\n",
    "\n",
    "                for val_step, (images, labels) in enumerate(dl_test):\n",
    "                    outputs = resnet18(images)\n",
    "                    loss = loss_fn(outputs, labels)\n",
    "                    val_loss += loss.item()\n",
    "\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    accuracy += sum((preds == labels).numpy())\n",
    "\n",
    "                val_loss /= (val_step + 1)\n",
    "                accuracy = accuracy/len(test_dataset)\n",
    "                print(f'Validation Loss: {val_loss:.4f}, Accuracy: {accuracy:.4f}')\n",
    "\n",
    "                show_preds()\n",
    "\n",
    "                resnet18.train()\n",
    "\n",
    "                if accuracy >= 0.95:\n",
    "                    print('Performance condition satisfied, stopping..')\n",
    "                    return\n",
    "\n",
    "        train_loss /= (train_step + 1)\n",
    "\n",
    "        print(f'Training Loss: {train_loss:.4f}')\n",
    "    print('Training complete..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "train(epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_preds()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}