CovidAID.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CovidAID.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aaekay/CovidAID/blob/master/CovidAID.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZV6y4tzLax6J",
        "colab_type": "text"
      },
      "source": [
        "### Enable GPU first"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mh6Kx4aR_6JY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kU4vzrHy-nVv",
        "colab_type": "text"
      },
      "source": [
        "removing the directory if already exists\n",
        "then downloading the contents"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_huWwDAtAGS2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -rf ./CovidAID\n",
        "!git clone --recursive https://github.com/aaekay/CovidAID.git\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nX5zFsgTBVCW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 224
        },
        "outputId": "213d8c9a-d30e-4b88-d7da-2bcb194a789f"
      },
      "source": [
        "!wget https://sogppg.bl.files.1drv.com/y4m7xCVzOJm2tZO9UttZgEidLVEYUsYqN4zjak7oGSlj9k75JYCOwQNF-NXmkolYeS0oUJ8STXyq5wKG7SHTBXF_uRXefR4FbA0x8-b-lQVHNVqp5Ts72XLkxvG4R6Z_EonfAX-XYeSRpkSgmOVqXGVCCOB-GQAbnxrpi0TdvZB4--kBMLqtLe7ESNTaruxnOpjt7z466AvplCvYc8UtdO86Q\n",
        "!mv /content/y4m7xCVzOJm2tZO9UttZgEidLVEYUsYqN4zjak7oGSlj9k75JYCOwQNF-NXmkolYeS0oUJ8STXyq5wKG7SHTBXF_uRXefR4FbA0x8-b-lQVHNVqp5Ts72XLkxvG4R6Z_EonfAX-XYeSRpkSgmOVqXGVCCOB-GQAbnxrpi0TdvZB4--kBMLqtLe7ESNTaruxnOpjt7z466AvplCvYc8UtdO86Q pneumo.zip\n"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-06-24 11:10:57--  https://sogppg.bl.files.1drv.com/y4m7xCVzOJm2tZO9UttZgEidLVEYUsYqN4zjak7oGSlj9k75JYCOwQNF-NXmkolYeS0oUJ8STXyq5wKG7SHTBXF_uRXefR4FbA0x8-b-lQVHNVqp5Ts72XLkxvG4R6Z_EonfAX-XYeSRpkSgmOVqXGVCCOB-GQAbnxrpi0TdvZB4--kBMLqtLe7ESNTaruxnOpjt7z466AvplCvYc8UtdO86Q\n",
            "Resolving sogppg.bl.files.1drv.com (sogppg.bl.files.1drv.com)... 13.107.42.12\n",
            "Connecting to sogppg.bl.files.1drv.com (sogppg.bl.files.1drv.com)|13.107.42.12|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 358296270 (342M) [application/zip]\n",
            "Saving to: ‘y4m7xCVzOJm2tZO9UttZgEidLVEYUsYqN4zjak7oGSlj9k75JYCOwQNF-NXmkolYeS0oUJ8STXyq5wKG7SHTBXF_uRXefR4FbA0x8-b-lQVHNVqp5Ts72XLkxvG4R6Z_EonfAX-XYeSRpkSgmOVqXGVCCOB-GQAbnxrpi0TdvZB4--kBMLqtLe7ESNTaruxnOpjt7z466AvplCvYc8UtdO86Q’\n",
            "\n",
            "y4m7xCVzOJm2tZO9Utt 100%[===================>] 341.70M  25.5MB/s    in 14s     \n",
            "\n",
            "2020-06-24 11:11:12 (24.0 MB/s) - ‘y4m7xCVzOJm2tZO9UttZgEidLVEYUsYqN4zjak7oGSlj9k75JYCOwQNF-NXmkolYeS0oUJ8STXyq5wKG7SHTBXF_uRXefR4FbA0x8-b-lQVHNVqp5Ts72XLkxvG4R6Z_EonfAX-XYeSRpkSgmOVqXGVCCOB-GQAbnxrpi0TdvZB4--kBMLqtLe7ESNTaruxnOpjt7z466AvplCvYc8UtdO86Q’ saved [358296270/358296270]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V-qzLoVSBY4t",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip ./pneumo.zip -d ./CovidAID/chest-xray-pneumonia"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6dOHiUG8Yc2L",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#for installing the torch version 0.3.0\n",
        "!wget https://download.pytorch.org/whl/cu75/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BNW-2tWjYltQ",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 207
        },
        "outputId": "2cd23ee4-90d3-4226-c705-d918cc1cb853"
      },
      "source": [
        "!pip install /content/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Processing ./torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==0.3.0.post4) (1.18.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.6/dist-packages (from torch==0.3.0.post4) (3.13)\n",
            "\u001b[31mERROR: torchvision 0.6.1+cu101 has requirement torch==1.5.1, but you'll have torch 0.3.0.post4 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: fastai 1.0.61 has requirement torch>=1.0.0, but you'll have torch 0.3.0.post4 which is incompatible.\u001b[0m\n",
            "Installing collected packages: torch\n",
            "  Found existing installation: torch 1.5.1+cu101\n",
            "    Uninstalling torch-1.5.1+cu101:\n",
            "      Successfully uninstalled torch-1.5.1+cu101\n",
            "Successfully installed torch-0.3.0.post4\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IJj1gtqsUaAL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# this is for installing requirements\n",
        "!pip install -r /content/CovidAID/requirements.txt "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BqthpdkPCLug",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dKbU2qfyC8lm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4bonNxt0K0Z-",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 275
        },
        "outputId": "45b9ac32-eef0-4fcf-fea1-fa43bdd43b1e"
      },
      "source": [
        "!python3 ./CovidAID/tools/trainer.py --mode train --freeze --checkpoint ./CovidAID/models/CovidAID_4_class.pth --bs 16 --save ./new_models"
      ],
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.\n",
            "  import pandas.util.testing as tm\n",
            "Distributed training OFF\n",
            "Loading model from ./CovidAID/models/CovidAID_4_class.pth\n",
            "Traceback (most recent call last):\n",
            "  File \"./CovidAID/tools/trainer.py\", line 434, in <module>\n",
            "    start_epoch=args.start, save_path=args.save, freeze_feature_layers=args.freeze)\n",
            "  File \"./CovidAID/tools/trainer.py\", line 73, in train\n",
            "    combine_pneumonia=self.combine_pneumonia)\n",
            "  File \"/content/CovidAID/tools/read_data.py\", line 112, in __init__\n",
            "    self.loss_weight_minus = torch.FloatTensor([self.num_normal, self.num_bact, self.num_viral, self.num_covid]).unsqueeze(0).cuda() / self.total\n",
            "  File \"/usr/local/lib/python3.6/dist-packages/torch/tensor.py\", line 330, in __div__\n",
            "    return self.div(other)\n",
            "RuntimeError: invalid argument 3: divide by zero at /pytorch/torch/lib/THC/generic/THCTensorMathPairwise.cu:88\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CJVj9pjQt1ja",
        "colab_type": "text"
      },
      "source": [
        "### Now preparing for the covid data set"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YiYjJ8UbtzkN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from collections import Counter \n"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iP39_G19t8YL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "COVID_DATA_PATH='./CovidAID/covid-chestxray-dataset'\n",
        "METADATA_CSV = os.path.join(COVID_DATA_PATH, 'metadata.csv')\n",
        "TRAIN_FILE = './CovidAID/data/covid19/train_list.txt'\n",
        "VAL_FILE = './CovidAID/data/covid19/val_list.txt'\n",
        "TEST_FILE = './CovidAID/data/covid19/test_list.txt'\n"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "udPg8oxct-WU",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 255
        },
        "outputId": "f4689714-bcff-4b41-c423-0aa75638f412"
      },
      "source": [
        "# Load patient stats\n",
        "covids = dict()\n",
        "df = pd.read_csv(METADATA_CSV)\n",
        "df = df[(df['finding'] == 'COVID-19') & (df['modality'] == 'X-ray') & (\n",
        "                (df['view'] == 'PA') | (df['view'] == 'AP') | (df['view'] == 'AP Supine')\n",
        "            )]\n",
        "print(df)"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "     patientid  offset  ... other_notes  Unnamed: 26\n",
            "0            2     0.0  ...         NaN          NaN\n",
            "1            2     3.0  ...         NaN          NaN\n",
            "2            2     5.0  ...         NaN          NaN\n",
            "3            2     6.0  ...         NaN          NaN\n",
            "4            4     0.0  ...         NaN          NaN\n",
            "..         ...     ...  ...         ...          ...\n",
            "274        155     NaN  ...         NaN          NaN\n",
            "275        156     NaN  ...         NaN          NaN\n",
            "276        157     NaN  ...         NaN          NaN\n",
            "277        158     NaN  ...         NaN          NaN\n",
            "278        159     NaN  ...         NaN          NaN\n",
            "\n",
            "[155 rows x 27 columns]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LfyLsMQuuA27",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "7b0bfabb-a433-49b2-9ea9-3b69cea05ccd"
      },
      "source": [
        "patient_ids = Counter(df['patientid'].tolist())\n",
        "covids = {k: v for k, v in sorted(patient_ids.items(), key=lambda item: item[1])}\n",
        "total_data = sum([v for k,v in covids.items()])\n",
        "print (\"Patient-#X-Rays statistics:\")\n",
        "print (covids)\n",
        "print (\"Total Images:\", total_data, '\\n')"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Patient-#X-Rays statistics:\n",
            "{11: 1, 12: 1, 14: 1, 16: 1, 18: 1, 20: 1, 21: 1, 22: 1, 23: 1, 24: 1, 32: 1, 33: 1, 34: 1, 35: 1, 39: 1, 40: 1, 41: 1, 42: 1, 43: 1, 44: 1, 45: 1, 46: 1, 47: 1, 48: 1, 49: 1, 50: 1, 52: 1, 53: 1, 54: 1, 55: 1, 56: 1, 60: 1, 61: 1, 62: 1, 63: 1, 64: 1, 65: 1, 72: 1, 74: 1, 75: 1, 76: 1, 77: 1, 78: 1, 79: 1, 80: 1, 81: 1, 82: 1, 83: 1, 84: 1, 85: 1, 93: 1, 94: 1, 99: 1, 118: 1, 137: 1, 138: 1, 140: 1, 141: 1, 145: 1, 146: 1, 147: 1, 148: 1, 149: 1, 151: 1, 152: 1, 153: 1, 154: 1, 155: 1, 156: 1, 157: 1, 158: 1, 159: 1, 4: 2, 6: 2, 36: 2, 37: 2, 51: 2, 57: 2, 59: 2, 67: 2, 68: 2, 69: 2, 70: 2, 73: 2, 86: 2, 95: 2, 98: 2, 112: 2, 113: 2, 114: 2, 115: 2, 139: 2, 143: 2, 144: 2, 150: 2, 15: 3, 17: 3, 58: 3, 71: 3, 116: 3, 117: 3, 132: 3, 2: 4, 13: 4, 19: 4, 142: 4}\n",
            "Total Images: 155 \n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmC2fZUgukVU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Assign train-val-test split\n",
        "test_patients = set({4, 15, 86, 59, 6, 82, 80, 78, 76, 65, 36, 32, 50, 18, 115, 152, 138, 70, 116})\n",
        "val_patients = set({73, 51, 48, 11, 43, 24, 112})\n"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_DzDlQMUul80",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        },
        "outputId": "1f8a57e7-7b74-48e2-9dc3-5e7513342574"
      },
      "source": [
        "print ('#Train patients:', len(set(covids.keys()).difference(test_patients.union(val_patients))))\n",
        "print ('#Test patients:', len(test_patients))\n",
        "print ('#Val patients:', len(val_patients))\n",
        "print ()\n",
        "print ('#Train data points:', sum([v for k, v in covids.items() if int(k) not in test_patients.union(val_patients)]))\n",
        "print ('#Test data points:', sum([v for k, v in covids.items() if int(k) in test_patients]))\n",
        "print ('#Val data points:', sum([v for k, v in covids.items() if int(k) in val_patients]))"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "#Train patients: 80\n",
            "#Test patients: 19\n",
            "#Val patients: 7\n",
            "\n",
            "#Train data points: 115\n",
            "#Test data points: 30\n",
            "#Val data points: 10\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zGbZHnJBuve9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Construct the split lists\n",
        "train_list = []\n",
        "test_list = []\n",
        "val_list = []"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tvdjf9dkuzjr",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "d9b839e8-49f6-4eff-94bc-9405c3072c05"
      },
      "source": [
        "for i, row in df.iterrows():\n",
        "    patient_id = row['patientid']\n",
        "    filename = os.path.join(row['folder'], row['filename'])\n",
        "\n",
        "    if int(patient_id) in test_patients:\n",
        "        test_list.append(filename)\n",
        "    elif int(patient_id) in val_patients:\n",
        "        val_list.append(filename)\n",
        "    else:\n",
        "        train_list.append(filename)\n",
        "\n",
        "print (len(train_list), len(test_list), len(val_list))"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "115 30 10\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "68W-5zLYu55D",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Write image list in file\n",
        "def make_img_list(data_file, img_file_list):\n",
        "    with open(data_file, 'w') as f:\n",
        "        for imgfile in img_file_list:\n",
        "            try: \n",
        "                assert os.path.isfile(os.path.join(COVID_DATA_PATH, imgfile))\n",
        "                f.write(\"%s\\n\" % imgfile)\n",
        "            except: \n",
        "                print (\"Image %s NOT FOUND\" % imgfile)"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A03By7JzvBmM",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "d26cb3db-a583-48bf-cc3b-cc52d31ecc3f"
      },
      "source": [
        "make_img_list(TRAIN_FILE, train_list)\n",
        "make_img_list(VAL_FILE, val_list)\n",
        "make_img_list(TEST_FILE, test_list)\n",
        "TRAIN_FILE, train_list"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('./CovidAID/data/covid19/train_list.txt',\n",
              " ['images/auntminnie-a-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg',\n",
              "  'images/auntminnie-b-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg',\n",
              "  'images/auntminnie-c-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg',\n",
              "  'images/auntminnie-d-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg',\n",
              "  'images/nCoV-radiol.2020200269.fig1-day7.jpeg',\n",
              "  'images/nejmoa2001191_f1-PA.jpeg',\n",
              "  'images/nejmoa2001191_f3-PA.jpeg',\n",
              "  'images/nejmoa2001191_f4.jpeg',\n",
              "  'images/nejmoa2001191_f5-PA.jpeg',\n",
              "  'images/ryct.2020200034.fig2.jpeg',\n",
              "  'images/ryct.2020200028.fig1a.jpeg',\n",
              "  'images/jkms-35-e79-g001-l-a.jpg',\n",
              "  'images/jkms-35-e79-g001-l-b.jpg',\n",
              "  'images/jkms-35-e79-g001-l-c.jpg',\n",
              "  'images/1-s2.0-S0929664620300449-gr2_lrg-a.jpg',\n",
              "  'images/1-s2.0-S0929664620300449-gr2_lrg-b.jpg',\n",
              "  'images/1-s2.0-S0929664620300449-gr2_lrg-c.jpg',\n",
              "  'images/1-s2.0-S0929664620300449-gr2_lrg-d.jpg',\n",
              "  'images/covid-19-pneumonia-15-PA.jpg',\n",
              "  'images/covid-19-pneumonia-2.jpg',\n",
              "  'images/covid-19-pneumonia-7-PA.jpg',\n",
              "  'images/covid-19-pneumonia-14-PA.png',\n",
              "  'images/7C69C012-7479-493F-8722-ABC29C60A2DD.jpeg',\n",
              "  'images/2C10A413-AABE-4807-8CCE-6A2025594067.jpeg',\n",
              "  'images/E1724330-1866-4581-8CD8-CEC9B8AFEDDE.jpeg',\n",
              "  'images/F2DE909F-E19C-4900-92F5-8F435B031AC6.jpeg',\n",
              "  'images/31BA3780-2323-493F-8AED-62081B9C383B.jpeg',\n",
              "  'images/1312A392-67A3-4EBF-9319-810CF6DA5EF6.jpeg',\n",
              "  'images/23E99E2E-447C-46E5-8EB2-D35D12473C39.png',\n",
              "  'images/7AF6C1AF-D249-4BD2-8C26-449304105D03.jpeg',\n",
              "  'images/1B734A89-A1BF-49A8-A1D3-66FAFA4FAC5D.jpeg',\n",
              "  'images/85E52EB3-56E9-4D67-82DA-DEA247C82886.jpeg',\n",
              "  'images/6CB4EFC6-68FA-4CD5-940C-BEFA8DAFE9A7.jpeg',\n",
              "  'images/01E392EE-69F9-4E33-BFCE-E5C968654078.jpeg',\n",
              "  'images/F63AB6CE-1968-4154-A70F-913AF154F53D.jpeg',\n",
              "  'images/93FE0BB1-022D-4F24-9727-987A07975FFB.jpeg',\n",
              "  'images/5CBC2E94-D358-401E-8928-965CCD965C5C.jpeg',\n",
              "  'images/446B2CB6-B572-40AB-B01F-1910CA07086A.jpeg',\n",
              "  'images/80446565-E090-4187-A031-9D3CEAA586C8.jpeg',\n",
              "  'images/7E335538-2F86-424E-A0AB-6397783A38D0.jpeg',\n",
              "  'images/D7AF463C-2369-492D-908D-BE1911CCD74C.jpeg',\n",
              "  'images/all14238-fig-0001-m-b.jpg',\n",
              "  'images/all14238-fig-0001-m-c.jpg',\n",
              "  'images/all14238-fig-0002-m-d.jpg',\n",
              "  'images/all14238-fig-0002-m-e.jpg',\n",
              "  'images/all14238-fig-0002-m-f.jpg',\n",
              "  'images/ciaa199.pdf-001-a.png',\n",
              "  'images/ciaa199.pdf-001-b.png',\n",
              "  'images/ciaa199.pdf-001-c.png',\n",
              "  'images/kjr-21-e24-g001-l-a.jpg',\n",
              "  'images/kjr-21-e24-g002-l-a.jpg',\n",
              "  'images/1-s2.0-S1684118220300682-main.pdf-002-a1.png',\n",
              "  'images/1-s2.0-S1684118220300682-main.pdf-002-a2.png',\n",
              "  'images/1-s2.0-S1684118220300682-main.pdf-003-b1.png',\n",
              "  'images/1-s2.0-S1684118220300682-main.pdf-003-b2.png',\n",
              "  'images/gr1_lrg-a.jpg',\n",
              "  'images/gr1_lrg-b.jpg',\n",
              "  'images/171CB377-62FF-4B76-906C-F3787A01CB2E.jpeg',\n",
              "  'images/5931B64A-7B97-485D-BE60-3F1EA76BC4F0.jpeg',\n",
              "  'images/C6EA0BE5-B01E-4113-B194-18D956675E25.jpeg',\n",
              "  'images/7EF28E12-F628-4BEC-A8C5-E6277C2E4F60.png',\n",
              "  'images/5e6dd879fde9502400e58b2f.jpeg',\n",
              "  'images/covid-19-pneumonia-19.jpg',\n",
              "  'images/03BF7561-A9BA-4C3C-B8A0-D3E585F73F3C.jpeg',\n",
              "  'images/figure1-5e73d7ae897e27ff066a30cb-98.jpeg',\n",
              "  'images/figure1-5e71be566aa8714a04de3386-98-left.jpeg',\n",
              "  'images/5A78BCA9-5B7A-440D-8A4E-AE7710EA6EAD.jpeg',\n",
              "  'images/2B8649B2-00C4-4233-85D5-1CE240CF233B.jpeg',\n",
              "  'images/2966893D-5DDF-4B68-9E2B-4979D5956C8E.jpeg',\n",
              "  'images/covid-19-pneumonia-30-PA.jpg',\n",
              "  'images/6b44464d-73a7-4cf3-bbb6-ffe7168300e3.annot.original.jpeg',\n",
              "  'images/58cb9263f16e94305c730685358e4e_jumbo.jpeg',\n",
              "  'images/9fdd3c3032296fd04d2cad5d9070d4_jumbo.jpeg',\n",
              "  'images/covid-19-infection-exclusive-gastrointestinal-symptoms-pa.png',\n",
              "  'images/covid-19-infection-exclusive-gastrointestinal-symptoms-l.png',\n",
              "  'images/covid-19-pneumonia-28.png',\n",
              "  'images/1.CXRCTThoraximagesofCOVID-19fromSingapore.pdf-001-fig2a.png',\n",
              "  'images/1.CXRCTThoraximagesofCOVID-19fromSingapore.pdf-001-fig2b.png',\n",
              "  'images/1.CXRCTThoraximagesofCOVID-19fromSingapore.pdf-002-fig3a.png',\n",
              "  'images/1.CXRCTThoraximagesofCOVID-19fromSingapore.pdf-002-fig3b.png',\n",
              "  'images/covid-19-pneumonia-evolution-over-a-week-1-day0-PA.jpg',\n",
              "  'images/covid-19-pneumonia-evolution-over-a-week-1-day3.jpg',\n",
              "  'images/covid-19-pneumonia-evolution-over-a-week-1-day4.jpg',\n",
              "  'images/53EC07C9-5CC6-4BE4-9B6F-D7B0D72AAA7E.jpeg',\n",
              "  'images/covid-19-pneumonia-23-day1.jpg',\n",
              "  'images/covid-19-pneumonia-23-day3.jpg',\n",
              "  'images/covid-19-pneumonia-23-day9.jpg',\n",
              "  'images/A7E260CE-8A00-4C5F-A7F5-27336527A981.jpeg',\n",
              "  'images/covid-19-caso-70-1-PA.jpg',\n",
              "  'images/covid-19-caso-70-2-APS.jpg',\n",
              "  'images/4e43e48d52c9e2d4c6c1fb9bc1544f_jumbo.jpeg',\n",
              "  'images/covid-19-pneumonia-53.jpg',\n",
              "  'images/post-intubuation-pneumomediastium-and-pneumothorax-background-covid-19-pneumonia-day1.jpg',\n",
              "  'images/post-intubuation-pneumomediastium-and-pneumothorax-background-covid-19-pneumonia-day6-1.jpg',\n",
              "  'images/post-intubuation-pneumomediastium-and-pneumothorax-background-covid-19-pneumonia-day6-2.jpg',\n",
              "  'images/post-intubuation-pneumomediastium-and-pneumothorax-background-covid-19-pneumonia-day7.jpg',\n",
              "  'images/covid-19-pneumonia-49-day4.jpg',\n",
              "  'images/covid-19-pneumonia-49-day8.jpg',\n",
              "  'images/covid-19-pneumonia-43-day0.jpeg',\n",
              "  'images/covid-19-pneumonia-43-day2.jpeg',\n",
              "  'images/covid-19-pneumonia-40.jpg',\n",
              "  'images/covid-19-pneumonia-42.jpeg',\n",
              "  'images/da9e9aac-de8c-44c7-ba57-e7cc8e4caaba.annot.original.jpeg',\n",
              "  'images/4ad30bc6-2da0-4f84-bc9b-62acabfd518a.annot.original.png',\n",
              "  'images/fff49165-b22d-4bb4-b9d1-d5d62c52436c.annot.original.png',\n",
              "  'images/figure1-5e7c1b8d98c29ab001275405-98.jpeg',\n",
              "  'images/figure1-5e7c1b8d98c29ab001275405-98-later.jpeg',\n",
              "  'images/radiol.2020201160.fig2a.jpeg',\n",
              "  'images/radiol.2020201160.fig2c.jpeg',\n",
              "  'images/radiol.2020201160.fig2d.jpeg',\n",
              "  'images/radiol.2020201160.fig3a.jpeg',\n",
              "  'images/radiol.2020201160.fig3b.jpeg',\n",
              "  'images/radiol.2020201160.fig3c.jpeg',\n",
              "  'images/radiol.2020201160.fig3d.jpeg',\n",
              "  'images/radiol.2020201160.fig6b.jpeg'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g8rt09f8vjTU",
        "colab_type": "text"
      },
      "source": [
        "### Prepare combined dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vSG5xslOiEoL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "dfd25271-fd02-4450-bc18-2a6f7efb4d7a"
      },
      "source": [
        "\"\"\"\n",
        "Script to prepare combined dataset\n",
        "Class 0: Normal\n",
        "Class 1: Bacterial Pneumonia\n",
        "Class 2: Viral Pneumonia\n",
        "Class 3: COVID-19\n",
        "\"\"\"\n",
        "import glob\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "\"\"\"parser = argparse.ArgumentParser()\n",
        "parser.add_argument(\"--combine_pneumonia\", action='store_true', default=False)\n",
        "args = parser.parse_args()\n",
        "\"\"\"\n",
        "\n",
        "COVID19_DATA_PATH = \"./CovidAID/data/covid19\"\n",
        "COVID19_IMGS_PATH = \"./CovidAID/covid-chestxray-dataset\"\n",
        "PNEUMONIDA_DATA_PATH = \"./CovidAID/chest-xray-pneumonia\"\n",
        "DATA_PATH = \"./CovidAID/data\"\n",
        "\n",
        "# Assert that the data directories are present\n",
        "for d in [COVID19_DATA_PATH, COVID19_IMGS_PATH, PNEUMONIDA_DATA_PATH, DATA_PATH]:\n",
        "    try:\n",
        "        assert os.path.isdir(d) \n",
        "    except:\n",
        "        print (\"Directory %s does not exists\" % d)\n",
        "\n",
        "def create_list (split):\n",
        "    assert split in ['train', 'test', 'val']\n",
        "\n",
        "    l = []\n",
        "    # Prepare list using kaggle pneumonia dataset\n",
        "    for f in glob.glob(os.path.join(PNEUMONIDA_DATA_PATH, split, 'NORMAL', '*')):\n",
        "        l.append((f, 0)) # Class 0\n",
        "\n",
        "    for f in glob.glob(os.path.join(PNEUMONIDA_DATA_PATH, split, 'PNEUMONIA', '*')):\n",
        "       \"\"\" if args.combine_pneumonia:\n",
        "            l.append((f, 1)) # Class 1\n",
        "        else:\n",
        "            if 'bacteria' in f:\n",
        "                l.append((f, 1)) # Class 1\n",
        "            else:\n",
        "                l.append((f, 2)) # Class 2\"\"\"\n",
        "\n",
        "    # Prepare list using covid dataset\n",
        "    covid_file = os.path.join(COVID19_DATA_PATH, '%s_list.txt'%split)\n",
        "    with open(covid_file, 'r') as cf:\n",
        "        for f in cf.readlines():\n",
        "            f = os.path.join(COVID19_IMGS_PATH, f.strip())\n",
        "            \"\"\"if args.combine_pneumonia:\n",
        "                l.append((f, 2)) # Class 2\n",
        "            else:\n",
        "                l.append((f, 3)) # Class 3\"\"\"\n",
        "\n",
        "    with open(os.path.join(DATA_PATH, '%s.txt'%split), 'w') as f:\n",
        "        for item in l:\n",
        "            f.write(\"%s %d\\n\" % item)\n",
        "\n",
        "for split in ['train', 'test', 'val']:\n",
        "    create_list(split)"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Directory ./CovidAID/chest-xray-pneumonia does not exists\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yr-dLuZZzkk2",
        "colab_type": "text"
      },
      "source": [
        "### covidaid"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oPIzPtRazmQx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "The main CovidAID and CheXNet implementation\n",
        "\"\"\"\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "\n",
        "class DenseNet121(nn.Module):\n",
        "    \"\"\"Model modified.\n",
        "\n",
        "    The architecture of our model is the same as standard DenseNet121\n",
        "    except the classifier layer which has an additional sigmoid function.\n",
        "\n",
        "    \"\"\"\n",
        "    def __init__(self, out_size):\n",
        "        super(DenseNet121, self).__init__()\n",
        "        self.densenet121 = torchvision.models.densenet121(pretrained=True)\n",
        "        num_ftrs = self.densenet121.classifier.in_features\n",
        "        self.densenet121.classifier = nn.Sequential(\n",
        "            nn.Linear(num_ftrs, out_size),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.densenet121(x)\n",
        "        return x\n",
        "    \n",
        "\n",
        "class CovidAID(DenseNet121):\n",
        "    \"\"\"\n",
        "    Modified DenseNet network with 4 classes\n",
        "    \"\"\"\n",
        "    def __init__(self, combine_pneumonia=False):\n",
        "        NUM_CLASSES = 3 if combine_pneumonia else 4\n",
        "        super(CovidAID, self).__init__(NUM_CLASSES)\n",
        "\n",
        "\n",
        "class CheXNet(DenseNet121):\n",
        "    \"\"\"\n",
        "    Modified DenseNet network with 14 classes\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super(CheXNet, self).__init__(14)\n"
      ],
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VOXYmprP19ph",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!cp /content/CovidAID/tools/covidaid.py /content/covidaid.py"
      ],
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SdqByWKnqW5p",
        "colab_type": "text"
      },
      "source": [
        "### Prepare pretrained layer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PZSUi_nVqbdh",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 350
        },
        "outputId": "154486b1-cb65-40ec-cd45-79624397c5a2"
      },
      "source": [
        "\"\"\"\n",
        "Code to transfer weights from CheXNet (torch 0.3) to CovidAID\n",
        "\"\"\"\n",
        "\n",
        "import sys\n",
        "from covidaid import CovidAID, CheXNet\n",
        "import torch\n",
        "import argparse\n",
        "\n",
        "'''parser = argparse.ArgumentParser()\n",
        "parser.add_argument(\"--combine_pneumonia\", action='store_true', default=False)\n",
        "parser.add_argument(\"--chexnet_model_checkpoint\", \"--old\", type=str, default=\"./CovidAID/data/CheXNet_model.pth.tar\")\n",
        "parser.add_argument(\"--covidaid_model_trained_checkpoint\", \"--new\", type=str, default=\"./CovidAID/models/CovidAID_transfered.pth.tar\")\n",
        "args = parser.parse_args()'''\n",
        "\n",
        "'''chexnet_model_checkpoint = args.chexnet_model_checkpoint\n",
        "covidaid_model_trained_checkpoint = args.covidaid_model_trained_checkpoint\n",
        "\n",
        "model = CovidAID(combine_pneumonia=args.combine_pneumonia)'''\n",
        "chexnet_model_checkpoint = '/content/CovidAID/data/CheXNet_model.pth.tar'\n",
        "covidaid_model_trained_checkpoint = '/content/CovidAID/models/CovidAID_4_class.pth'\n",
        "\n",
        "model ='/content/CovidAID/models/CovidAID_4_class.pth'\n",
        "\n",
        "def load_weights(checkpoint_pth, state_dict=True):\n",
        "    model = torch.load(checkpoint_pth)\n",
        "    \n",
        "    if state_dict:\n",
        "        return model['state_dict']\n",
        "    else:\n",
        "        return model\n",
        "\n",
        "def get_top_keys(model, depth=0):\n",
        "    return set({w.split('.')[depth] for w in model.keys()})\n",
        "\n",
        "chexnet_model = load_weights(chexnet_model_checkpoint)\n",
        "template =  load_weights(covidaid_model_trained_checkpoint)\n",
        "\n",
        "assert get_top_keys(chexnet_model, depth=2) == set({'features', 'classifier'})\n",
        "assert get_top_keys(template, depth=1) == set({'features', 'classifier'})\n",
        "\n",
        "# print (chexnet_model.keys())\n",
        "# print (template.keys())\n",
        "# print (model.state_dict().keys())\n",
        "\n",
        "c_keys = {k for k in chexnet_model.keys()}\n",
        "t_keys = {'module.' + k for k in template.keys()}\n",
        "\n",
        "assert len(c_keys.difference(t_keys)) == 0\n",
        "assert len(t_keys.difference(c_keys)) == 0\n",
        "\n",
        "\n",
        "# Transfer the feature weights\n",
        "for k, w in template.items():\n",
        "    chex_key = 'module.' + k\n",
        "\n",
        "    if k.split('.')[1] == 'classifier':\n",
        "        # Uncomment below to copy trained weights of pneumonia\n",
        "        # 6th class is pneumonia in CheXNet => Copy it's weights to pneumonia classes\n",
        "        # for c in [1, 2, 3]:\n",
        "        #     template[k][c, ...] = chexnet_model[chex_key][6, ...]\n",
        "        print ('doing nothing for', k)\n",
        "    else:\n",
        "        # print (type(template[k]), template[k].size())\n",
        "        # print (type(chexnet_model[chex_key]), chexnet_model[chex_key].size())\n",
        "        assert chexnet_model[chex_key].size() == template[k].size()\n",
        "        template[k] = chexnet_model[chex_key]\n",
        "\n",
        "torch.save(template, covidaid_model_trained_checkpoint)\n"
      ],
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-46-c83b70dd2444>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     35\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0mchexnet_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_weights\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchexnet_model_checkpoint\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 37\u001b[0;31m \u001b[0mtemplate\u001b[0m \u001b[0;34m=\u001b[0m  \u001b[0mload_weights\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcovidaid_model_trained_checkpoint\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     38\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[0;32massert\u001b[0m \u001b[0mget_top_keys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchexnet_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdepth\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m'features'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'classifier'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-46-c83b70dd2444>\u001b[0m in \u001b[0;36mload_weights\u001b[0;34m(checkpoint_pth, state_dict)\u001b[0m\n\u001b[1;32m     27\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mstate_dict\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 29\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state_dict'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     30\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyError\u001b[0m: 'state_dict'"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JbXd7BsBq6KJ",
        "colab_type": "text"
      },
      "source": [
        "### training the classifier layer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "knr2eKizq-QY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Trainer for training and testing the networks\n",
        "\"\"\"\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import datetime\n",
        "import json\n",
        "import argparse\n",
        "import glob\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.backends.cudnn as cudnn\n",
        "from torch import optim\n",
        "from torch.utils.data.distributed import DistributedSampler\n",
        "from torch.utils.data import DataLoader\n",
        "import torchvision.transforms as transforms\n",
        "from read_data import ChestXrayDataSet, ChestXrayDataSetTest\n",
        "from sklearn.metrics import roc_auc_score, confusion_matrix, roc_curve, auc, f1_score\n",
        "import seaborn as sn\n",
        "import pandas as pd\n",
        "from scipy import interp\n",
        "from itertools import cycle\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from covidaid import CovidAID\n",
        "from tqdm import tqdm\n",
        "\n",
        "class Trainer:\n",
        "    def __init__ (self, local_rank=None, checkpoint=None, combine_pneumonia=False):\n",
        "        \"\"\"\n",
        "        Trainer for the CovidAID\n",
        "        \"\"\"\n",
        "        self.distributed = True if local_rank is not None else False\n",
        "        print (\"Distributed training %s\" % ('ON' if self.distributed else 'OFF'))\n",
        "        if self.distributed:\n",
        "            raise NotImplementedError(\"Currently distributed training not supported\")\n",
        "            self.device = torch.cuda.device('cuda', local_rank)\n",
        "        else:\n",
        "            self.device = torch.cuda.device('cuda')\n",
        "\n",
        "        # Using 2 classes for pneumonia vs 1 class\n",
        "        self.combine_pneumonia = combine_pneumonia\n",
        "\n",
        "        # self.net = CovidAID().to(self.device)\n",
        "        self.net = CovidAID(combine_pneumonia).cuda()\n",
        "        if self.distributed:\n",
        "            self.net = torch.nn.parallel.DistributedDataParallel(self.net,\n",
        "                                                            device_ids=[local_rank],\n",
        "                                                            output_device=local_rank)\n",
        "        \n",
        "        # load model\n",
        "        if checkpoint is not None:\n",
        "            self.load_model(checkpoint)\n",
        "\n",
        "    def train(self, TRAIN_IMAGE_LIST, VAL_IMAGE_LIST, NUM_EPOCHS=10, LR=0.001, BATCH_SIZE=64,\n",
        "                start_epoch=0, logging=True, save_path=None, freeze_feature_layers=True):\n",
        "        \"\"\"\n",
        "        Train the CovidAID\n",
        "        \"\"\"\n",
        "        normalize = transforms.Normalize([0.485, 0.456, 0.406],\n",
        "                                     [0.229, 0.224, 0.225])\n",
        "\n",
        "        train_dataset = ChestXrayDataSet(image_list_file=TRAIN_IMAGE_LIST,\n",
        "                                        transform=transforms.Compose([\n",
        "                                            transforms.Resize(256),\n",
        "                                            transforms.TenCrop(224),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([normalize(crop) for crop in crops]))\n",
        "                                        ]),\n",
        "                                        combine_pneumonia=self.combine_pneumonia)\n",
        "        if self.distributed:\n",
        "            sampler = DistributedSampler(train_dataset)\n",
        "            train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE,\n",
        "                                    shuffle=False, num_workers=8, pin_memory=True,\n",
        "                                    sampler=sampler)\n",
        "        else:\n",
        "            train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE,\n",
        "                                    shuffle=True, num_workers=8, pin_memory=True)\n",
        "\n",
        "        val_dataset = ChestXrayDataSet(image_list_file=VAL_IMAGE_LIST,\n",
        "                                        transform=transforms.Compose([\n",
        "                                            transforms.Resize(256),\n",
        "                                            transforms.TenCrop(224),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([normalize(crop) for crop in crops]))\n",
        "                                        ]),\n",
        "                                        combine_pneumonia=self.combine_pneumonia)\n",
        "        if self.distributed:\n",
        "            sampler = DistributedSampler(val_dataset)\n",
        "            val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE,\n",
        "                                    shuffle=False, num_workers=8, pin_memory=True,\n",
        "                                    sampler=sampler)\n",
        "        else:\n",
        "            val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE,\n",
        "                                    shuffle=True, num_workers=8, pin_memory=True)\n",
        "\n",
        "        # Freeze heads and create optimizer\n",
        "        if freeze_feature_layers:\n",
        "            print (\"Freezing feature layers\")\n",
        "            for param in self.net.densenet121.features.parameters():\n",
        "                param.requires_grad = False\n",
        "\n",
        "        # optimizer = optim.SGD(filter(lambda p: p.requires_grad, self.net.parameters()),\n",
        "        #                 lr=LR, momentum=0.9)\n",
        "        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.net.parameters()), lr=LR)\n",
        "\n",
        "\n",
        "        for epoch in range(start_epoch, NUM_EPOCHS):\n",
        "            # switch to train mode\n",
        "            self.net.train()\n",
        "            tot_loss = 0.0\n",
        "            for i, (inputs, target) in tqdm(enumerate(train_loader), total=len(train_dataset)/BATCH_SIZE):\n",
        "                # inputs = inputs.to(self.device)\n",
        "                # target = target.to(self.device)\n",
        "                inputs = inputs.cuda()\n",
        "                target = target.cuda()\n",
        "\n",
        "                # Shape of input == [BATCH_SIZE, NUM_CROPS=19, CHANNELS=3, HEIGHT=224, WIDTH=244]\n",
        "                bs, n_crops, c, h, w = inputs.size()\n",
        "                inputs = inputs.view(-1, c, h, w)\n",
        "                inputs = torch.autograd.Variable(inputs.view(-1, c, h, w))\n",
        "                target = torch.autograd.Variable(target)\n",
        "                preds = self.net(inputs).view(bs, n_crops, -1).mean(dim=1)\n",
        "\n",
        "                # loss = torch.sum(torch.abs(preds - target) ** 2)    \n",
        "                loss = train_dataset.loss(preds, target)  \n",
        "                # exit()          \n",
        "                tot_loss += float(loss.data)\n",
        "\n",
        "                optimizer.zero_grad()\n",
        "                loss.backward()\n",
        "                optimizer.step()\n",
        "\n",
        "            tot_loss /= len(train_dataset)\n",
        "\n",
        "            # Clear cache\n",
        "            torch.cuda.empty_cache()\n",
        "            \n",
        "            # Running on validation set\n",
        "            self.net.eval()\n",
        "            val_loss = 0.0\n",
        "            for i, (inputs, target) in tqdm(enumerate(val_loader), total=len(val_dataset)/BATCH_SIZE):\n",
        "                # inputs = inputs.to(self.device)\n",
        "                # target = target.to(self.device)\n",
        "                inputs = inputs.cuda()\n",
        "                target = target.cuda()\n",
        "\n",
        "                # Shape of input == [BATCH_SIZE, NUM_CROPS=19, CHANNELS=3, HEIGHT=224, WIDTH=244]\n",
        "                bs, n_crops, c, h, w = inputs.size()\n",
        "                inputs = inputs.view(-1, c, h, w)\n",
        "                inputs = torch.autograd.Variable(inputs.view(-1, c, h, w), volatile=True)\n",
        "                target = torch.autograd.Variable(target, volatile=True)\n",
        "\n",
        "                preds = self.net(inputs).view(bs, n_crops, -1).mean(1)\n",
        "                # loss = torch.sum(torch.abs(preds - target) ** 2)\n",
        "                loss = val_dataset.loss(preds, target) \n",
        "                \n",
        "                val_loss += float(loss.data)\n",
        "\n",
        "            val_loss /= len(val_dataset)\n",
        "\n",
        "            # Clear cache\n",
        "            torch.cuda.empty_cache()\n",
        "\n",
        "            # logging statistics\n",
        "            timestamp = str(datetime.datetime.now()).split('.')[0]\n",
        "            log = json.dumps({\n",
        "                'timestamp': timestamp,\n",
        "                'epoch': epoch+1,\n",
        "                'train_loss': float('%.5f' % tot_loss),\n",
        "                'val_loss': float('%.5f' % val_loss),\n",
        "                'lr': float('%.6f' % LR)\n",
        "            })\n",
        "            if logging:\n",
        "                print (log)\n",
        "\n",
        "            log_file = os.path.join(save_path, 'train.log')\n",
        "            if log_file is not None:\n",
        "                with open(log_file, 'a') as f:\n",
        "                    f.write(\"{}\\n\".format(log))\n",
        "\n",
        "            model_path = os.path.join(save_path, 'epoch_%d.pth'%(epoch+1))\n",
        "            self.save_model(model_path)\n",
        "            \n",
        "        print ('Finished Training')\n",
        "\n",
        "    def predict(self, TEST_IMAGE_LIST, BATCH_SIZE=64):\n",
        "        \"\"\"\n",
        "        Predict the task labels corresponding to the input images\n",
        "        \"\"\"\n",
        "        normalize = transforms.Normalize([0.485, 0.456, 0.406],\n",
        "                                     [0.229, 0.224, 0.225])\n",
        "\n",
        "        test_dataset = ChestXrayDataSetTest(image_list_file=TEST_IMAGE_LIST,\n",
        "                                        transform=transforms.Compose([\n",
        "                                            transforms.Resize(256),\n",
        "                                            transforms.TenCrop(224),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),\n",
        "                                            transforms.Lambda\n",
        "                                            (lambda crops: torch.stack([normalize(crop) for crop in crops]))\n",
        "                                        ]),\n",
        "                                        combine_pneumonia=self.combine_pneumonia)\n",
        "        if self.distributed:\n",
        "            sampler = DistributedSampler(test_dataset)\n",
        "            test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE,\n",
        "                                shuffle=False, num_workers=8, pin_memory=True,\n",
        "                                sampler=sampler)\n",
        "        else:\n",
        "            test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE,\n",
        "                                shuffle=True, num_workers=8, pin_memory=True)\n",
        "\n",
        "        # initialize the ground truth and output tensor\n",
        "        gt = torch.FloatTensor().cuda()\n",
        "        pred = torch.FloatTensor().cuda()\n",
        "\n",
        "        # switch to evaluate mode\n",
        "        self.net.eval()\n",
        "\n",
        "        for i, (inputs, target) in tqdm(enumerate(test_loader), total=len(test_loader)):\n",
        "            # inputs = inputs.to(self.device)\n",
        "            # target = target.to(self.device)\n",
        "            inputs = inputs.cuda()\n",
        "            target = target.cuda()\n",
        "            gt = torch.cat((gt, target), 0)\n",
        "\n",
        "            # Shape of input == [BATCH_SIZE, NUM_CROPS=19, CHANNELS=3, HEIGHT=224, WIDTH=244]\n",
        "            bs, n_crops, c, h, w = inputs.size()\n",
        "            inputs = torch.autograd.Variable(inputs.view(-1, c, h, w), volatile=True)\n",
        "\n",
        "            # Pass through the network and take average prediction from all the crops\n",
        "            output = self.net(inputs)\n",
        "            output_mean = output.view(bs, n_crops, -1).mean(1)\n",
        "            pred = torch.cat((pred, output_mean.data), 0)\n",
        "        gt = gt.cpu().numpy()\n",
        "        pred = pred.cpu().numpy()\n",
        "\n",
        "        return gt, pred\n",
        "\n",
        "    def evaluate(self, TEST_IMAGE_LIST, BATCH_SIZE=64, cm_path='cm', roc_path='roc'):\n",
        "        \"\"\"\n",
        "        Evaluate on the test set plotting confusion matrix and roc curves\n",
        "        \"\"\"\n",
        "        gt, pred = self.predict(TEST_IMAGE_LIST, BATCH_SIZE=64)\n",
        "        print (pred)\n",
        "\n",
        "        # Compute ROC scores\n",
        "        labels = ['Normal', 'Bacterial', 'Viral', 'COVID-19']\n",
        "        if self.combine_pneumonia:\n",
        "            labels = ['Normal', 'Pneumonia', 'COVID-19']\n",
        "        self.compute_AUC_scores(gt, pred, labels)\n",
        "\n",
        "        # Plot ROC scores\n",
        "        self.plot_ROC_curve(gt, pred, labels, roc_path)\n",
        "\n",
        "        # Treat the max. output as prediction. \n",
        "        # Plot Confusion Matrix\n",
        "        gt = gt.argmax(axis=1)\n",
        "        pred = pred.argmax(axis=1)\n",
        "        self.plot_confusion_matrix(gt, pred, labels, cm_path)\n",
        "\n",
        "    def F1(self, TEST_DIR, out_file, BATCH_SIZE=64):\n",
        "        \"\"\"\n",
        "        Evaluate on multiple test sets and compute F1 scores\n",
        "        \"\"\"\n",
        "        f = open(out_file, 'w')\n",
        "        for test_file in glob.glob(os.path.join(TEST_DIR, '*.txt')):\n",
        "            print (test_file)\n",
        "            gt, pred = self.predict(test_file, BATCH_SIZE)\n",
        "            # Treat the max. output as prediction. \n",
        "            gt = gt.argmax(axis=1)\n",
        "            pred = pred.argmax(axis=1)\n",
        "            f1 = f1_score(gt, pred, average='macro')\n",
        "            f.write('%s %.6f\\n' % (test_file, f1))\n",
        "        f.close()\n",
        "\n",
        "    def plot_confusion_matrix(self, y_true, y_pred, labels, cm_path):\n",
        "        norm_cm = confusion_matrix(y_true, y_pred, normalize='true')\n",
        "        norm_df_cm = pd.DataFrame(norm_cm, index=labels, columns=labels)\n",
        "        plt.figure(figsize = (10,7))\n",
        "        sn.heatmap(norm_df_cm, annot=True, fmt='.2f', square=True, cmap=plt.cm.Blues)\n",
        "        plt.xlabel(\"Predicted\")\n",
        "        plt.ylabel(\"Ground Truth\")\n",
        "        matplotlib.rcParams.update({'font.size': 14})\n",
        "        plt.savefig('%s_norm.png' % cm_path, pad_inches = 0, bbox_inches='tight')\n",
        "        \n",
        "        cm = confusion_matrix(y_true, y_pred)\n",
        "        # Finding the annotations\n",
        "        cm = cm.tolist()\n",
        "        norm_cm = norm_cm.tolist()\n",
        "        annot = [\n",
        "            [(\"%d (%.2f)\" % (c, nc)) for c, nc in zip(r, nr)]\n",
        "            for r, nr in zip(cm, norm_cm)\n",
        "        ]\n",
        "        plt.figure(figsize = (10,7))\n",
        "        sn.heatmap(norm_df_cm, annot=annot, fmt='', cbar=False, square=True, cmap=plt.cm.Blues)\n",
        "        plt.xlabel(\"Predicted\")\n",
        "        plt.ylabel(\"Ground Truth\")\n",
        "        matplotlib.rcParams.update({'font.size': 14})\n",
        "        plt.savefig('%s.png' % cm_path, pad_inches = 0, bbox_inches='tight')\n",
        "        print (cm)\n",
        "\n",
        "        accuracy = np.sum(y_true == y_pred) / len(y_true)\n",
        "        print (\"Accuracy: %.5f\" % accuracy)\n",
        "\n",
        "    def compute_AUC_scores(self, y_true, y_pred, labels):\n",
        "        \"\"\"\n",
        "        Computes the Area Under the Curve (AUC) from prediction scores\n",
        "        y_true.shape  = [n_samples, n_classes]\n",
        "        y_preds.shape = [n_samples, n_classes]\n",
        "        labels.shape  = [n_classes]\n",
        "        \"\"\"\n",
        "        AUROC_avg = roc_auc_score(y_true, y_pred)\n",
        "        print('The average AUROC is {AUROC_avg:.4f}'.format(AUROC_avg=AUROC_avg))\n",
        "        for y, pred, label in zip(y_true.transpose(), y_pred.transpose(), labels):\n",
        "            print('The AUROC of {0:} is {1:.4f}'.format(label, roc_auc_score(y, pred)))\n",
        "\n",
        "    def plot_ROC_curve(self, y_true, y_pred, labels, roc_path): \n",
        "        \"\"\"\n",
        "        Plots the ROC curve from prediction scores\n",
        "        y_true.shape  = [n_samples, n_classes]\n",
        "        y_preds.shape = [n_samples, n_classes]\n",
        "        labels.shape  = [n_classes]\n",
        "        \"\"\"\n",
        "        n_classes = len(labels)\n",
        "        # Compute ROC curve and ROC area for each class\n",
        "        fpr = dict()\n",
        "        tpr = dict()\n",
        "        roc_auc = dict()\n",
        "        for y, pred, label in zip(y_true.transpose(), y_pred.transpose(), labels):\n",
        "            fpr[label], tpr[label], _ = roc_curve(y, pred)\n",
        "            roc_auc[label] = auc(fpr[label], tpr[label])\n",
        "\n",
        "        # First aggregate all false positive rates\n",
        "        all_fpr = np.unique(np.concatenate([fpr[label] for label in labels]))\n",
        "\n",
        "        # Then interpolate all ROC curves at this points\n",
        "        mean_tpr = np.zeros_like(all_fpr)\n",
        "        for label in labels:\n",
        "            mean_tpr += interp(all_fpr, fpr[label], tpr[label])\n",
        "\n",
        "        # Finally average it and compute AUC\n",
        "        mean_tpr /= n_classes\n",
        "\n",
        "        # Compute micro-average ROC curve and ROC area\n",
        "        fpr[\"micro\"], tpr[\"micro\"], _ = roc_curve(y_true.ravel(), y_pred.ravel())\n",
        "        roc_auc[\"micro\"] = auc(fpr[\"micro\"], tpr[\"micro\"])\n",
        "\n",
        "        fpr[\"macro\"] = all_fpr\n",
        "        tpr[\"macro\"] = mean_tpr\n",
        "        roc_auc[\"macro\"] = auc(fpr[\"macro\"], tpr[\"macro\"])\n",
        "\n",
        "        # Plot all ROC curves\n",
        "        plt.figure()\n",
        "        lw = 2\n",
        "        plt.plot(fpr[\"micro\"], tpr[\"micro\"],\n",
        "                label='micro-average ROC curve (area = {0:0.3f})'\n",
        "                    ''.format(roc_auc[\"micro\"]),\n",
        "                color='deeppink', linestyle=':', linewidth=2)\n",
        "\n",
        "        plt.plot(fpr[\"macro\"], tpr[\"macro\"],\n",
        "                label='macro-average ROC curve (area = {0:0.3f})'\n",
        "                    ''.format(roc_auc[\"macro\"]),\n",
        "                color='navy', linestyle=':', linewidth=2)\n",
        "\n",
        "        if len(labels) == 4:\n",
        "            colors = ['green', 'cornflowerblue', 'darkorange', 'darkred']\n",
        "        else:\n",
        "            colors = ['green', 'cornflowerblue', 'darkred']\n",
        "        for label, color in zip(labels, cycle(colors)):\n",
        "            plt.plot(fpr[label], tpr[label], color=color, lw=lw,\n",
        "                    label='ROC curve of {0} (area = {1:0.3f})'\n",
        "                    ''.format(label, roc_auc[label]))\n",
        "\n",
        "        plt.plot([0, 1], [0, 1], 'k--', lw=lw)\n",
        "        plt.xlim([0.0, 1.0])\n",
        "        plt.ylim([0.0, 1.05])\n",
        "        plt.xlabel('False Positive Rate')\n",
        "        plt.ylabel('True Positive Rate')\n",
        "        plt.title('ROC curve')\n",
        "        plt.legend(loc=\"lower right\")\n",
        "        matplotlib.rcParams.update({'font.size': 14})\n",
        "        plt.savefig('%s.png' % roc_path, pad_inches = 0, bbox_inches='tight')\n",
        "\n",
        "    def save_model(self, checkpoint_path, model=None):\n",
        "        if model is None: model = self.net\n",
        "        torch.save(model.state_dict(), checkpoint_path)\n",
        "    \n",
        "    def load_model(self, checkpoint_path, model=None):\n",
        "        print (\"Loading model from %s\" % checkpoint_path)\n",
        "        if model is None: model = self.net\n",
        "        model.load_state_dict(torch.load(checkpoint_path))\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument(\"--local_rank\", type=int) # For distributed processing\n",
        "    parser.add_argument(\"--mode\", choices=['train', 'test', 'f1'], required=True)\n",
        "    parser.add_argument(\"--checkpoint\", type=str, required=True)\n",
        "    parser.add_argument(\"--combine_pneumonia\", action='store_true', default=False)\n",
        "    parser.add_argument(\"--save\", type=str)\n",
        "    parser.add_argument(\"--start\", type=int, default=0)\n",
        "    parser.add_argument(\"--lr\", type=float, default=1e-4)\n",
        "    parser.add_argument(\"--freeze\", action='store_true', default=False)\n",
        "    parser.add_argument(\"--bs\", type=int, default=8)\n",
        "    parser.add_argument(\"--cm_path\", type=str, default='plots/cm')\n",
        "    parser.add_argument(\"--roc_path\", type=str, default='plots/roc')\n",
        "\n",
        "    # parser.add_argment(\"--torch_version\", \"--tv\", choices=[\"0.3\", \"new\"], default=\"0.3\")\n",
        "    args = parser.parse_args()\n",
        "\n",
        "    TRAIN_IMAGE_LIST = './data/train.txt'\n",
        "    VAL_IMAGE_LIST = './data/val.txt'\n",
        "    TEST_IMAGE_LIST = './data/test.txt'\n",
        "    TEST_DIR = './data/samples'\n",
        "\n",
        "    # Initializes the distributed backend which will take care of sychronizing nodes/GPUs\n",
        "    if args.local_rank is not None:\n",
        "        torch.distributed.init_process_group(backend='nccl')\n",
        "\n",
        "    trainer = Trainer(local_rank=args.local_rank, checkpoint=args.checkpoint, combine_pneumonia=args.combine_pneumonia)\n",
        "    if args.mode == 'test':\n",
        "        trainer.evaluate(TEST_IMAGE_LIST, cm_path=args.cm_path, roc_path=args.roc_path)\n",
        "    elif args.mode == 'train':\n",
        "        assert args.save is not None\n",
        "        trainer.train(TRAIN_IMAGE_LIST, VAL_IMAGE_LIST, BATCH_SIZE=args.bs, NUM_EPOCHS=300, LR=args.lr,\n",
        "                        start_epoch=args.start, save_path=args.save, freeze_feature_layers=args.freeze)\n",
        "    else:\n",
        "        trainer.F1(TEST_DIR, 'models/samples.txt')\n",
        "\n",
        "# Run command for distributed\n",
        "# python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=\"192.168.1.1\" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)\n",
        "# python -m torch.distributed.launch --nproc_per_node=2 trainer.py --mode train"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p8pSoZhZrHWJ",
        "colab_type": "text"
      },
      "source": [
        "### inference layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e7piK_mwrS4R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Script to process set of images and output predictions\n",
        "\"\"\"\n",
        "from RISE.visualize import visualize\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "import torchvision.transforms as transforms\n",
        "from PIL import Image\n",
        "import os\n",
        "import glob\n",
        "import argparse\n",
        "from covidaid import CovidAID\n",
        "from tqdm import tqdm\n",
        "import termtables as tt\n",
        "\n",
        "\n",
        "\n",
        "class CovidDataLoader(Dataset):\n",
        "    \"\"\"\n",
        "    Read images and corresponding labels.\n",
        "    \"\"\"\n",
        "    def __init__(self, image_dir, transform=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            image_dir: path to image directory.\n",
        "            transform: optional transform to be applied on a sample.\n",
        "        \"\"\"\n",
        "        self.image_names = [img for img in glob.glob(os.path.join(image_dir, '*'))]\n",
        "        self.transform = transform\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            index: the index of item\n",
        "        Returns:\n",
        "            image and its name\n",
        "        \"\"\"\n",
        "        image_name = self.image_names[index]\n",
        "        image = Image.open(image_name).convert('RGB')\n",
        "        if self.transform is not None:\n",
        "            image = self.transform(image)\n",
        "        return image, image_name.split('/')[-1]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.image_names)\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument(\"--img_dir\", type=str, required=True)\n",
        "    parser.add_argument(\"--checkpoint\", type=str, required=True)\n",
        "    parser.add_argument(\"--combine_pneumonia\", action='store_true', default=False)\n",
        "    parser.add_argument(\"--visualize_dir\", type=str, default=None)\n",
        "    args = parser.parse_args()\n",
        "\n",
        "    # Load the model\n",
        "    model = CovidAID(args.combine_pneumonia).cuda()\n",
        "    model.load_state_dict(torch.load(args.checkpoint))\n",
        "\n",
        "    # Load the data\n",
        "    normalize = transforms.Normalize([0.485, 0.456, 0.406],\n",
        "                                     [0.229, 0.224, 0.225])\n",
        "    \n",
        "    test_dataset = CovidDataLoader(image_dir=args.img_dir,\n",
        "            transform=transforms.Compose([\n",
        "                transforms.Resize(256),\n",
        "                transforms.TenCrop(224),\n",
        "                transforms.Lambda\n",
        "                (lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),\n",
        "                transforms.Lambda\n",
        "                (lambda crops: torch.stack([normalize(crop) for crop in crops]))\n",
        "            ])\n",
        "    )\n",
        "    test_loader = DataLoader(dataset=test_dataset, batch_size=64,\n",
        "                    shuffle=False, num_workers=8, pin_memory=True)\n",
        "\n",
        "    # initialize the output tensor\n",
        "    pred = torch.FloatTensor().cuda()\n",
        "    pred_names = []\n",
        "\n",
        "    # switch to evaluate mode\n",
        "    model.eval()\n",
        "\n",
        "    for i, (inputs, names) in tqdm(enumerate(test_loader), total=len(test_loader)):\n",
        "        inputs = inputs.cuda()\n",
        "\n",
        "        # Shape of input == [BATCH_SIZE, NUM_CROPS=10, CHANNELS=3, HEIGHT=224, WIDTH=244]\n",
        "        bs, n_crops, c, h, w = inputs.size()\n",
        "        inputs = torch.autograd.Variable(inputs.view(-1, c, h, w), volatile=True)\n",
        "\n",
        "        # Pass through the network and take average prediction from all the crops\n",
        "        output = model(inputs)\n",
        "        output_mean = output.view(bs, n_crops, -1).mean(1)\n",
        "        pred = torch.cat((pred, output_mean.data), 0)\n",
        "        pred_names += names\n",
        "\n",
        "    pred = pred.cpu().numpy()\n",
        "\n",
        "    assert len(pred) == len(pred_names)\n",
        "\n",
        "    scores = []\n",
        "    for p, n in zip(pred, pred_names):\n",
        "        p = [\"%.1f %%\" % (i * 100) for i in p]\n",
        "        scores.append([n] + p)\n",
        "\n",
        "    header=['Name', 'Normal', 'Bacterial', 'Viral', 'COVID-19']\n",
        "    alignment=\"c\"*5\n",
        "    if args.combine_pneumonia:\n",
        "        header = ['Name', 'Normal', 'Pneumonia', 'COVID-19']\n",
        "        alignment = \"c\"*4\n",
        "\n",
        "    string = tt.to_string(\n",
        "        scores,\n",
        "        header=header,\n",
        "        style=tt.styles.ascii_thin_double,\n",
        "        padding=(0, 1),\n",
        "        alignment=alignment\n",
        "    )\n",
        "\n",
        "    print (string)\n",
        "    \n",
        "    # RISE Visualization\n",
        "    if args.visualize_dir:\n",
        "        visualize(model,args.img_dir,args.visualize_dir,CovidDataLoader)\n",
        "        print(\"Visualizations generated at \"+str(args.visualize_dir))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}