Extract both CT+PT for 'deep' dataset and CT+PT+segmentation for 'radiomics' dataset

01_preprocessing/extract_roi_volume.py

@@ -11,25 +11,28 @@ import dicom_utils.dicom_utils as du
 import dicom_utils.dicom_utils.visualize as viz
 import gc
 
-from utils import sitk_to_numpy, filter_outbounds
 from config import get_project_root
 
-# os.chdir(PATH)
+from utils import sitk_to_numpy, filter_outbounds
+
 #%%
 PROJECT_ROOT = get_project_root()
 
-DATASET_NAME = 'HN_BZ'
+DATASET_NAME = 'HN_val'
 PROJECT_DATA_PATH = PROJECT_ROOT / 'data' / DATASET_NAME
 PROCESSED_DATA_PATH = PROJECT_DATA_PATH / 'processed'
 CLINICAL_DATA_FILENAME_CLEAN = f'clinical_{DATASET_NAME}.csv'
 
 DCM_DIR = PROCESSED_DATA_PATH / 'dcm'
 BBOX_OUT_DIR = PROCESSED_DATA_PATH / 'bbox'
-BBOX_SUBDATASET = 'bbox_64'
+BBOX_SUBDATASET_DEEP = 'bbox_64'
+BBOX_SUBDATASET_RAD = 'bbox_64_rad'
 
-OUTDIR = BBOX_OUT_DIR / BBOX_SUBDATASET
+OUTDIR_DEEP = BBOX_OUT_DIR / BBOX_SUBDATASET_DEEP
+OUTDIR_RAD = BBOX_OUT_DIR / BBOX_SUBDATASET_RAD
 
-os.makedirs(OUTDIR, exist_ok=False)
+os.makedirs(OUTDIR_DEEP, exist_ok=False)
+os.makedirs(OUTDIR_RAD, exist_ok=False)
 
 
 CT_ranges = [-1050, 3050]
@@ -94,22 +97,35 @@ for i, row in tqdm(clinical.iterrows()):
         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 = sitk.Resample(segmentation, scan_CT_box)
+
         # segmentation_box = gaussian.Execute(segmentation_box)
-        # out = segmentation_box>0.5
+        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 = np.stack([sitk_to_numpy(scan_CT_box), sitk_to_numpy(scan_PT_box)], axis=0)
+        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 / filename_out, out)
+        np.save(OUTDIR_DEEP / filename_out, out_deep)
+        np.save(OUTDIR_RAD / filename_out, out_rad)
 
-        del scan_CT_box, scan_PT_box, segmentation, out
+        del scan_CT_box, scan_PT_box, segmentation, mask_box, out_deep, out_rad
         gc.collect()
 
     except Exception as e: