create_COVIDx_CT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# COVIDx CT Dataset Constructor\n",
    "This notebook constructs the COVIDx CT dataset from scratch using the raw data. See [docs/dataset.md](docs/dataset.md) for more details on manual steps which must be completed beforehand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import glob\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Paths\n",
    "Paths to the source data and output location should be set here. Note that the window width and window level for processing scans in Hounsfield units are defined in [data_utils.py](data_utils.py) as `HU_WINDOW_WIDTH = 1500` and `HU_WINDOW_LEVEL = -600`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select dataset version. Major version (e.g., 2) should match current release.\n",
    "# Checkout earlier versions of the repo to construct previous dataset versions.\n",
    "DATASET_VERSION = '3A'\n",
    "\n",
    "# CNCB AI Diagnosis paths\n",
    "CNCB_EXCLUDE_FILE = 'dataset_construction/metadata/cncb_exclude_list.txt'\n",
    "CNCB_EXTRA_LESION_FILE = 'dataset_construction/metadata/cncb_extra_lesions_slices.csv'\n",
    "CNCB_DIR = 'data/CNCB_AIDiagnosis'\n",
    "\n",
    "# Radiopaedia/coronacases paths\n",
    "RADIOPAEDIA_CORONACASES_CT_DIR = 'data/Coronacases_Radiopaedia/COVID-19-CT-Seg_20cases'\n",
    "RADIOPAEDIA_CORONACASES_SEG_DIR = 'data/Coronacases_Radiopaedia/Infection_Mask'\n",
    "RADIOPAEDIA_META_CSV = 'dataset_construction/metadata/radiopaedia_metadata.csv'\n",
    "RADIOPAEDIA_EXCLUDE_FILE = 'dataset_construction/metadata/radiopaedia_exclude_list.txt'\n",
    "\n",
    "# LIDC-IDRI paths\n",
    "LIDC_META_CSV = 'dataset_construction/metadata/lidc_idri_metadata.csv'\n",
    "\n",
    "# COVID-19-20 paths\n",
    "COVID_19_20_DIR = 'data/COVID-19-20_v2/Train'\n",
    "\n",
    "# TCIA COVID-19 paths\n",
    "TCIA_COVID_META_CSV = 'dataset_construction/metadata/tcia_covid_metadata.csv'\n",
    "TCIA_DIR = 'data/CT_Images_in_COVID-19_August_2020'\n",
    "\n",
    "# COVID-CTset paths\n",
    "COVID_CTSET_META_CSV = 'data/COVID-CTset/Labels&Detailes/Patient_details.csv'\n",
    "COVID_CTSET_DIR = 'data/COVID-CTset/Train&Validation'\n",
    "\n",
    "# iCTCF paths\n",
    "ICTCF_META_CSV = 'dataset_construction/metadata/ictcf_metadata.csv'\n",
    "ICTCF_DIR = 'data/iCTCF'\n",
    "\n",
    "# COVID-CT-MD paths\n",
    "COVID_CT_MD_DIR = 'data/COVID-CT-MD'\n",
    "COVID_CT_MD_INDEX_CSV = 'data/COVID-CT-MD/Index.csv'\n",
    "COVID_CT_MD_META_CSV = 'dataset_construction/metadata/covid_ct_md_metadata.csv'\n",
    "COVID_CT_MD_LABEL_CSV = 'data/COVID-CT-MD/Slice-level-labels.npy'\n",
    "\n",
    "# STOIC paths\n",
    "STOIC_META_CSV = 'dataset_construction/metadata/stoic_metadata.csv'\n",
    "STOIC_DIR = 'data/STOIC/data/mha'\n",
    "\n",
    "# Stony Brook paths\n",
    "STONYBROOK_META_CSV = 'dataset_construction/metadata/stonybrook_metadata.csv'\n",
    "STONYBROOK_DIR = 'data/COVID-19-NY-SBU'\n",
    "\n",
    "# MosMedData paths\n",
    "MOSMED_CT_DIR = 'data/MosMedData/COVID19_1110/studies'\n",
    "MOSMED_SEG_DIR = 'data/MosMedData/COVID19_1110/masks'\n",
    "MOSMED_META_CSV = 'dataset_construction/metadata/mosmed_metadata.csv'\n",
    "\n",
    "# Output directory path\n",
    "OUTPUT_DIR = 'data/COVIDx_CT-{}'.format(DATASET_VERSION)  # directory to save the images in\n",
    "\n",
    "# Make output directory\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "These cells process the data from each source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filename and class lists\n",
    "filenames, classes = [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████| 115837/115837 [00:11<00:00, 9903.68it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process CNCB AI Diagnosis data\n",
    "from dataset_construction import cncb\n",
    "fnames, cls = cncb.process_cncb_data(CNCB_DIR, CNCB_EXCLUDE_FILE, OUTPUT_DIR, extra_lesion_files=[CNCB_EXTRA_LESION_FILE])\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 20/20 [00:47<00:00,  2.35s/it]\n"
     ]
    }
   ],
   "source": [
    "# Process segmented radiopaedia and coronacases data\n",
    "from dataset_construction import radiopaedia_coronacases as rc\n",
    "fnames, cls = rc.process_radiopaedia_and_coronacases_seg_data(\n",
    "    RADIOPAEDIA_CORONACASES_CT_DIR, RADIOPAEDIA_CORONACASES_SEG_DIR, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 118/118 [00:44<00:00,  2.68it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process additional radiopaedia data\n",
    "from dataset_construction import radiopaedia\n",
    "fnames, cls = radiopaedia.process_radiopaedia_data(RADIOPAEDIA_META_CSV, RADIOPAEDIA_EXCLUDE_FILE, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 39/39 [02:03<00:00,  3.16s/it]\n"
     ]
    }
   ],
   "source": [
    "# Process LIDC-IDRI data\n",
    "from dataset_construction import lidc_idri\n",
    "fnames, cls = lidc_idri.process_lidc_idri_data(LIDC_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 199/199 [03:35<00:00,  1.08s/it]\n"
     ]
    }
   ],
   "source": [
    "# Process COVID-19-20 challenge data\n",
    "from dataset_construction import covid_19_20\n",
    "fnames, cls = covid_19_20.process_covid_19_20_data(COVID_19_20_DIR, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 445/445 [06:33<00:00,  1.13it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process additional TCIA COVID-19 studies\n",
    "from dataset_construction import tcia_covid\n",
    "fnames, cls = tcia_covid.process_tcia_covid_data(TCIA_DIR, TCIA_COVID_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 377/377 [00:09<00:00, 40.33it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process COVID-CTset data\n",
    "from dataset_construction import covid_ctset\n",
    "fnames, cls = covid_ctset.process_covid_ctset_data(COVID_CTSET_DIR, COVID_CTSET_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 709/709 [01:03<00:00, 11.25it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process iCTCF data\n",
    "from dataset_construction import ictcf\n",
    "fnames, cls = ictcf.process_ictcf_data(ICTCF_DIR, ICTCF_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 301/301 [25:18<00:00,  5.05s/it]\n"
     ]
    }
   ],
   "source": [
    "# Process COVID-CT-MD data\n",
    "from dataset_construction import covid_ct_md as cctmd\n",
    "fnames, cls = cctmd.process_covid_ct_md_data(\n",
    "    COVID_CT_MD_DIR, COVID_CT_MD_INDEX_CSV, COVID_CT_MD_META_CSV, COVID_CT_MD_LABEL_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████| 1176/1176 [1:14:31<00:00,  3.80s/it]\n"
     ]
    }
   ],
   "source": [
    "# Process STOIC challenge data\n",
    "from dataset_construction import stoic\n",
    "fnames, cls = stoic.process_stoic_data(STOIC_DIR, STOIC_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 269/269 [03:47<00:00,  1.18it/s]\n"
     ]
    }
   ],
   "source": [
    "# Process Stony Brook data\n",
    "from dataset_construction import stonybrook as sbrook\n",
    "fnames, cls = sbrook.process_stonybrook_data(STONYBROOK_DIR, STONYBROOK_META_CSV, OUTPUT_DIR)\n",
    "filenames.extend(fnames)\n",
    "classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process MosMedData studies (only for \"B\" variant)\n",
    "if DATASET_VERSION[-1] == 'B':\n",
    "    from dataset_construction import mosmed\n",
    "    ct_1_dir = os.path.join(MOSMED_CT_DIR, 'CT-1')\n",
    "    fnames, cls = mosmed.process_mosmed_seg_data(ct_1_dir, MOSMED_SEG_DIR, OUTPUT_DIR)\n",
    "    filenames.extend(fnames)\n",
    "    classes.extend(cls)\n",
    "    \n",
    "    fnames, cls = mosmed.process_mosmed_unseg_data(MOSMED_META_CSV, MOSMED_CT_DIR, OUTPUT_DIR)\n",
    "    filenames.extend(fnames)\n",
    "    classes.extend(cls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image Counts\n",
      "Normal: []\n",
      "Pneumonia: []\n",
      "COVID-19: [14461]\n"
     ]
    }
   ],
   "source": [
    "# Print image counts\n",
    "from dataset_construction.utils import CLASS_MAP\n",
    "uniq_classes, counts = np.unique(classes, return_counts=True)\n",
    "uniq_classes = np.arange(3, dtype=np.int32)\n",
    "counts = []\n",
    "print('Image Counts')\n",
    "for name, cls in CLASS_MAP.items():\n",
    "    print('{}: {}'.format(name, counts[uniq_classes == cls]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verification\n",
    "Perform a check to ensure that all files are present (optional)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "425024/425024 files created, dataset successfully constructed!\n"
     ]
    }
   ],
   "source": [
    "# Get split files\n",
    "split_files = glob.glob(os.path.join('splits/v' + DATASET_VERSION[0], '*_COVIDx_CT-{}.txt'.format(DATASET_VERSION)))\n",
    "if not len(split_files):\n",
    "    raise ValueError('Split files for COVIDx CT-{} not found'.format(DATASET_VERSION))\n",
    "\n",
    "# Default to \"A\" variant when \"B\" variant files are not present\n",
    "if len(split_files) != 3:\n",
    "    a_var = DATASET_VERSION[0] + 'A'\n",
    "    existing_splits = set(os.path.basename(split_file).split('_')[0] for split_file in split_files)\n",
    "    split_files_a = glob.glob(os.path.join('splits/v' + DATASET_VERSION[0], '*_COVIDx_CT-{}.txt'.format(a_var)))\n",
    "    split_files_a = [f for f in split_files_a if os.path.basename(f).split('_')[0] not in existing_splits]\n",
    "    split_files = split_files + split_files_a\n",
    "\n",
    "# Check that all files from all splits are present in the constructed data\n",
    "count = 0\n",
    "total = 0\n",
    "incomplete = False\n",
    "for split_file in split_files:\n",
    "    with open(split_file, 'r') as f:\n",
    "        for line in f.readlines():\n",
    "            fname, cls = line.strip('\\n').split()[:2]\n",
    "            fpath = os.path.join(OUTPUT_DIR, fname)\n",
    "            \n",
    "            total += 1\n",
    "            if os.path.exists(fpath):\n",
    "                count += 1\n",
    "            else:\n",
    "                print('Missing', fname)\n",
    "                incomplete = True\n",
    "if incomplete:\n",
    "    print('{}/{} files are missing, dataset is incomplete!'.format(count, total))\n",
    "else:\n",
    "    print('{}/{} files created, dataset successfully constructed!'.format(count, total))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}