Extract boxes 128x128x128 and resize to 64x64x64 for deep learning and save...

Extract boxes 128x128x128 and resize to 64x64x64 for deep learning and save three 128x128x128 images (CT+PT+mask) for radiomics

01_preprocessing/extract_roi_volume.py

 #%%
+import gc
 import os
-import numpy as np
-import SimpleITK as sitk
-import pandas as pd
-from pathlib import Path
 import sys
+from pathlib import Path
 
+import numpy as np
+import pandas as pd
+import SimpleITK as sitk
 from tqdm import tqdm
+
 import dicom_utils.dicom_utils as du
 import dicom_utils.dicom_utils.visualize as viz
-import gc
-
 from config import get_project_root
+from dicom_utils.dicom_utils import augmentation as aug
+from dicom_utils.dicom_utils import processing as dup
+from utils import filter_outbounds, sitk_to_numpy
+
+
+def resize(image, size):
+    N_CHANNELS = image.shape[0]
+    image_seq = [sitk.GetImageFromArray(image[i, :, :, :]) for i in range(N_CHANNELS)]
+    image_seq = [dup.resample(image, size=(size, size, size)) for image in image_seq]
+    image_seq = [sitk.GetArrayFromImage(image) for image in image_seq]
+    image = np.stack(image_seq, axis=0)
+    return image
 
-from utils import sitk_to_numpy, filter_outbounds
 
 #%%
 PROJECT_ROOT = get_project_root()
 
+BOX_SIZE = 128  # mm == pixels for voxel size of 1mm3
+RESIZE = True
+
 DATASET_NAME = 'HN_val'
 PROJECT_DATA_PATH = PROJECT_ROOT / 'data' / DATASET_NAME
 PROCESSED_DATA_PATH = PROJECT_DATA_PATH / 'processed'
@@ -25,8 +39,8 @@ CLINICAL_DATA_FILENAME_CLEAN = f'clinical_{DATASET_NAME}.csv'
 
 DCM_DIR = PROCESSED_DATA_PATH / 'dcm'
 BBOX_OUT_DIR = PROCESSED_DATA_PATH / 'bbox'
-BBOX_SUBDATASET_DEEP = 'bbox_64'
-BBOX_SUBDATASET_RAD = 'bbox_64_rad'
+BBOX_SUBDATASET_DEEP = f'bbox_{BOX_SIZE//2}_deep'
+BBOX_SUBDATASET_RAD = f'bbox_rad'
 
 OUTDIR_DEEP = BBOX_OUT_DIR / BBOX_SUBDATASET_DEEP
 OUTDIR_RAD = BBOX_OUT_DIR / BBOX_SUBDATASET_RAD
@@ -37,9 +51,9 @@ os.makedirs(OUTDIR_RAD, exist_ok=False)
 
 CT_ranges = [-1050, 3050]
 PT_ranges = [0, 50]
-BOX_SIZE = 64  # mm == pixels for voxel size of 1mm3
+
 VOXEL_SIZE = [1, 1, 1]
-HALF_BOX = BOX_SIZE // 2
+
 
 # gaussian = sitk.SmoothingRecursiveGaussianImageFilter()
 # gaussian.SetSigma(2)
@@ -57,11 +71,12 @@ errors = []
 #%%
 for i, row in tqdm(clinical.iterrows()):
 
+    patient = row['patient']
+    filename_out = row['filename']
+    roi_name = row['ROI_name']
+    roi_modality = row['ROI_modality']
+
     try:
-        patient = row['patient']
-        filename_out = row['filename']
-        roi_name = row['ROI_name']
-        roi_modality = row['ROI_modality']
 
         scan_CT = du.load_series(DCM_DIR / patient / 'CT')
         scan_PT = du.load_SUV(DCM_DIR / patient / 'PT')
@@ -69,7 +84,6 @@ for i, row in tqdm(clinical.iterrows()):
         rs_filename = os.listdir(DCM_DIR / patient / 'RTSTRUCT')[0]
 
         if roi_modality == 'CT':
-
             segmentation = du.load_roi(
                 DCM_DIR / patient / 'RTSTRUCT' / rs_filename, roi_name, scan_CT
             )
@@ -80,53 +94,49 @@ for i, row in tqdm(clinical.iterrows()):
         else:
             raise ValueError('ROI modality must be `CT` or `PT`.')
 
+        # Calculate segmentation mask vertices
         start_mm, stop_mm = du.get_bbox_vertices(segmentation)
         center_mm = (np.array(stop_mm) + np.array(start_mm)) / 2
 
+        # Calculate vertices for deep learning bbox centered at the center of the segmentation mask
         start_mm = (
-            center_mm - HALF_BOX - 5
+            center_mm - BOX_SIZE // 2 - 5
         )  # add margin to allow a better interpolation
-        stop_mm = center_mm + HALF_BOX + 5
+        stop_mm = center_mm + BOX_SIZE // 2 + 5
 
         scan_CT_box = du.extract_volume(scan_CT, start_mm, stop_mm)
+
         scan_CT_box = filter_outbounds(scan_CT_box, CT_ranges)
 
-        # upsample and register
         scan_CT_box = du.processing.resample(scan_CT_box, spacing=VOXEL_SIZE)
-
         scan_PT_box = sitk.Resample(scan_PT, scan_CT_box)
         scan_PT_box = filter_outbounds(scan_PT_box, PT_ranges)
 
-        segmentation_box = sitk.Resample(segmentation, scan_CT_box)
-
-        # segmentation_box = gaussian.Execute(segmentation_box)
-        mask_box = segmentation_box > 0.5
-
         start_mm = start_mm + 5  # remove margin
         stop_mm = stop_mm - 5  # remove margin
 
         scan_CT_box = du.extract_volume(scan_CT_box, start_mm, stop_mm)
         scan_PT_box = du.extract_volume(scan_PT_box, start_mm, stop_mm)
-        mask_box = du.extract_volume(mask_box, start_mm, stop_mm)
 
         out_deep = np.stack(
             [sitk_to_numpy(scan_CT_box), sitk_to_numpy(scan_PT_box)], axis=0
         )
-        out_rad = np.stack(
-            [
-                sitk_to_numpy(scan_CT_box),
-                sitk_to_numpy(scan_PT_box),
-                sitk_to_numpy(mask_box),
-            ],
-            axis=0,
-        )
 
-        # save
-        np.save(OUTDIR_DEEP / filename_out, out_deep)
-        np.save(OUTDIR_RAD / filename_out, out_rad)
+        out_deep_resized = resize(out_deep, size=BOX_SIZE // 2)
+        # Save
+        np.save(OUTDIR_DEEP / filename_out, out_deep_resized)
+
+        segmentation_box = sitk.Resample(segmentation, scan_CT_box)
+        # segmentation_box = gaussian.Execute(segmentation_box)
+        mask_box = segmentation_box > 0.5
+        mask_box = du.extract_volume(mask_box, start_mm, stop_mm)
+
+        # Save
+        os.makedirs(OUTDIR_RAD / patient, exist_ok=False)
 
-        del scan_CT_box, scan_PT_box, segmentation, mask_box, out_deep, out_rad
-        gc.collect()
+        sitk.WriteImage(scan_CT_box, str(OUTDIR_RAD / patient / 'CT.nrrd'))
+        sitk.WriteImage(scan_PT_box, str(OUTDIR_RAD / patient / 'PT.nrrd'))
+        sitk.WriteImage(mask_box, str(OUTDIR_RAD / patient / 'mask.nrrd'))
 
     except Exception as e:
         print(e)