create_COVIDx.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-19 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 covid-19 dataset from https://github.com/agchung/Actualmed-COVID-chestxray-dataset\n",
    "actmed_imgpath = '../Actualmed-COVID-chestxray-dataset/images'\n",
    "actmed_csvpath = '../Actualmed-COVID-chestxray-dataset/metadata.csv'\n",
    "\n",
    "# path to covid-19 dataset from https://www.kaggle.com/tawsifurrahman/covid19-radiography-database\n",
    "sirm_imgpath = '../COVID-19-Radiography-Database/COVID-19'\n",
    "sirm_csvpath = '../COVID-19-Radiography-Database/COVID-19.metadata.xlsx'\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",
    "actmed_csv = pd.read_csv(actmed_csvpath, nrows=None)\n",
    "\n",
    "sirm_csv = pd.read_excel(sirm_csvpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data distribution from covid datasets:\n",
      "{'normal': 0, 'pneumonia': 33, 'COVID-19': 358}\n"
     ]
    }
   ],
   "source": [
    "# get non-COVID19 viral, bacteria, and COVID-19 infections from covid-chestxray-dataset, figure1 and actualmed\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",
    "covid_ds = {'cohen': [], 'fig1': [], 'actmed': [], 'sirm': []}\n",
    "\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], 'cohen']\n",
    "        filename_label[mapping[f]].append(entry)\n",
    "        if mapping[f] == 'COVID-19':\n",
    "            covid_ds['cohen'].append(str(row['patientid']))\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], 'fig1']\n",
    "            elif os.path.exists(os.path.join(fig1_imgpath, row['patientid'] + '.png')):\n",
    "                entry = [row['patientid'], row['patientid'] + '.png', mapping[f], 'fig1']\n",
    "            filename_label[mapping[f]].append(entry)\n",
    "            if mapping[f] == 'COVID-19':\n",
    "                covid_ds['fig1'].append(row['patientid'])\n",
    "\n",
    "for index, row in actmed_csv.iterrows():\n",
    "    if not str(row['finding']) == 'nan':\n",
    "        f = row['finding'].split(',')[0]\n",
    "        if f in mapping:\n",
    "            count[mapping[f]] += 1\n",
    "            entry = [row['patientid'], row['imagename'], mapping[f], 'actmed']\n",
    "            filename_label[mapping[f]].append(entry)\n",
    "            if mapping[f] == 'COVID-19':\n",
    "                covid_ds['actmed'].append(row['patientid'])\n",
    "    \n",
    "sirm = set(sirm_csv['URL'])\n",
    "cohen = set(cohen_csv['url'])\n",
    "discard = ['100', '101', '102', '103', '104', '105', \n",
    "           '110', '111', '112', '113', '122', '123', \n",
    "           '124', '125', '126', '217']\n",
    "\n",
    "for idx, row in sirm_csv.iterrows():\n",
    "    patientid = row['FILE NAME']\n",
    "    if row['URL'] not in cohen and patientid[patientid.find('(')+1:patientid.find(')')] not in discard:\n",
    "        count[mapping['COVID-19']] += 1\n",
    "        imagename = patientid + '.' + row['FORMAT'].lower()\n",
    "        if not os.path.exists(os.path.join(sirm_imgpath, imagename)):\n",
    "            imagename = patientid.split('(')[0] + ' ('+ patientid.split('(')[1] + '.' + row['FORMAT'].lower()\n",
    "        entry = [patientid, imagename, mapping['COVID-19'], 'sirm']\n",
    "        filename_label[mapping['COVID-19']].append(entry)\n",
    "        covid_ds['sirm'].append(patientid)\n",
    "    \n",
    "print('Data distribution from covid datasets:')\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', '174', '175', '179', '190', '191COVID-00024', 'COVID-00025', 'COVID-00026', 'COVID-00027', 'COVID-00029', 'COVID-00030', 'COVID-00032', 'COVID-00033', 'COVID-00035', 'COVID-00036', 'COVID-00037', 'COVID-00038', 'ANON24', 'ANON45', 'ANON126', 'ANON106', 'ANON67', 'ANON153', 'ANON135', 'ANON44', 'ANON29', 'ANON201', 'ANON191', 'ANON234', 'ANON110', 'ANON112', 'ANON73', 'ANON220', 'ANON189', 'ANON30', 'ANON53', 'ANON46', 'ANON218', 'ANON240', 'ANON100', 'ANON237', 'ANON158', 'ANON174', 'ANON19', 'ANON195', 'COVID-19(119)', 'COVID-19(87)', 'COVID-19(70)', 'COVID-19(94)', 'COVID-19(215)', 'COVID-19(77)', 'COVID-19(213)', 'COVID-19(81)', 'COVID-19(216)', 'COVID-19(72)', 'COVID-19(106)', 'COVID-19(131)', 'COVID-19(107)', 'COVID-19(116)', 'COVID-19(95)', 'COVID-19(214)', 'COVID-19(129)']\n",
      "test count:  {'normal': 0, 'pneumonia': 5, 'COVID-19': 100}\n",
      "train count:  {'normal': 0, 'pneumonia': 28, 'COVID-19': 258}\n"
     ]
    }
   ],
   "source": [
    "# add covid-chestxray-dataset, figure1 and actualmed into COVIDx dataset\n",
    "# since these datasets don't have test dataset, split into train/test by patientid\n",
    "# for covid-chestxray-dataset:\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",
    "# for figure 1:\n",
    "# patients 24, 25, 27, 29, 30, 32, 33, 36, 37, 38\n",
    "\n",
    "ds_imgpath = {'cohen': cohen_imgpath, 'fig1': fig1_imgpath, 'actmed': actmed_imgpath, 'sirm': sirm_imgpath}\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",
    "    # random.sample(list(arr[:,0]), num_test)\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', '174', '175', '179', '190', '191'\n",
    "                         'COVID-00024', 'COVID-00025', 'COVID-00026', 'COVID-00027', 'COVID-00029',\n",
    "                         'COVID-00030', 'COVID-00032', 'COVID-00033', 'COVID-00035', 'COVID-00036',\n",
    "                         'COVID-00037', 'COVID-00038',\n",
    "                         'ANON24', 'ANON45', 'ANON126', 'ANON106', 'ANON67',\n",
    "                         'ANON153', 'ANON135', 'ANON44', 'ANON29', 'ANON201', \n",
    "                         'ANON191', 'ANON234', 'ANON110', 'ANON112', 'ANON73', \n",
    "                         'ANON220', 'ANON189', 'ANON30', 'ANON53', 'ANON46',\n",
    "                         'ANON218', 'ANON240', 'ANON100', 'ANON237', 'ANON158',\n",
    "                         'ANON174', 'ANON19', 'ANON195',\n",
    "                         'COVID-19(119)', 'COVID-19(87)', 'COVID-19(70)', 'COVID-19(94)', \n",
    "                         'COVID-19(215)', 'COVID-19(77)', 'COVID-19(213)', 'COVID-19(81)', \n",
    "                         'COVID-19(216)', 'COVID-19(72)', 'COVID-19(106)', 'COVID-19(131)', \n",
    "                         'COVID-19(107)', 'COVID-19(116)', 'COVID-19(95)', 'COVID-19(214)', \n",
    "                         'COVID-19(129)']\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",
    "            if patient[3] == 'sirm':\n",
    "                image = cv2.imread(os.path.join(ds_imgpath[patient[3]], patient[1]))\n",
    "                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "                patient[1] = patient[1].replace(' ', '')\n",
    "                cv2.imwrite(os.path.join(savepath, 'test', patient[1]), gray)\n",
    "            else:\n",
    "                copyfile(os.path.join(ds_imgpath[patient[3]], patient[1]), os.path.join(savepath, 'test', patient[1]))\n",
    "            test.append(patient)\n",
    "            test_count[patient[2]] += 1\n",
    "        else:\n",
    "            if patient[3] == 'sirm':\n",
    "                image = cv2.imread(os.path.join(ds_imgpath[patient[3]], patient[1]))\n",
    "                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "                patient[1] = patient[1].replace(' ', '')\n",
    "                cv2.imwrite(os.path.join(savepath, 'train', patient[1]), gray)\n",
    "            else:\n",
    "                copyfile(os.path.join(ds_imgpath[patient[3]], 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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test count:  {'normal': 885, 'pneumonia': 594, 'COVID-19': 100}\n",
      "train count:  {'normal': 7966, 'pneumonia': 5451, 'COVID-19': 258}\n"
     ]
    }
   ],
   "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, 'rsna'])\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, 'rsna'])\n",
    "            train_count[key] += 1\n",
    "\n",
    "print('test count: ', test_count)\n",
    "print('train count: ', train_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final stats\n",
      "Train count:  {'normal': 7966, 'pneumonia': 5451, 'COVID-19': 258}\n",
      "Test count:  {'normal': 885, 'pneumonia': 594, 'COVID-19': 100}\n",
      "Total length of train:  13675\n",
      "Total length of test:  1579\n"
     ]
    }
   ],
   "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": 8,
   "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.txt\",'w') \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.txt\", 'w')\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 (covid)",
   "language": "python",
   "name": "covid"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}