ProstateX preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ProstateX(2) preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will read and preprocess all ProstateX and ProstateX2 images so that they can be read by the Retina UNet for the purposes of lesion detection and classification.\n",
    "\n",
    "Before running this notebook, the following things will be needed:\n",
    "\n",
    "- Download TRAIN and / or TEST images from ProstateX challenge website: \n",
    "  - Go to: https://wiki.cancerimagingarchive.net/display/Public/SPIE-AAPM-NCI+PROSTATEx+Challenges\n",
    "  - Go to the \"Detailed Description\" tab \n",
    "  - In Section \"PROSTATEx Challenge (November 21, 2016 to February 16, 2017)\" download all the data, except for the .bmp files\n",
    "  - In Section \"PROSTATEx-2 — SPIE-AAPM-NCI Prostate MR Gleason Grade Group Challenge (May 15, 2017 to August 3, 2017)\" download only lesion information (.zip), as the images are exactly the same as the ones we have already downloaded.\n",
    "  - Once all images have been downloaded, please modify the path below to point to them\n",
    "  \n",
    "  \n",
    "- We will need ``plot_lib`` to visualize the images as they are processed:\n",
    " - `plot_lib`: https://github.com/OscarPellicer/plot_lib\n",
    " \n",
    "- Edit the following cell according to where the downloaded data is located.\n",
    "\n",
    "- Finally, make sure to read the Initial Setup Section (below) for further instructions.\n",
    "\n",
    "\n",
    "- Please note that this Notebook will have to be run twice, once for the ProstateX TRAIN data, and other time for the ProstateX test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Path to the ProstateX dataset to be processed (can be either the train or test set)\n",
    "DS_PATH= r'D:\\oscar\\Prostate Images\\ProstateX\\TRAIN'\n",
    "#DS_PATH= r'D:\\oscar\\Prostate Images\\ProstateX\\TEST'\n",
    "\n",
    "#TRAIN is boolean indicating whether we are using the TRAIN or the TEST data\n",
    "TRAIN= True\n",
    "\n",
    "#Main configuration\n",
    "verbose= True #Show extra information during the process\n",
    "apply_registration= True #Use transformations in transforms_path to register the images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intial setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`plot_lib` setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.output { flex-direction: row; flex-wrap: wrap; }\n",
       "         .widget-hslider { width: auto; } </style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Import plot_lib\n",
    "from pathlib import Path\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(Path.home(), 'plot_lib'))\n",
    "from plot_lib import plot_alpha, plot_multi_mask, plot,  plot4\n",
    "\n",
    "#Some CSS to allow images to display side by side by default\n",
    "br= lambda: print(' '*100) \n",
    "from IPython.display import display, HTML\n",
    "CSS = \"\"\".output { flex-direction: row; flex-wrap: wrap; }\n",
    "         .widget-hslider { width: auto; } \"\"\"\n",
    "HTML('<style>{}</style>'.format(CSS))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import other required libraries. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import pydicom\n",
    "import glob\n",
    "import pickle\n",
    "from preprocessing_lib import (info, grow_regions_sitk,\n",
    "                              join_sitk_images, join_masks, read_prostatex_patient,\n",
    "                              rescale_intensity, center_image, get_blank_image,\n",
    "                              get_lesion_mask_id_seed, grow_lesions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to set all the **required paths**. Please, feel free to modify them to point to the correct place if needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#1) Data that must have been downloaded following instructions above\n",
    "#Path where DICOMS are stored\n",
    "dicom_path= os.path.join(DS_PATH, 'Images')\n",
    "\n",
    "#Path where ktrans images are stored\n",
    "ktrans_path= os.path.join(DS_PATH, 'Images Ktrans')\n",
    "\n",
    "#Path where csvs are stored\n",
    "prostateX_csv_path= os.path.join(DS_PATH, 'Lesion Information', \n",
    "                                 'ProstateX-Findings-%s.csv'%('Train' if TRAIN else 'Test'))\n",
    "prostateX2_csv_path= os.path.join(DS_PATH.replace('ProstateX', 'ProstateX2'), 'Lesion Information', \n",
    "                                 'ProstateX-Findings-%s.csv'%('Train' if TRAIN else 'Test'))\n",
    "\n",
    "#2) Data provided alongside the repository:\n",
    "#Path to where ProstateX_masks are stored\n",
    "masks_path= r'ProstateX_masks'\n",
    "\n",
    "#Path where registration transformations are stored\n",
    "transforms_path= r'ProstateX_transforms'\n",
    "\n",
    "#3) Output data to be generated by this Notebook:\n",
    "#Path where data to be read by the Retina UNet model will be stored\n",
    "#We create the directory if it did not exist\n",
    "output_path= 'out' + ('_unregistered' if not apply_registration else '')\n",
    "if not os.path.exists(output_path): os.makedirs(output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Description of the datasets and their labels\n",
    "\n",
    "**ProstateX** contains 204 train and 142 test images of prostates with lesion locations and two significance levels (only in train, in test there are no significance levels):\n",
    " - **False**: Gleason Score <= 3+3\n",
    " - **True**: Gleason Score >= 3+4\n",
    "     \n",
    "References:\n",
    " - https://prostatex.grand-challenge.org/challenge_description/\n",
    " - https://wiki.cancerimagingarchive.net/display/Public/SPIE-AAPM-NCI+PROSTATEx+Challenges\n",
    "     \n",
    "     \n",
    "**ProstateX2** contains the same images (both train and test), but with much more complete label information:\n",
    " - **Grade Group 1 (Gleason score <= 3+3)**: Only individual discrete well-formed glands\n",
    " - **Grade Group 2 (Gleason score 3+4)**: Predominantly well-formed glands with lesser component of poorly-formed/fused/cribriform glands\n",
    " - **Grade Group 3 (Gleason score 4+3)**: Predominantly poorly formed/fused/cribriform glands with lesser component of well-formed glands\n",
    " - **Grade Group 4 (Gleason score 4+4, 3+5, 5+3)**: (1) Only poorly-formed/fused/cribriform glands or (2) predominantly well-formed glands and lesser component lacking glands or (3) predominantly lacking glands and lesser component of well-formed glands\n",
    " - **Grade Group 5 (Higher Gleason scores)**: Lacks gland formation (or with necrosis) with or without poorly formed/fused/cribriform glands\n",
    "\n",
    "References:\n",
    " - Epstein JI, Egevad L, Amin MB, Delahunt B, Srigley JR, Humphrey PA, the Grading Committee. The 2014 International Society of Urologic Pathology (ISUP) Consensus Conference on Gleason Grading of Prostatic Carcinoma: Definition of grading  - patterns and proposal for a new grading system. Am J Surg Pathol, (40)244-252, 2016\n",
    "https://www.aapm.org/GrandChallenge/PROSTATEx-2/\n",
    "\n",
    "As it turns out, the images that appear in ProstateX, but not in ProsateX2, are those that have a GS < 6 (i.e.: benign). \n",
    "\n",
    "Therefore, by combining the information from both datasets, we can create an **extended labelling system** defined as follows:\n",
    " - 0: Normal prostate\n",
    " - 1-5: GG{1-5}\n",
    " - 10: Benign lesion or normal prostate\n",
    " - 20: Unknown (test set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fid</th>\n",
       "      <th>pos</th>\n",
       "      <th>zone</th>\n",
       "      <th>ClinSig</th>\n",
       "      <th>ggg</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProxID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ProstateX-0000</th>\n",
       "      <td>1</td>\n",
       "      <td>25.7457 31.8707 -38.511</td>\n",
       "      <td>PZ</td>\n",
       "      <td>True</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0001</th>\n",
       "      <td>1</td>\n",
       "      <td>-40.5367071921656 29.320722668457 -16.70766907...</td>\n",
       "      <td>AS</td>\n",
       "      <td>False</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0002</th>\n",
       "      <td>1</td>\n",
       "      <td>-27.0102 41.5467 -26.0469</td>\n",
       "      <td>PZ</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0002</th>\n",
       "      <td>2</td>\n",
       "      <td>-2.058 38.6752 -34.6104</td>\n",
       "      <td>PZ</td>\n",
       "      <td>False</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0003</th>\n",
       "      <td>1</td>\n",
       "      <td>22.1495 31.2717 -2.45933</td>\n",
       "      <td>TZ</td>\n",
       "      <td>False</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0003</th>\n",
       "      <td>2</td>\n",
       "      <td>-21.2871 19.3995 19.7429</td>\n",
       "      <td>TZ</td>\n",
       "      <td>False</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0004</th>\n",
       "      <td>1</td>\n",
       "      <td>-7.69665 3.64226 23.1659</td>\n",
       "      <td>AS</td>\n",
       "      <td>False</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0005</th>\n",
       "      <td>0</td>\n",
       "      <td>-14.5174331665039 49.4428329467773 20.78152465...</td>\n",
       "      <td>PZ</td>\n",
       "      <td>True</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0005</th>\n",
       "      <td>1</td>\n",
       "      <td>-38.6276 42.2781 21.4084</td>\n",
       "      <td>PZ</td>\n",
       "      <td>True</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ProstateX-0005</th>\n",
       "      <td>1</td>\n",
       "      <td>-22.0892639160156 25.4668045043945 22.87915420...</td>\n",
       "      <td>TZ</td>\n",
       "      <td>False</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                fid                                                pos zone  \\\n",
       "ProxID                                                                        \n",
       "ProstateX-0000    1                            25.7457 31.8707 -38.511   PZ   \n",
       "ProstateX-0001    1  -40.5367071921656 29.320722668457 -16.70766907...   AS   \n",
       "ProstateX-0002    1                          -27.0102 41.5467 -26.0469   PZ   \n",
       "ProstateX-0002    2                            -2.058 38.6752 -34.6104   PZ   \n",
       "ProstateX-0003    1                           22.1495 31.2717 -2.45933   TZ   \n",
       "ProstateX-0003    2                           -21.2871 19.3995 19.7429   TZ   \n",
       "ProstateX-0004    1                           -7.69665 3.64226 23.1659   AS   \n",
       "ProstateX-0005    0  -14.5174331665039 49.4428329467773 20.78152465...   PZ   \n",
       "ProstateX-0005    1                           -38.6276 42.2781 21.4084   PZ   \n",
       "ProstateX-0005    1  -22.0892639160156 25.4668045043945 22.87915420...   TZ   \n",
       "\n",
       "                ClinSig  ggg  \n",
       "ProxID                        \n",
       "ProstateX-0000     True    3  \n",
       "ProstateX-0001    False    1  \n",
       "ProstateX-0002     True    2  \n",
       "ProstateX-0002    False    1  \n",
       "ProstateX-0003    False   10  \n",
       "ProstateX-0003    False   10  \n",
       "ProstateX-0004    False    1  \n",
       "ProstateX-0005     True    2  \n",
       "ProstateX-0005     True    3  \n",
       "ProstateX-0005    False   10  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Read csvs\n",
    "l_info= pd.read_csv(prostateX_csv_path, index_col='ProxID')\n",
    "l_info_2= pd.read_csv(prostateX2_csv_path, index_col='ProxID')\n",
    "lesion_info= l_info.reset_index().merge(l_info_2.reset_index(), how=\"left\", \n",
    "                                        on=['ProxID', 'pos', 'zone', 'fid']).set_index('ProxID')\n",
    "\n",
    "if TRAIN: \n",
    "    lesion_info.loc[lesion_info.ggg.isna(), 'ggg']= 10\n",
    "    lesion_info.ggg= lesion_info.ggg.astype(int);\n",
    "\n",
    "lesion_info.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read, process and save the images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each patient, we will produce three ouputs to be saved `output_path` for use in the Retina UNet (`ID` represents the actual patient ID):\n",
    "\n",
    "- `ID_img.npy`: It contains the prostate mpMRI with dimensions 160(x) x 160(y) x 24(z) x 7 (channels) and spacing 0.5 x 0.5 x 3 mm. The channels are the following:\n",
    " - 0: T2\n",
    " - 1: B500\n",
    " - 2: B800\n",
    " - 3: ADC\n",
    " - 4: ktrans\n",
    " - 5: Prostate mask\n",
    " - 6: Central Zone mask\n",
    " - 7: Peripheral Zone mask\n",
    " \n",
    " \n",
    "- `ID_rois.npy`: It contains a mask of the lesion IDs\n",
    "\n",
    "\n",
    "- `meta_info_ID.pickle`: It contains a dictionary with items: \n",
    " - 'pid': String with the ID of the patient. E.g.: 'ProstateX-0000'\n",
    " - 'class_target': 1D array with the class associated with each of the lesions. E.g.: np.array([1,10])\n",
    " - 'spacing': tuple with the spacing of the image. E.g.: (0.5, 0.5, 3)\n",
    " - 'fg_slices': z slices where there is at least one lesion. E.g.: [5,6,7,8,9]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ID in os.listdir(dicom_path):\n",
    "    try:\n",
    "        #Print patient ID\n",
    "        print('\\n%s'%ID) \n",
    "        \n",
    "        #Check if registration transform exists and load it\n",
    "        if apply_registration and os.path.exists(os.path.join(transforms_path, ID + '.tfm')):\n",
    "            transform= sitk.ReadTransform(os.path.join(transforms_path, ID + '.tfm'))\n",
    "        else:\n",
    "            transform= sitk.Euler3DTransform()\n",
    "            print('No transform was found (or apply_registration is off). Image might be unregistered')\n",
    "\n",
    "        #---------------Read all images and masks, and then combine them---------------\n",
    "        \n",
    "        #Read all the modalities for a given ProstateX patient ID\n",
    "        #There might be multiple directories (or not)\n",
    "        patient_directories= os.listdir(os.path.join(dicom_path, ID))\n",
    "        if len(patient_directories) != 1: \n",
    "             print(' - Warning: Multiple directories!')\n",
    "        patient_directories= patient_directories[0]\n",
    "        images_path= os.path.join(dicom_path, ID, patient_directories)\n",
    "        images_list= read_prostatex_patient(ID, images_path, ktrans_path, verbose=True)\n",
    "\n",
    "        #Read Prostate segmentation mask\n",
    "        mask = sitk.ReadImage(os.path.join(masks_path, ID + '_msk.nrrd'))\n",
    "\n",
    "        #Read CZ segmentation mask\n",
    "        cz_mask= sitk.ReadImage(os.path.join(masks_path, ID + '_cz_msk.nrrd'))\n",
    "        \n",
    "        #Combine all images in images_list as single multichannel image using the first as reference\n",
    "        img_final= join_sitk_images(images_list, resampler=sitk.sitkBSpline, cast_type=sitk.sitkFloat32)\n",
    "        \n",
    "        #Join all masks\n",
    "        prostate_mask= join_masks(mask, cz_mask > 1.5, mode='append')\n",
    "\n",
    "        #---------------Lesion region growing preparation---------------\n",
    "        \n",
    "        #Load lesions information for current ID\n",
    "        try:\n",
    "            lesions= lesion_info.loc[ID].values\n",
    "            lesions= lesions[np.newaxis, ] if len(lesions.shape) == 1 else lesions\n",
    "            positions= np.array([np.fromstring(p[1], dtype=np.float32, sep=' ') for p in lesions])\n",
    "            positions_img= np.array([images_list[0].TransformPhysicalPointToContinuousIndex(p.astype(np.float64)) \n",
    "                                         for p in positions])\n",
    "            significances= [p[4] for p in lesions] if TRAIN else [20 for _ in lesions]\n",
    "            print(' - Lesion positions and significances:')\n",
    "            for i, (pos, sig) in enumerate(zip(positions_img, significances)):\n",
    "                print('   - %d: %s, Sig: %d'%(i+1, str(pos), sig))\n",
    "            \n",
    "        except Exception as e:\n",
    "            positions, positions_img, significances= [], [], []\n",
    "            print(' - Error: No lesion information found!')\n",
    "            raise e\n",
    "            \n",
    "        #We now create a small mask around the positions where lesions are located,\n",
    "        #which will be used as the seed to grow the lesions later\n",
    "        lesion_mask_id_seed= get_lesion_mask_id_seed(positions_img, mask)\n",
    "        prostate_mask_intermediate= join_masks(prostate_mask, lesion_mask_id_seed, mode='append')\n",
    "            \n",
    "        #---------------Get ROI and apply computed registration (if available)---------------\n",
    "               \n",
    "        #Rescale intensity (must be converted to numpy first)\n",
    "        img_backup= sitk.Image(img_final)\n",
    "        img_array= sitk.GetArrayFromImage(img_final)\n",
    "        img_array= rescale_intensity(img_array)\n",
    "        img_final= sitk.GetImageFromArray(img_array, isVector=True)\n",
    "        img_final.CopyInformation(img_backup)\n",
    "        \n",
    "        #Keep only the central area of size (160,160,24) around the prostate centroid \n",
    "        img_final, prostate_mask_intermediate= center_image(\n",
    "                                    img_final, prostate_mask_intermediate, center_around_roi=True,\n",
    "                                    size=(160,160,24), spacing=(0.5,0.5,3),\n",
    "                                    transform_channels=[1,2,3], per_channel_transform=transform,\n",
    "                                    pre_mask_growth_mm=2, pre_mask_growth_mm_channels= [2])\n",
    "        \n",
    "        #---------------Actual lesion region growing---------------\n",
    "        \n",
    "        #Create the automatic lesion segmentation mask\n",
    "        lesion_mask_id, _= grow_lesions( prostate_mask_intermediate, img_final, significances, transform, \n",
    "                                            iters_max=120, factors= [2.5,2.5,3.5,4] )\n",
    "        \n",
    "        #Join all masks\n",
    "        prostate_mask= join_masks(sitk.VectorIndexSelectionCast(prostate_mask_intermediate, 0), \n",
    "                                  sitk.VectorIndexSelectionCast(prostate_mask_intermediate, 1), mode='append')\n",
    "        prostate_mask= join_masks(prostate_mask, lesion_mask_id, mode='append')\n",
    "        \n",
    "        #---------------Plot results---------------\n",
    "        \n",
    "        #info(img_final)\n",
    "        for c in [0,3]:\n",
    "            plot_multi_mask(sitk.VectorIndexSelectionCast(img_final, c), prostate_mask, title='All masks')\n",
    "\n",
    "        #---------------Save them---------------\n",
    "        \n",
    "        #Generate needed arrays and information\n",
    "        img_final_arr= sitk.GetArrayFromImage(img_final)\n",
    "        prostate_mask_arr= sitk.GetArrayFromImage(prostate_mask)  \n",
    "\n",
    "        final_rois= prostate_mask_arr[...,2][..., np.newaxis]\n",
    "        img_arr= np.concatenate([img_final_arr, #C0-4 Image \n",
    "                                 prostate_mask_arr[...,[0,1]], #C5-6: Prostate & CZ mask\n",
    "                                 (prostate_mask_arr[...,0] - prostate_mask_arr[...,1])[...,np.newaxis]], #C7: PZ mask\n",
    "                                 axis=-1)\n",
    "        fg_slices= [ii for ii in np.unique(np.argwhere(final_rois != 0)[:, 0])]    \n",
    "\n",
    "        #Save all information: ID_rois.npy, ID_img.npy, meta_info_ID.pickle\n",
    "        np.save(os.path.join(output_path, '{}_rois.npy'.format(ID)), final_rois)\n",
    "        np.save(os.path.join(output_path, '{}_img.npy'.format(ID)), img_arr)\n",
    "        with open(os.path.join(output_path, 'meta_info_{}.pickle'.format(ID)), 'wb') as handle:\n",
    "            meta_info_dict = {'pid': ID, 'class_target': significances, \n",
    "                              'spacing': img_final.GetSpacing(), 'fg_slices': fg_slices}\n",
    "            pickle.dump(meta_info_dict, handle)\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(' - Error: Unhandled exception: %s'%e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Aggregated meta info to df with length 204\n"
     ]
    }
   ],
   "source": [
    "#Generate info_df.pickle for all the images\n",
    "files = [os.path.join(output_path, f) for f in os.listdir(output_path) if 'meta_info' in f]\n",
    "df = pd.DataFrame(columns=['pid', 'class_target', 'spacing', 'fg_slices'])\n",
    "for f in files:\n",
    "    with open(f, 'rb') as handle:\n",
    "        df.loc[len(df)] = pickle.load(handle)\n",
    "df.to_pickle(os.path.join(output_path, 'info_df.pickle'))\n",
    "print (\"Aggregated meta info to df with length\", len(df))"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}