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\n",
    "import glob"
   ]
  },
  {
   "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'\n",
    "sirm_csvpath = '../COVID-19-Radiography-Database/COVID.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",
    "# path to ricord covid-19 images created by create_ricord_dataset/create_ricord_dataset.ipynb\n",
    "# run create_ricord_dataset.ipynb before this notebook\n",
    "ricord_imgpath = 'create_ricord_dataset/ricord_images'\n",
    "ricord_txt = 'create_ricord_dataset/ricord_data_set.txt'\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['E.Coli'] = '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": 5,
   "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data distribution from covid datasets:\n",
      "{'normal': 0, 'pneumonia': 57, 'COVID-19': 2449}\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': [], 'ricord': []}\n",
    "\n",
    "for index, row in cohen_csv.iterrows():\n",
    "    f = row['finding'].split('/')[-1] # take final finding in hierarchy, for the case of COVID-19, ARDS\n",
    "    if f == 'COVID-19' and ('eurorad.org' in row['url'] or 'ml-workgroup' in row['url'] or 'sirm.org' in row['url']):\n",
    "        # skip COVID-19 positive images from eurorad to not duplicate sirm images\n",
    "        pass\n",
    "    elif 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",
    "cohen.add('https://github.com/ieee8023/covid-chestxray-dataset')  # Add base URL to remove sirm images from ieee dataset\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 = \"COVID ({}).png\".format(imagename.rsplit(\".png\")[0].split(\"COVID \")[1])\n",
    "        entry = [patientid, imagename, mapping['COVID-19'], 'sirm']\n",
    "        filename_label[mapping['COVID-19']].append(entry)\n",
    "        covid_ds['sirm'].append(patientid)\n",
    "\n",
    "# get ricord file names \n",
    "with open(ricord_txt) as f:\n",
    "    ricord_file_names = [line.split()[0] for line in f]\n",
    "\n",
    "for imagename in ricord_file_names:\n",
    "    count[mapping['COVID-19']] += 1 # since RICORD data is all COVID-19 postive images\n",
    "    patientid = imagename.split('-')[3] + '-' + imagename.split('-')[4]\n",
    "    \n",
    "    entry = [patientid, imagename, mapping['COVID-19'], 'ricord']\n",
    "    filename_label[mapping['COVID-19']].append(entry)\n",
    "    \n",
    "    covid_ds['ricord'].append(patientid)\n",
    "    \n",
    "print('Data distribution from covid datasets:')\n",
    "print(count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "176\n"
     ]
    }
   ],
   "source": [
    "# Create list of RICORD patients to be added to test, equal to 200 images\n",
    "\n",
    "# We want to prevent patients present in both train and test\n",
    "# Get list of patients who have one image\n",
    "ricord_patients = []\n",
    "for label in filename_label['COVID-19']:\n",
    "    if label[3] == 'ricord':\n",
    "        ricord_patients.append(label[0])\n",
    "\n",
    "pt_with_one_image = [x for x in ricord_patients if ricord_patients.count(x) == 1] # contrains 176 patients\n",
    "print(len(pt_with_one_image))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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 119', 'COVID 87', 'COVID 70', 'COVID 94', 'COVID 215', 'COVID 77', 'COVID 213', 'COVID 81', 'COVID 216', 'COVID 72', 'COVID 106', 'COVID 131', 'COVID 107', 'COVID 116', 'COVID 95', 'COVID 214', 'COVID 129', '419639-000573', '419639-001015', '419639-001070', '419639-001433', '419639-001571', '419639-001603', '419639-001629', '419639-001762', '419639-001882', '419639-001932', '419639-002043', '419639-002050', '419639-002088', '419639-002374', '419639-002376', '419639-002733', '419639-003021', '419639-003024', '419639-003027', '419639-003060', '419639-003080', '419639-003099', '419639-003116', '419639-003169', '419639-003187', '419639-003199', '419639-003203', '419639-003216', '419639-003232', '419639-003251', '419639-003319', '419639-003336', '419639-003365', '419639-003380', 'SITE2-000002', 'SITE2-000003', 'SITE2-000014', 'SITE2-000029', 'SITE2-000034', 'SITE2-000068', 'SITE2-000072', 'SITE2-000094', 'SITE2-000096', 'SITE2-000097', 'SITE2-000098', 'SITE2-000099', 'SITE2-000108', 'SITE2-000115', 'SITE2-000120', 'SITE2-000127', 'SITE2-000131', 'SITE2-000135', 'SITE2-000141', 'SITE2-000142', 'SITE2-000147', 'SITE2-000159', 'SITE2-000162', 'SITE2-000168', 'SITE2-000170', 'SITE2-000171', 'SITE2-000173', 'SITE2-000178', 'SITE2-000202', 'SITE2-000203', 'SITE2-000211', 'SITE2-000217', 'SITE2-000229', 'SITE2-000230', 'SITE2-000237', 'SITE2-000239', 'SITE2-000244', 'SITE2-000265', 'SITE2-000270', 'SITE2-000272', 'SITE2-000291', 'SITE2-000154', '419639-000002', '419639-000082', '419639-000086', '419639-000235', '419639-000299', '419639-000402', '419639-000432', '419639-000538', '419639-000582', '419639-000586', '419639-000785', '419639-000918', '419639-001007', '419639-001085', '419639-001155', '419639-001222', '419639-001287', '419639-001288', '419639-001374', '419639-001425', '419639-001505', '419639-001661', '419639-001699', '419639-001735', '419639-001782', '419639-001829', '419639-001903', '419639-002114', '419639-002281', '419639-002309', '419639-002310', '419639-002401', '419639-002463', '419639-002472', '419639-002528', '419639-002601', '419639-002724', '419639-002757', '419639-002768', '419639-002784', '419639-002871', '419639-002889', '419639-002895', '419639-002929', '419639-002997', '419639-003102', '419639-003105', '419639-003157', '419639-003186', '419639-003191', '419639-003257', '419639-003318', '419639-003334', '419639-003338', '419639-003340', '419639-003360', '419639-003371', '440808-000009', '440808-000016', '440808-000019', '440808-000020', 'SITE2-000001', 'SITE2-000015', 'SITE2-000018', 'SITE2-000023', 'SITE2-000028', 'SITE2-000035', 'SITE2-000045', 'SITE2-000048', 'SITE2-000058', 'SITE2-000066', 'SITE2-000067', 'SITE2-000105', 'SITE2-000106', 'SITE2-000116', 'SITE2-000124', 'SITE2-000125', 'SITE2-000136', 'SITE2-000149', 'SITE2-000153', 'SITE2-000155', 'SITE2-000176', 'SITE2-000179', 'SITE2-000181', 'SITE2-000182', 'SITE2-000185', 'SITE2-000188', 'SITE2-000200', 'SITE2-000204', 'SITE2-000206', 'SITE2-000207', 'SITE2-000208', 'SITE2-000210', 'SITE2-000220', 'SITE2-000221', 'SITE2-000235', 'SITE2-000238', 'SITE2-000279', 'SITE2-000283', 'SITE2-000297', '419639-000025', '419639-001464']\n",
      "test count:  {'normal': 0, 'pneumonia': 5, 'COVID-19': 291}\n",
      "train count:  {'normal': 0, 'pneumonia': 52, 'COVID-19': 2158}\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, 'ricord':ricord_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 119', 'COVID 87', 'COVID 70', 'COVID 94', \n",
    "                         'COVID 215', 'COVID 77', 'COVID 213', 'COVID 81', \n",
    "                         'COVID 216', 'COVID 72', 'COVID 106', 'COVID 131', \n",
    "                         'COVID 107', 'COVID 116', 'COVID 95', 'COVID 214', \n",
    "                         'COVID 129']\n",
    "#         Add 178 RICORD patients to COVID-19, equal to 200 images\n",
    "        test_patients.extend(pt_with_one_image)\n",
    "        test_patients.extend(['419639-000025', '419639-001464'])\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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test count:  {'normal': 885, 'pneumonia': 594, 'COVID-19': 291}\n",
      "train count:  {'normal': 7966, 'pneumonia': 5475, 'COVID-19': 2158}\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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final stats\n",
      "Train count:  {'normal': 7966, 'pneumonia': 5475, 'COVID-19': 2158}\n",
      "Test count:  {'normal': 885, 'pneumonia': 594, 'COVID-19': 291}\n",
      "Total length of train:  15599\n",
      "Total length of test:  1770\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": 11,
   "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 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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}