Resampling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: SimpleITK in /opt/anaconda/lib/python3.7/site-packages (1.2.0)\n",
      "Requirement already satisfied: pydicom in /opt/anaconda/lib/python3.7/site-packages (1.2.2)\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "The code here:\n",
    "    1. resamples CT and PET volumes to match a fixed size \n",
    "    2. isolates the GTV segmentation in RTSS files\n",
    "    3. resamples labelmaps which were converted from the RTSS files using the BatchStructureSetConversion module in 3D slicer\n",
    "    4. converts the CT/PET volumes and labelmaps into numpy arrays\n",
    "'''\n",
    "\n",
    "!pip install SimpleITK\n",
    "!pip install pydicom\n",
    "\n",
    "import pydicom\n",
    "import SimpleITK as sitk\n",
    "import os\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import gzip\n",
    "import shutil\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num ct structs: 194\n",
      "num pet structs: 194\n",
      "num ct: 194\n",
      "num pet: 194\n"
     ]
    }
   ],
   "source": [
    "# Trawl through files system for RTSS and volume files\n",
    "\n",
    "def get_processed_dir(path):\n",
    "    return path.replace('Head-Neck-PET-CT', 'Head-Neck-PET-CT-Processed')\n",
    "\n",
    "def get_unprocessed_dir(path):\n",
    "    return path.replace('Head-Neck-PET-CT-Processed', 'Head-Neck-PET-CT')\n",
    "\n",
    "ct_structs = []\n",
    "pet_structs = []\n",
    "ct_vols = set()\n",
    "pet_vols = set()\n",
    "\n",
    "for root, dirs, files in os.walk('/home/jzhe0882/datasets/Head-Neck-PET-CT'):\n",
    "    for name in files:\n",
    "        file_path = os.path.join(root, name)\n",
    "        \n",
    "        relevant_dir = True #whether the directory contains a file to be processed\n",
    "        if 'RTstructCTsim-CTPET-CT' in root:\n",
    "            ct_structs.append(file_path) # CT RTSS file\n",
    "            \n",
    "        elif 'RTstructCTsim-PETPET-CT' in root:\n",
    "            pet_structs.append(file_path) # PET RTSS file\n",
    "            \n",
    "        elif os.path.getsize(file_path) < 200000: \n",
    "            pet_vols.add(root) # Folder containing series of PET images\n",
    "\n",
    "        elif os.path.getsize(file_path) < 1000000: \n",
    "            ct_vols.add(root) # Folder containing series of CT images\n",
    "            \n",
    "        else: #otherwise it is an RTDose file\n",
    "            relevant_dir = False\n",
    "            \n",
    "        #create the relevant directories if they don't exist\n",
    "        processed_dir = get_processed_dir(root)\n",
    "        if relevant_dir and not os.path.isdir(processed_dir):\n",
    "            os.makedirs(processed_dir) \n",
    "            \n",
    "#all should be equal\n",
    "print('num ct structs:', len(ct_structs))\n",
    "print('num pet structs:', len(pet_structs))\n",
    "print('num ct:', len(ct_vols))\n",
    "print('num pet:', len(pet_vols))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vol2rtss size: 388\n"
     ]
    }
   ],
   "source": [
    "# Associate each CT/PET volume with their RTSS file\n",
    "\n",
    "vol2rtss = {} # associates each CT/PET volume directory with its paired RTSS' file path\n",
    "\n",
    "for pet_path in pet_vols:    \n",
    "    base_path = os.path.abspath(os.path.join(pet_path, '..'))\n",
    "    for struct_path in pet_structs:\n",
    "        if base_path in struct_path:\n",
    "            vol2rtss[pet_path] = struct_path\n",
    "            break\n",
    "    else:\n",
    "        print('RTSS not found for', root)\n",
    "        \n",
    "for ct_path in ct_vols:    \n",
    "    base_path = os.path.abspath(os.path.join(ct_path, '..'))\n",
    "    for struct_path in ct_structs:\n",
    "        if base_path in struct_path:\n",
    "            vol2rtss[ct_path] = struct_path\n",
    "            break\n",
    "    else:\n",
    "        print('RTSS not found for', root)\n",
    "        \n",
    "print('vol2rtss size:', len(vol2rtss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "set()\n"
     ]
    }
   ],
   "source": [
    "# Remove all segmentations except for those corresponding to GTV's\n",
    "# Associate the segmentations with the newly resampled CT/PET scans\n",
    "\n",
    "#reference: https://dicom.innolitics.com/ciods/rt-structure-set/roi-contour/30060039/30060084\n",
    "\n",
    "#check for potential gtv tag names if they aren't contained in gtv_names\n",
    "missing_types = set()\n",
    "\n",
    "#these are obtained from the names of gtv contours in https://wiki.cancerimagingarchive.net/display/Public/Head-Neck-PET-CT#5395a20070db41d49745c83e3312ff86\n",
    "gtv_names = {'GTV', 'GTV1', 'GTV-P', 'GTV 67', 'gtv pet', 'GTV-BOT', 'GTV 67.5', 'GTV 67.5Gy', 'GTV primary_70GY',\n",
    "           'GTV_67.5gy', 'GTV_70Gy', 'GTV_69Gy', 'GTVt', 'GTV primaire', 'GTV p', 'GTV Primaire 70', 'GTV T irm',\n",
    "           'GTV 70', 'GTV et GG', 'GTV n', 'GTV_70gy', 'GTV_70GY', 'GTV_67.5', 'GTV-p', 'GTV_67.5Gy', 'GTV_67.5GY',\n",
    "            'GTV_T_67.5gy'}\n",
    "\n",
    "def configure_rtss(file, slice_sop_uids, patient_id):\n",
    "    ds = pydicom.read_file(file)\n",
    "    \n",
    "    #there are cases where the patient id/name of the RTSS files are mismatched, so we use the id from the PET/CT volume metadata\n",
    "    #patient_id = patient_name in this dataset\n",
    "    ds.PatientID = patient_id\n",
    "    ds.PatientName = patient_id\n",
    "    \n",
    "    gtv_num = None\n",
    "    \n",
    "    #print(dir(ds))\n",
    "    #print(file)\n",
    "    #print(ds.ReferencedFrameOfReferenceSequence)\n",
    "    \n",
    "    #Assign new SOP UID's for each of the slices\n",
    "    for referenced_frame_of_reference in ds.ReferencedFrameOfReferenceSequence:       \n",
    "        for RT_referenced_study in referenced_frame_of_reference.RTReferencedStudySequence:       \n",
    "            for RT_referenced_series in RT_referenced_study.RTReferencedSeriesSequence:\n",
    "                \n",
    "                contour_image_sequence = RT_referenced_series.ContourImageSequence\n",
    "                sop_class = contour_image_sequence[0].ReferencedSOPClassUID\n",
    "                contour_image_sequence.clear()\n",
    "                                \n",
    "                for sop_uid in slice_sop_uids:\n",
    "                    element = pydicom.Dataset()\n",
    "                    element.ReferencedSOPClassUID = sop_class\n",
    "                    element.ReferencedSOPInstanceUID = sop_uid\n",
    "                    contour_image_sequence.append(element)\n",
    "        break\n",
    "    \n",
    "    #look for the id number for the gross tumour volume tag\n",
    "    for contour_type in ds.StructureSetROISequence:\n",
    "        if contour_type.ROIName in gtv_names:\n",
    "            gtv_num = contour_type.ROINumber\n",
    "            break\n",
    "     \n",
    "    #shouldn't happen\n",
    "    if gtv_num is None:\n",
    "        \n",
    "        for contour_type in ds.StructureSetROISequence:\n",
    "            missing_types.add(contour_type.ROIName)\n",
    "        \n",
    "        print('GTV not found for', file)\n",
    "        \n",
    "    #remove all segmentations except the GTV segmentation\n",
    "    else:       \n",
    "        gtv_contours = []\n",
    "        \n",
    "        for contour in ds.ROIContourSequence:\n",
    "            if contour.ReferencedROINumber == gtv_num:\n",
    "                gtv_contours.append(contour)\n",
    "                \n",
    "        ds.ROIContourSequence.clear()\n",
    "                \n",
    "        for contour in gtv_contours:\n",
    "            ds.ROIContourSequence.append(contour)\n",
    "            \n",
    "        ds.save_as(get_processed_dir(file))\n",
    "        \n",
    "#should be empty\n",
    "print(missing_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "resampled 20\n",
      "resampled 40\n",
      "resampled 60\n",
      "resampled 80\n",
      "resampled 100\n",
      "resampled 120\n",
      "resampled 140\n",
      "resampled 160\n",
      "resampled 180\n",
      "resampled 200\n",
      "resampled 220\n",
      "resampled 240\n",
      "resampled 260\n",
      "resampled 280\n",
      "resampled 300\n",
      "resampled 320\n",
      "resampled 340\n",
      "resampled 360\n",
      "resampled 380\n",
      "resampled 388\n"
     ]
    }
   ],
   "source": [
    "# Resample the CT and PET volumes\n",
    "\n",
    "#Read-writing DICOM files is a delicate operation: https://simpleitk.readthedocs.io/en/master/Examples/DicomSeriesReadModifyWrite/Documentation.html\n",
    "\n",
    "def resample_volume(path, ct=True):    \n",
    "    target_ct_size = (128,128,96)\n",
    "    target_pet_size = (128,128,96)\n",
    "    \n",
    "    reader = sitk.ImageSeriesReader()\n",
    "    files = reader.GetGDCMSeriesFileNames(path)\n",
    "    reader.SetFileNames(files)\n",
    "    reader.MetaDataDictionaryArrayUpdateOn() #load tags\n",
    "    reader.LoadPrivateTagsOn() #load private tags (these aren't loaded by default)\n",
    "    volume = reader.Execute()\n",
    "    \n",
    "    #keep everything the same except for the spacing and size\n",
    "    #'size' here refers to number of slices along each dimension\n",
    "    #hence, total volume = spacing * size\n",
    "    spacing = volume.GetSpacing()\n",
    "    size = volume.GetSize()\n",
    "    \n",
    "    #In Python, variables introduced in if-else blocks preserve scope\n",
    "    #Wtf\n",
    "    if ct:\n",
    "        new_size = target_ct_size\n",
    "        new_spacing = ((size[0] / target_ct_size[0]) * spacing[0], \n",
    "                       (size[1] / target_ct_size[1]) * spacing[1], \n",
    "                       (size[2] / target_ct_size[2]) * spacing[2])\n",
    "    else:\n",
    "        new_size = target_pet_size\n",
    "        new_spacing = ((size[0] / target_pet_size[0]) * spacing[0], \n",
    "                       (size[1] / target_pet_size[1]) * spacing[1], \n",
    "                       (size[2] / target_pet_size[2]) * spacing[2])\n",
    "    \n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetInterpolator(sitk.sitkLinear)\n",
    "    resampler.SetReferenceImage(volume)\n",
    "    resampler.SetSize(new_size)\n",
    "    resampler.SetOutputSpacing(new_spacing)\n",
    "    \n",
    "    new_volume = resampler.Execute(volume)\n",
    "        \n",
    "    writer = sitk.ImageFileWriter()  \n",
    "    \n",
    "    #by default, the image writer does not preserve meta-data from original DICOM series; we copy the following tags across\n",
    "    writer.KeepOriginalImageUIDOn() \n",
    "    tags_to_copy = ['0002|0002', # Media Storage SOP Class UID\n",
    "                '0008|0016', # SOP Class UID\n",
    "                '0008|0022', # Acquisition Date\n",
    "                '0008|0023', # Content Date\n",
    "                '0008|0032', # Acquisition Time\n",
    "                '0008|0033', # Content Time\n",
    "                '0020|0012', # Acquisition Number\n",
    "                \"0010|0020\", # Patient ID\n",
    "                \"0010|0030\", # Patient Birth Date\n",
    "                \"0020|000d\", # Study Instance UID, for machine consumption\n",
    "                \"0020|0010\", # Study ID, for human consumption\n",
    "                \"0008|0020\", # Study Date\n",
    "                \"0008|0030\", # Study Time\n",
    "                \"0008|0050\", # Accession Number\n",
    "                \"0008|0060\", # Modality\n",
    "                '0028|1050', # Window Centre\n",
    "                '0028|1051' # Window Width    \n",
    "                ]\n",
    "        \n",
    "    modification_time = time.strftime(\"%H%M%S\")\n",
    "    modification_date = time.strftime(\"%Y%m%d\")\n",
    "        \n",
    "    tag_values = [(k, reader.GetMetaData(0,k)) for k in tags_to_copy if reader.HasMetaDataKey(0,k)] + \\\n",
    "                 [(\"0010|0010\", reader.GetMetaData(0,\"0010|0020\")), # Patient ID is assigned to Patient Name \n",
    "                  (\"0008|0031\",modification_time), # Series Time\n",
    "                  (\"0008|0021\",modification_date), # Series Date\n",
    "                  (\"0008|0008\",\"DERIVED\\\\SECONDARY\"), # Image Type\n",
    "                  ('0018|0050', str(new_spacing[2])), # Slice Thickness\n",
    "                  ('0028|0030', '{}/{}'.format(new_spacing[0], new_spacing[1])), # Pixel Spacing\n",
    "                  (\"0020|000e\", \"1.2.826.0.1.3680043.2.1125.\"+modification_date+\".1\"+modification_time), # Series Instance UID\n",
    "                  (\"0008|103e\", reader.GetMetaData(0,\"0008|103e\") + \" Processed-SimpleITK\")] # Series Description\n",
    "    \n",
    "    #Some volumes need to cast from float to int before writing to the file system\n",
    "    #Otherwise this error is thrown: A Floating point buffer was passed but the stored pixel type was not specified.\n",
    "    castFilter = sitk.CastImageFilter()\n",
    "    castFilter.SetOutputPixelType(sitk.sitkInt32)\n",
    "    new_volume = castFilter.Execute(new_volume)\n",
    "    \n",
    "    #Write DICOM series while assigning metadata values\n",
    "    save_directory = get_processed_dir(path)\n",
    "    sop_uids = []\n",
    "    for i in range(new_size[2]):\n",
    "        image_slice = new_volume[:,:,i]\n",
    "                \n",
    "        #tags shared across all slices\n",
    "        for tag, value in tag_values:\n",
    "            image_slice.SetMetaData(tag, value)\n",
    "        \n",
    "        #slice-specific tags\n",
    "        image_slice.SetMetaData(\"0008|0012\", time.strftime(\"%Y%m%d\")) # Instance Creation Date\n",
    "        image_slice.SetMetaData(\"0008|0013\", time.strftime(\"%H%M%S\")) # Instance Creation Time\n",
    "        image_slice.SetMetaData(\"0020|0032\", '\\\\'.join(map(str,new_volume.TransformIndexToPhysicalPoint((0,0,i))))) # Image Position (Patient)\n",
    "        image_slice.SetMetaData(\"0020|0013\", str(i)) # Instance Number\n",
    "        \n",
    "        file_path = os.path.join(save_directory, '{0:06d}.dcm'.format(i))\n",
    "        writer.SetFileName(file_path)\n",
    "        writer.Execute(image_slice)\n",
    "        \n",
    "        #The SITK library can't read hierarchical data; we use pydicom to assign Referenced SOP Instance UID values\n",
    "        #This is needed so that 3D Slicer recognises that this CT/PET scan is the source volume for RTSS segmentations\n",
    "        source_volume = pydicom.read_file(files[0]) #the source volume can be any image slice as the following attributes are shared across all slices\n",
    "        ds = pydicom.read_file(file_path)\n",
    "        ds.SeriesInstanceUID = source_volume.SeriesInstanceUID\n",
    "        ds.FrameOfReferenceUID = source_volume.FrameOfReferenceUID;\n",
    "        if 'ReferencedStudySequence' in source_volume: #this field is sometimes omitted in some volumes\n",
    "            ds.ReferencedStudySequence = source_volume.ReferencedStudySequence\n",
    "\n",
    "        sop_uids.append(ds.SOPInstanceUID)\n",
    "        ds.save_as(file_path)\n",
    "    \n",
    "    patient_id = reader.GetMetaData(0,\"0010|0020\")\n",
    "    configure_rtss(vol2rtss[path], sop_uids, patient_id)\n",
    "    \n",
    "    #print('completed {}'.format(save_directory))\n",
    "    #print('Size:', volume.GetSize(), '->', new_volume.GetSize())\n",
    "    #print('Spacing:', volume.GetSpacing(), '->', new_volume.GetSpacing())\n",
    "        \n",
    "        \n",
    "#resample_volume(path='/home/jzhe0882/datasets/Head-Neck-PET-CT/HN-CHUM-007/08-27-1885-CT pour planification -ONCO-09347/837620-LOR-RAMLA-92958', ct=False)\n",
    "\n",
    "resampling_progress = 0\n",
    "for ct_path in ct_vols:\n",
    "    resample_volume(ct_path, ct=True)\n",
    "    \n",
    "    resampling_progress += 1\n",
    "    if resampling_progress % 20 == 0:\n",
    "        print('resampled', resampling_progress)\n",
    "        \n",
    "for pet_path in pet_vols:\n",
    "    resample_volume(pet_path, ct=False)\n",
    "    \n",
    "    resampling_progress += 1\n",
    "    if resampling_progress % 20 == 0:\n",
    "        print('resampled', resampling_progress)\n",
    "        \n",
    "print('resampled', resampling_progress)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mask not found in /home/jzhe0882/datasets/LabelMaps\n",
      "mask2pet: 194\n",
      "mask2ct: 194\n"
     ]
    }
   ],
   "source": [
    "# Scan file system for masks generated by 3D Slicer and associate them with corresponding CT/PET volumes\n",
    "\n",
    "def get_processed_mask_dir(path):\n",
    "    return path.replace('LabelMaps', 'LabelMaps-Processed')\n",
    "\n",
    "mask2pet = {} #maps masks to resampled pet images' directories\n",
    "mask2ct = {} #maps masks to resampled ct images' directories\n",
    "\n",
    "for root, dirs, files in os.walk('/home/jzhe0882/datasets/LabelMaps'):\n",
    "    patient = os.path.basename(root) #e.g. HN-CHUM-001\n",
    "    \n",
    "    #choose the labelmap with the largest filesize to be the mask for both the CT and PET volume\n",
    "    mask = None\n",
    "    max_size = 0\n",
    "    for name in files:\n",
    "        file_path = os.path.join(root, name)\n",
    "        file_size = os.path.getsize(file_path)\n",
    "        if file_size > max_size:\n",
    "            max_size = file_size\n",
    "            mask = file_path\n",
    "            \n",
    "        #create the directory to save the resampled label maps \n",
    "        processed_dir = get_processed_mask_dir(root)\n",
    "        if not os.path.isdir(processed_dir):\n",
    "            os.makedirs(processed_dir) \n",
    "            \n",
    "    #map the ct and pet volume to the mask\n",
    "    if mask is not None:\n",
    "        for volume in ct_vols:\n",
    "            if patient in volume:\n",
    "                mask2ct[mask] = get_processed_dir(volume)\n",
    "                break\n",
    "                \n",
    "        for volume in pet_vols:\n",
    "            if patient in volume:\n",
    "                mask2pet[mask] = get_processed_dir(volume)\n",
    "                break\n",
    "    else:\n",
    "        print('mask not found in', root)\n",
    "                \n",
    "#should be equal\n",
    "print('mask2pet:', len(mask2pet))\n",
    "print('mask2ct:', len(mask2ct))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resample_mask(mask_dir, reference_dir):\n",
    "    \n",
    "    #load mask volume\n",
    "    mask_reader = sitk.ImageFileReader()\n",
    "    mask_reader.SetFileName(mask_dir)\n",
    "    mask_volume = mask_reader.Execute()\n",
    "    \n",
    "    #load ct/pet volume\n",
    "    reference_reader = sitk.ImageSeriesReader()\n",
    "    reference_files = reference_reader.GetGDCMSeriesFileNames(reference_dir)\n",
    "    reference_reader.SetFileNames(reference_files)\n",
    "    #reference_reader.MetaDataDictionaryArrayUpdateOn() #load tags\n",
    "    #reference_reader.LoadPrivateTagsOn() #load private tags (these aren't loaded by default)\n",
    "    reference_volume = reference_reader.Execute()\n",
    "    \n",
    "    #resample mask volume to match reference image parameters\n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetInterpolator(sitk.sitkNearestNeighbor)\n",
    "    resampler.SetReferenceImage(reference_volume)\n",
    "    new_mask = resampler.Execute(mask_volume)\n",
    "    \n",
    "    '''castFilter = sitk.CastImageFilter()\n",
    "    castFilter.SetOutputPixelType(sitk.sitkInt32)\n",
    "    new_mask = castFilter.Execute(new_mask)'''\n",
    "    \n",
    "    writer = sitk.ImageFileWriter()  \n",
    "    writer.KeepOriginalImageUIDOn()\n",
    "    writer.SetFileName(get_processed_mask_dir(mask_dir))\n",
    "    writer.Execute(new_mask)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "resampled 20\n",
      "resampled 40\n",
      "resampled 60\n",
      "resampled 80\n",
      "resampled 100\n",
      "resampled 120\n",
      "resampled 140\n",
      "resampled 160\n",
      "resampled 180\n",
      "resampled 194\n"
     ]
    }
   ],
   "source": [
    "# Resample label maps that were converted from the RTSS files using 3D slicer\n",
    "\n",
    "resampling_progress = 0\n",
    "\n",
    "#use the PET images for mask generation, as it has lower resolution\n",
    "for mask, volume in mask2pet.items():\n",
    "    resample_mask(mask, volume)\n",
    "    \n",
    "    resampling_progress += 1\n",
    "    if resampling_progress % 20 == 0:\n",
    "        print('resampled', resampling_progress)\n",
    "        \n",
    "print('resampled', resampling_progress)\n",
    "#resample_mask('/home/jzhe0882/datasets/LabelMaps/HN-CHUM-007/1  RTSTRUCT  RTstruct_1_GTV.nrrd',\n",
    "#             '/home/jzhe0882/datasets/Head-Neck-PET-CT/HN-CHUM-007/08-27-1885-CT pour planification -ONCO-09347/837620-LOR-RAMLA-92958')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converted 20\n",
      "converted 40\n",
      "converted 60\n",
      "converted 80\n",
      "converted 100\n",
      "converted 120\n",
      "converted 140\n",
      "converted 160\n",
      "converted 180\n",
      "converted 194\n"
     ]
    }
   ],
   "source": [
    "# Convert PET/CT volumes and segmentation masks to numpy arrays\n",
    "\n",
    "def convert(source_path, target_path, source_type='volume', target_type='numpy'):\n",
    "    \n",
    "    reader = None\n",
    "    if source_type == 'volume':\n",
    "        reader = sitk.ImageSeriesReader()\n",
    "        file_names = reader.GetGDCMSeriesFileNames(source_path)\n",
    "        reader.SetFileNames(file_names)\n",
    "    elif source_type == 'single_image':\n",
    "        reader = sitk.ImageFileReader()\n",
    "        reader.SetFileName(source_path)\n",
    "    else:\n",
    "        print('incorrect source type')\n",
    "        return\n",
    "        \n",
    "    volume = reader.Execute()\n",
    "    \n",
    "    if target_type == 'numpy':\n",
    "        target_path = '{}.npy'.format(target_path)\n",
    "        array = sitk.GetArrayFromImage(volume)\n",
    "        np.save(target_path, np.transpose(array))\n",
    "    elif target_type == 'nifti':\n",
    "        target_path = '{}.nii'.format(target_path)\n",
    "        writer = sitk.ImageFileWriter()\n",
    "        writer.SetFileName(target_path)\n",
    "        writer.Execute(volume)\n",
    "        \n",
    "        with open(target_path, 'rb') as f_in: #compress into .gz file as 3D-DSN code only accepts .nii.gz files :/\n",
    "            with gzip.open('{}.gz'.format(target_path), 'wb') as f_out:\n",
    "                shutil.copyfileobj(f_in, f_out)\n",
    "    \n",
    "    \n",
    "index = 0\n",
    "for mask, pet in mask2pet.items():\n",
    "    ct = mask2ct[mask]\n",
    "    mask = get_processed_mask_dir(mask)\n",
    "    \n",
    "    #e.g. HN-CHUM-001\n",
    "    patient = os.path.basename(os.path.abspath(os.path.join(ct, '../..')))\n",
    "    \n",
    "    #numpy conversion\n",
    "    numpy_ct_path = os.path.join('/home/jzhe0882/numpydata/CT', patient)\n",
    "    numpy_pet_path = os.path.join('/home/jzhe0882/numpydata/PET', patient)\n",
    "    numpy_mask_path = os.path.join('/home/jzhe0882/numpydata/Mask', patient)\n",
    "    \n",
    "    convert(ct, numpy_ct_path, source_type='volume', target_type='numpy')\n",
    "    convert(pet, numpy_pet_path, source_type='volume', target_type='numpy')\n",
    "    convert(mask, numpy_mask_path, source_type='single_image', target_type='numpy')\n",
    "    \n",
    "    #also convert files to a NIFTI format for other frameworks; VERY SLOW BLYAT\n",
    "    nifti_ct_path = os.path.join('/home/jzhe0882/datasets/Nifti', 'training_axial_crop_pat{}'.format(index))\n",
    "    #nifti_pet_path = os.path.join('/home/jzhe0882/datasets/Nifti', 'training_axial_crop_pat{}'.format(index))\n",
    "    nifti_mask_path = os.path.join('/home/jzhe0882/datasets/Nifti', 'training_axial_crop_pat{}-label'.format(index))\n",
    "    \n",
    "    convert(ct, nifti_ct_path, source_type='volume', target_type='nifti')\n",
    "    #convert(pet, nifti_pet_path, source_type='volume', target_type='nifti')\n",
    "    convert(mask, nifti_mask_path, source_type='single_image', target_type='nifti')\n",
    "    \n",
    "    index += 1\n",
    "    if index % 20 == 0:\n",
    "        print('converted', index)\n",
    "        \n",
    "print('converted', index)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}