create_COVIDx_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import random \n",
    "from shutil import copyfile\n",
    "import pydicom as dicom\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set parameters here\n",
    "savepath = 'data'\n",
    "seed = 0\n",
    "np.random.seed(seed) # Reset the seed so all runs are the same.\n",
    "random.seed(seed)\n",
    "MAXVAL = 255  # Range [0 255]\n",
    "\n",
    "# path to covid-19 dataset from https://github.com/ieee8023/covid-chestxray-dataset\n",
    "cohen_imgpath = '../covid-chestxray-dataset/images' \n",
    "cohen_csvpath = '../covid-chestxray-dataset/metadata.csv'\n",
    "\n",
    "# path to covid-14 dataset from https://github.com/agchung/Figure1-COVID-chestxray-dataset\n",
    "fig1_imgpath = '../Figure1-COVID-chestxray-dataset/images'\n",
    "fig1_csvpath = '../Figure1-COVID-chestxray-dataset/metadata.csv'\n",
    "\n",
    "# path to https://www.kaggle.com/c/rsna-pneumonia-detection-challenge\n",
    "rsna_datapath = '../rsna-pneumonia-detection-challenge'\n",
    "# get all the normal from here\n",
    "rsna_csvname = 'stage_2_detailed_class_info.csv' \n",
    "# get all the 1s from here since 1 indicate pneumonia\n",
    "# found that images that aren't pneunmonia and also not normal are classified as 0s\n",
    "rsna_csvname2 = 'stage_2_train_labels.csv' \n",
    "rsna_imgpath = 'stage_2_train_images'\n",
    "\n",
    "# parameters for COVIDx dataset\n",
    "train = []\n",
    "test = []\n",
    "test_count = {'normal': 0, 'pneumonia': 0, 'COVID-19': 0}\n",
    "train_count = {'normal': 0, 'pneumonia': 0, 'COVID-19': 0}\n",
    "\n",
    "mapping = dict()\n",
    "mapping['COVID-19'] = 'COVID-19'\n",
    "mapping['SARS'] = 'pneumonia'\n",
    "mapping['MERS'] = 'pneumonia'\n",
    "mapping['Streptococcus'] = 'pneumonia'\n",
    "mapping['Klebsiella'] = 'pneumonia'\n",
    "mapping['Chlamydophila'] = 'pneumonia'\n",
    "mapping['Legionella'] = 'pneumonia'\n",
    "mapping['Normal'] = 'normal'\n",
    "mapping['Lung Opacity'] = 'pneumonia'\n",
    "mapping['1'] = 'pneumonia'\n",
    "\n",
    "# train/test split\n",
    "split = 0.1\n",
    "\n",
    "# to avoid duplicates\n",
    "patient_imgpath = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://github.com/mlmed/torchxrayvision/blob/master/torchxrayvision/datasets.py#L814\n",
    "cohen_csv = pd.read_csv(cohen_csvpath, nrows=None)\n",
    "#idx_pa = csv[\"view\"] == \"PA\"  # Keep only the PA view\n",
    "views = [\"PA\", \"AP\", \"AP Supine\", \"AP semi erect\", \"AP erect\"]\n",
    "cohen_idx_keep = cohen_csv.view.isin(views)\n",
    "cohen_csv = cohen_csv[cohen_idx_keep]\n",
    "\n",
    "fig1_csv = pd.read_csv(fig1_csvpath, encoding='ISO-8859-1', nrows=None)\n",
    "#fig1_idx_keep = fig1_csv.view.isin(views)\n",
    "#fig1_csv = fig1_csv[fig1_idx_keep]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data distribution from covid-chestxray-dataset:\n",
      "{'normal': 0, 'pneumonia': 33, 'COVID-19': 286}\n"
     ]
    }
   ],
   "source": [
    "# get non-COVID19 viral, bacteria, and COVID-19 infections from covid-chestxray-dataset\n",
    "# stored as patient id, image filename and label\n",
    "filename_label = {'normal': [], 'pneumonia': [], 'COVID-19': []}\n",
    "count = {'normal': 0, 'pneumonia': 0, 'COVID-19': 0}\n",
    "for index, row in cohen_csv.iterrows():\n",
    "    f = row['finding'].split(',')[0] # take the first finding, for the case of COVID-19, ARDS\n",
    "    if f in mapping: # \n",
    "        count[mapping[f]] += 1\n",
    "        entry = [str(row['patientid']), row['filename'], mapping[f], row['view']]\n",
    "        filename_label[mapping[f]].append(entry)\n",
    "        \n",
    "for index, row in fig1_csv.iterrows():\n",
    "    if not str(row['finding']) == 'nan':\n",
    "        f = row['finding'].split(',')[0] # take the first finding\n",
    "        if f in mapping: # \n",
    "            count[mapping[f]] += 1\n",
    "            if os.path.exists(os.path.join(fig1_imgpath, row['patientid'] + '.jpg')):\n",
    "                entry = [row['patientid'], row['patientid'] + '.jpg', mapping[f]]\n",
    "            elif os.path.exists(os.path.join(fig1_imgpath, row['patientid'] + '.png')):\n",
    "                entry = [row['patientid'], row['patientid'] + '.png', mapping[f]]\n",
    "            filename_label[mapping[f]].append(entry)\n",
    "\n",
    "print('Data distribution from covid-chestxray-dataset:')\n",
    "print(count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key:  pneumonia\n",
      "Test patients:  ['8', '31']\n",
      "Key:  COVID-19\n",
      "Test patients:  ['19', '20', '36', '42', '86', '94', '97', '117', '132', '138', '144', '150', '163', '169']\n",
      "test count:  {'normal': 0, 'pneumonia': 5, 'COVID-19': 31}\n",
      "train count:  {'normal': 0, 'pneumonia': 28, 'COVID-19': 251}\n"
     ]
    }
   ],
   "source": [
    "# add covid-chestxray-dataset into COVIDx dataset\n",
    "# since covid-chestxray-dataset doesn't have test dataset\n",
    "# split into train/test by patientid\n",
    "# for COVIDx:\n",
    "# patient 8 is used as non-COVID19 viral test\n",
    "# patient 31 is used as bacterial test\n",
    "# patients 19, 20, 36, 42, 86 are used as COVID-19 viral test\n",
    "\n",
    "for key in filename_label.keys():\n",
    "    arr = np.array(filename_label[key])\n",
    "    if arr.size == 0:\n",
    "        continue\n",
    "    # split by patients\n",
    "    # num_diff_patients = len(np.unique(arr[:,0]))\n",
    "    # num_test = max(1, round(split*num_diff_patients))\n",
    "    # select num_test number of random patients\n",
    "    if key == 'pneumonia':\n",
    "        test_patients = ['8', '31']\n",
    "    elif key == 'COVID-19':\n",
    "        test_patients = ['19', '20', '36', '42', '86', \n",
    "                         '94', '97', '117', '132', \n",
    "                         '138', '144', '150', '163', '169'] # random.sample(list(arr[:,0]), num_test)\n",
    "    else: \n",
    "        test_patients = []\n",
    "    print('Key: ', key)\n",
    "    print('Test patients: ', test_patients)\n",
    "    # go through all the patients\n",
    "    for patient in arr:\n",
    "        if patient[0] not in patient_imgpath:\n",
    "            patient_imgpath[patient[0]] = [patient[1]]\n",
    "        else:\n",
    "            if patient[1] not in patient_imgpath[patient[0]]:\n",
    "                patient_imgpath[patient[0]].append(patient[1])\n",
    "            else:\n",
    "                continue  # skip since image has already been written\n",
    "        if patient[0] in test_patients:\n",
    "            copyfile(os.path.join(cohen_imgpath, patient[1]), os.path.join(savepath, 'test', patient[1]))\n",
    "            test.append(patient)\n",
    "            test_count[patient[2]] += 1\n",
    "        else:\n",
    "            if 'COVID' in patient[0]:\n",
    "                copyfile(os.path.join(fig1_imgpath, patient[1]), os.path.join(savepath, 'train', patient[1]))\n",
    "            else:\n",
    "                copyfile(os.path.join(cohen_imgpath, patient[1]), os.path.join(savepath, 'train', patient[1]))\n",
    "            train.append(patient)\n",
    "            train_count[patient[2]] += 1\n",
    "\n",
    "print('test count: ', test_count)\n",
    "print('train count: ', train_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add normal and rest of pneumonia cases from https://www.kaggle.com/c/rsna-pneumonia-detection-challenge\n",
    "csv_normal = pd.read_csv(os.path.join(rsna_datapath, rsna_csvname), nrows=None)\n",
    "csv_pneu = pd.read_csv(os.path.join(rsna_datapath, rsna_csvname2), nrows=None)\n",
    "patients = {'normal': [], 'pneumonia': []}\n",
    "\n",
    "for index, row in csv_normal.iterrows():\n",
    "    if row['class'] == 'Normal':\n",
    "        patients['normal'].append(row['patientId'])\n",
    "\n",
    "for index, row in csv_pneu.iterrows():\n",
    "    if int(row['Target']) == 1:\n",
    "        patients['pneumonia'].append(row['patientId'])\n",
    "\n",
    "for key in patients.keys():\n",
    "    arr = np.array(patients[key])\n",
    "    if arr.size == 0:\n",
    "        continue\n",
    "    # split by patients \n",
    "    # num_diff_patients = len(np.unique(arr))\n",
    "    # num_test = max(1, round(split*num_diff_patients))\n",
    "    test_patients = np.load('rsna_test_patients_{}.npy'.format(key)) # random.sample(list(arr), num_test), download the .npy files from the repo.\n",
    "    # np.save('rsna_test_patients_{}.npy'.format(key), np.array(test_patients))\n",
    "    for patient in arr:\n",
    "        if patient not in patient_imgpath:\n",
    "            patient_imgpath[patient] = [patient]\n",
    "        else:\n",
    "            continue  # skip since image has already been written\n",
    "                \n",
    "        ds = dicom.dcmread(os.path.join(rsna_datapath, rsna_imgpath, patient + '.dcm'))\n",
    "        pixel_array_numpy = ds.pixel_array\n",
    "        imgname = patient + '.png'\n",
    "        if patient in test_patients:\n",
    "            cv2.imwrite(os.path.join(savepath, 'test', imgname), pixel_array_numpy)\n",
    "            test.append([patient, imgname, key])\n",
    "            test_count[key] += 1\n",
    "        else:\n",
    "            cv2.imwrite(os.path.join(savepath, 'train', imgname), pixel_array_numpy)\n",
    "            train.append([patient, imgname, key])\n",
    "            train_count[key] += 1\n",
    "\n",
    "print('test count: ', test_count)\n",
    "print('train count: ', train_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# final stats\n",
    "print('Final stats')\n",
    "print('Train count: ', train_count)\n",
    "print('Test count: ', test_count)\n",
    "print('Total length of train: ', len(train))\n",
    "print('Total length of test: ', len(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export to train and test csv\n",
    "# format as patientid, filename, label, separated by a space\n",
    "train_file = open(\"train_split_v3.txt\",\"a\") \n",
    "for sample in train:\n",
    "    if len(sample) == 4:\n",
    "        info = str(sample[0]) + ' ' + sample[1] + ' ' + sample[2] + ' ' + sample[3] + '\\n'\n",
    "    else:\n",
    "        info = str(sample[0]) + ' ' + sample[1] + ' ' + sample[2] + '\\n'\n",
    "    train_file.write(info)\n",
    "\n",
    "train_file.close()\n",
    "\n",
    "test_file = open(\"test_split_v3.txt\", \"a\")\n",
    "for sample in test:\n",
    "    if len(sample) == 4:\n",
    "        info = str(sample[0]) + ' ' + sample[1] + ' ' + sample[2] + ' ' + sample[3] + '\\n'\n",
    "    else:\n",
    "        info = str(sample[0]) + ' ' + sample[1] + ' ' + sample[2] + '\\n'\n",
    "    test_file.write(info)\n",
    "\n",
    "test_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}