create_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c55ed9c",
   "metadata": {},
   "source": [
    "# Creating dataset from LIDC CT scans"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c682db2b",
   "metadata": {},
   "source": [
    "With this notebook volumes around the tumors and associated annotations are extracted. <br/>\n",
    "Extracted annotations include malignancy status and biomarkers such as diameter, spiculation, calcification .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b37134fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import glob\n",
    "import pylidc as pl\n",
    "from pylidc.utils import consensus\n",
    "import torchio as tio\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from scipy import ndimage\n",
    "import pickle\n",
    "import sys, os\n",
    "import ctypes\n",
    "from contextlib import contextmanager\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "528072f9",
   "metadata": {},
   "source": [
    "### define paths and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "815a87a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define path to LIDC-IDRI data\n",
    "base_path = \"/home/lbrocki/LIDC/LIDC-IDRI\"\n",
    "# define path to save the dataset\n",
    "save_path = \"/home/lbrocki/ConRad/dataset\"\n",
    "# define side length of tumor volume in mm\n",
    "h = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89adb922",
   "metadata": {
    "tags": []
   },
   "source": [
    "### helper code to catch warning from torchio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c0adeb4",
   "metadata": {
    "code_folding": [
     0
    ],
    "tags": []
   },
   "outputs": [],
   "source": [
    "# this code is only needed to catch the warning from torchio, samples with a warning are excluded from the dataset\n",
    "def flush(stream):\n",
    "    try:\n",
    "        ctypes.libc.fflush(None)\n",
    "        stream.flush()\n",
    "    except (AttributeError, ValueError, IOError):\n",
    "        pass  # unsupported\n",
    "def fileno(file_or_fd):\n",
    "    fd = getattr(file_or_fd, \"fileno\", lambda: file_or_fd)()\n",
    "    if not isinstance(fd, int):\n",
    "        raise ValueError(\"Expected a file (`.fileno()`) or a file descriptor\")\n",
    "    return fd\n",
    "@contextmanager\n",
    "def stdout_redirected(to=os.devnull, stdout=None):\n",
    "    if stdout is None:\n",
    "        stdout = sys.stdout\n",
    "    stdout_fd = fileno(stdout)\n",
    "\n",
    "    with os.fdopen(os.dup(stdout_fd), \"wb\") as copied:\n",
    "        flush(stdout)\n",
    "        try:\n",
    "            os.dup2(fileno(to), stdout_fd)  # $ exec >&to\n",
    "        except ValueError:  # filename\n",
    "            with open(to, \"wb\") as to_file:\n",
    "                os.dup2(to_file.fileno(), stdout_fd)  # $ exec > to\n",
    "        try:\n",
    "            yield stdout  # allow code to be run with the redirected stdout\n",
    "        finally:\n",
    "            flush(stdout)\n",
    "            os.dup2(copied.fileno(), stdout_fd)  # $ exec >&copied"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05ec75bb",
   "metadata": {},
   "source": [
    "### define the extraction of nodule volumes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8b48e9de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract slices of dimension h x h x h around nodule center \n",
    "# and the corresponding segmentations\n",
    "def process_nodule_vol(nodule, h=32):  \n",
    "    dicom = nodule[0].scan.get_path_to_dicom_files()\n",
    "    median_malig = np.median([nod.malignancy for nod in nodule])\n",
    "    \n",
    "    \n",
    "#     # catch errors from torchio and throw an exception so that this nodule is skipped\n",
    "    with open(\"output.txt\", \"w\") as f, stdout_redirected(f, stdout=sys.stderr):\n",
    "        tio_image = tio.ScalarImage(dicom)\n",
    "        spacing = tio_image.spacing\n",
    "        # resample isotropically to 1mm spacing\n",
    "        transform = tio.Resample(1)\n",
    "        res_image = transform(tio_image)\n",
    "        res_data = torch.movedim(res_image.data, (0,1,2,3), (0,2,1,3))\n",
    "        \n",
    "    with open(\"output.txt\") as f:\n",
    "        content = f.read()\n",
    "    if \"warning\" in content.lower():\n",
    "        raise RuntimeError(\"SimpleITK Warning .. skip!\")\n",
    "    open(\"output.txt\", \"w\").close()\n",
    "    \n",
    "    cmask,cbbox,masks = consensus(nodule, clevel=0.5)\n",
    "    \n",
    "    # resample cbbox accordingly\n",
    "    res_cbbox = [(round(cbbox[i].start*spacing[i]), \n",
    "                  round(cbbox[i].stop*spacing[i])) for i in range(3)]\n",
    "    \n",
    "    res_cmask = ndimage.zoom(cmask.astype(int), spacing)\n",
    "    \n",
    "    # center of cbbox\n",
    "    res_cbbox0 = [round((res_cbbox[i][0]+res_cbbox[i][1])/2) for i in range(3)]\n",
    "    \n",
    "    # cmask is given realtive to cbbox, express relative to original volume\n",
    "    g = np.zeros(res_data.shape[1:])\n",
    "    g[res_cbbox[0][0]:res_cbbox[0][0]+res_cmask.shape[0], \n",
    "      res_cbbox[1][0]:res_cbbox[1][0]+res_cmask.shape[1],\n",
    "      res_cbbox[2][0]:res_cbbox[2][0]+res_cmask.shape[2],] = res_cmask\n",
    "\n",
    "    # extract volumes of dimension 2k x 2k x 2k\n",
    "    k = int(h/2)\n",
    "    slices = (\n",
    "                slice(res_cbbox0[0]-k, res_cbbox0[0]+k),\n",
    "                slice(res_cbbox0[1]-k, res_cbbox0[1]+k),\n",
    "                slice(res_cbbox0[2]-k, res_cbbox0[2]+k)\n",
    "             )\n",
    "\n",
    "    crop = res_data[0][slices]\n",
    "\n",
    "    g = torch.tensor(g)\n",
    "    mask = g[slices]\n",
    "    \n",
    "    assert crop.shape == torch.Size([h,h,h])\n",
    "    assert mask.shape == torch.Size([h,h,h])\n",
    "    \n",
    "    return median_malig, crop, mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b36eba8f",
   "metadata": {},
   "source": [
    "### extract nodule volumes, associated tumor masks and annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "daacad10",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:16<00:00,  3.38s/it]\n"
     ]
    }
   ],
   "source": [
    "# labels: 1 -> malignant, 0 -> benign\n",
    "labels = {}\n",
    "new_id = 1\n",
    "match = []\n",
    "Path(f\"{save_path}/crops\").mkdir(parents=True, exist_ok=True)\n",
    "Path(f\"{save_path}/masks\").mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "attributes = [\n",
    "    \"subtlety\",\n",
    "    \"internalStructure\",\n",
    "    \"calcification\",\n",
    "    \"sphericity\",\n",
    "    \"margin\",\n",
    "    \"lobulation\",\n",
    "    \"spiculation\",\n",
    "    \"texture\"\n",
    "]\n",
    "\n",
    "d = glob.glob(f\"{base_path}/*\")\n",
    "d.sort()\n",
    "# list of patient ids\n",
    "pids = [i.split(\"/\")[-1].split(\"-\")[-1] for i in d]\n",
    "\n",
    "# diameter has to be treated separately because of how the annotations are organized\n",
    "avg_annotations = {}\n",
    "for att in attributes:\n",
    "    avg_annotations[att] = []\n",
    "avg_annotations[\"diameter\"] = []\n",
    "avg_annotations[\"target\"] = []\n",
    "avg_annotations[\"path\"] = []\n",
    "\n",
    "for patient_id in tqdm(pids):\n",
    "    #print(patient_id, new_id)\n",
    "    scan = pl.query(pl.Scan).filter(pl.Scan.patient_id == f\"LIDC-IDRI-{patient_id}\")[0]\n",
    "    nodules = scan.cluster_annotations()\n",
    "    if(len(nodules) == 0):\n",
    "        continue\n",
    "    k = 0\n",
    "    for nodule in nodules:\n",
    "        num_annotations = len(nodule)\n",
    "        # only nodules with > 2 annotations are considered\n",
    "        # if there are more than 4 annotations, clustering of annotations is ambiguous --> skip\n",
    "        if(num_annotations > 4):\n",
    "            print(\"skipping!\")\n",
    "        if(num_annotations > 2 and num_annotations <= 4):\n",
    "            try:\n",
    "                median_malig, crop, mask = process_nodule_vol(nodule, h=h)\n",
    "                str_new_id = str(new_id).zfill(4)\n",
    "                append = False\n",
    "                if(median_malig > 3):\n",
    "                    avg_annotations[\"target\"].append(1)\n",
    "                    append = True\n",
    "                    new_id += 1\n",
    "                elif(median_malig < 3):\n",
    "                    avg_annotations[\"target\"].append(0)\n",
    "                    append = True\n",
    "                    new_id += 1\n",
    "                if(append):\n",
    "                    avg_annotations[\"diameter\"].append(np.mean([ann.diameter for ann in nodule]))\n",
    "                    avg_annotations[\"path\"].append(f\"{str_new_id}.pt\")\n",
    "                    for att in attributes:\n",
    "                        avg_annotations[att].append(np.mean([vars(ann)[att] for ann in nodule]))\n",
    "                        \n",
    "                    match.append([patient_id, k, new_id])\n",
    "                    torch.save(crop.clone(), f\"{save_path}/crops/{str_new_id}.pt\")\n",
    "                    torch.save(mask.clone(), f\"{save_path}/masks/{str_new_id}.pt\")\n",
    "            # if creation of crop fails for any reason, skip to next nodule\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "        k += 1\n",
    "with open(f\"{save_path}/match.pkl\", \"wb\") as handle:\n",
    "    pickle.dump(match, handle)\n",
    "    \n",
    "with open(f\"{save_path}/annotations_new.pkl\", \"wb\") as handle:\n",
    "    pickle.dump(avg_annotations, handle)\n",
    "os.remove(\"output.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "83826656",
   "metadata": {},
   "outputs": [],
   "source": [
    "#create pandas dataframe from dictionary and save as pickle\n",
    "with open(f\"{save_path}/annotations_new.pkl\", \"rb\") as f:\n",
    "    avg_annotations = pickle.load(f)\n",
    "df_ann = pd.DataFrame.from_dict(avg_annotations)\n",
    "df_ann.to_pickle(f\"{save_path}/annotations_df.pkl\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}