extract_roi_volume.py 4.86 KB
Newer Older
Nicole Bussola's avatar
Nicole Bussola committed
1
# %%
2
import gc
3
4
import os
import sys
5
from pathlib import Path
6

7
8
9
import numpy as np
import pandas as pd
import SimpleITK as sitk
10
from tqdm import tqdm
11

12
13
import dicom_utils.dicom_utils as du
import dicom_utils.dicom_utils.visualize as viz
14
from config import get_project_root
15
16
17
18
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

Nicole Bussola's avatar
Nicole Bussola committed
19
20
from matplotlib import pyplot as plt

21
22
23
24
25
26
27
28

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
29

30

Nicole Bussola's avatar
Nicole Bussola committed
31
# %%
32
33
PROJECT_ROOT = get_project_root()

34
35
36
BOX_SIZE = 128  # mm == pixels for voxel size of 1mm3
RESIZE = True

37
DATASET_NAME = 'HN_val'
38
PROJECT_DATA_PATH = PROJECT_ROOT / 'data' / DATASET_NAME
39
40
41
42
43
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'
44
45
BBOX_SUBDATASET_DEEP = f'bbox_{BOX_SIZE//2}_deep'
BBOX_SUBDATASET_RAD = f'bbox_rad'
46

47
48
OUTDIR_DEEP = BBOX_OUT_DIR / BBOX_SUBDATASET_DEEP
OUTDIR_RAD = BBOX_OUT_DIR / BBOX_SUBDATASET_RAD
49

Nicole Bussola's avatar
Nicole Bussola committed
50
51
os.makedirs(OUTDIR_DEEP, exist_ok=True)
os.makedirs(OUTDIR_RAD, exist_ok=True)
52
53
54
55


CT_ranges = [-1050, 3050]
PT_ranges = [0, 50]
56

57
VOXEL_SIZE = [1, 1, 1]
58

59
60
61
62

# gaussian = sitk.SmoothingRecursiveGaussianImageFilter()
# gaussian.SetSigma(2)

Nicole Bussola's avatar
Nicole Bussola committed
63
# %%
64
65
66
67
68
69
70
71
72
clinical = pd.read_csv(PROCESSED_DATA_PATH / CLINICAL_DATA_FILENAME_CLEAN)

available_patients = patients = [
    f for f in os.listdir(DCM_DIR) if os.path.isdir(DCM_DIR / f)
]

clinical = clinical.loc[clinical['patient'].isin(available_patients)]

errors = []
Nicole Bussola's avatar
Nicole Bussola committed
73
74
# %%
for i, row in clinical.iterrows():
75

76
77
78
79
80
    patient = row['patient']
    filename_out = row['filename']
    roi_name = row['ROI_name']
    roi_modality = row['ROI_modality']

81
82
83
    try:

        scan_CT = du.load_series(DCM_DIR / patient / 'CT')
Nicole Bussola's avatar
Nicole Bussola committed
84
85
        # scan_PT = du.load_SUV(DCM_DIR / patient / 'PT')
        scan_PT = du.load_series(DCM_DIR / patient / 'PT')
86
87
88
89
90
91
92
93
94
95
96
        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
            )
        elif roi_modality == 'PT':
            segmentation = du.load_roi(
                DCM_DIR / patient / 'RTSTRUCT' / rs_filename, roi_name, scan_PT
            )
        else:
            raise ValueError('ROI modality must be `CT` or `PT`.')
97
        # Calculate segmentation mask vertices
Nicole Bussola's avatar
Nicole Bussola committed
98
        print(type(segmentation))
99
        start_mm, stop_mm = du.get_bbox_vertices(segmentation)
Nicole Bussola's avatar
Nicole Bussola committed
100
101
        # print(start_mm, stop_mm)

102
103
        center_mm = (np.array(stop_mm) + np.array(start_mm)) / 2

104
        # Calculate vertices for deep learning bbox centered at the center of the segmentation mask
105
        start_mm = (
106
            center_mm - BOX_SIZE // 2 - 5
107
        )  # add margin to allow a better interpolation
108
        stop_mm = center_mm + BOX_SIZE // 2 + 5
109
110

        scan_CT_box = du.extract_volume(scan_CT, start_mm, stop_mm)
Nicole Bussola's avatar
Nicole Bussola committed
111
        print(scan_CT_box.GetSize())
112

Nicole Bussola's avatar
Nicole Bussola committed
113
        # scan_CT_box = filter_outbounds(scan_CT_box, CT_ranges)
114
115

        scan_CT_box = du.processing.resample(scan_CT_box, spacing=VOXEL_SIZE)
Nicole Bussola's avatar
Nicole Bussola committed
116
117
        print(scan_CT_box.GetSize())

118
        scan_PT_box = sitk.Resample(scan_PT, scan_CT_box)
Nicole Bussola's avatar
Nicole Bussola committed
119
        #  scan_PT_box = filter_outbounds(scan_PT_box, PT_ranges)
120
121
122
123
124

        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)
Nicole Bussola's avatar
Nicole Bussola committed
125
126
        print(scan_CT_box.GetSize())

127
        scan_PT_box = du.extract_volume(scan_PT_box, start_mm, stop_mm)
128
129
130
131

        out_deep = np.stack(
            [sitk_to_numpy(scan_CT_box), sitk_to_numpy(scan_PT_box)], axis=0
        )
132

133
134
135
136
        out_deep_resized = resize(out_deep, size=BOX_SIZE // 2)
        # Save
        np.save(OUTDIR_DEEP / filename_out, out_deep_resized)

Nicole Bussola's avatar
Nicole Bussola committed
137
        # segmentation_box = sitk.Resample(segmentation, scan_CT_box)
138
        # segmentation_box = gaussian.Execute(segmentation_box)
Nicole Bussola's avatar
Nicole Bussola committed
139
140
141
        # mask_box = segmentation_box > 0.5
        # print(type(mask_box))
        # mask_box = du.extract_volume(mask_box, start_mm, stop_mm)
142
143

        # Save
Nicole Bussola's avatar
Nicole Bussola committed
144
145
146
147
148
        # os.makedirs(OUTDIR_RAD / patient, exist_ok=False)
    #
    # 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'))
149
150

    except Exception as e:
Nicole Bussola's avatar
Nicole Bussola committed
151
152
        print(f'\nError processing patient: {row["patient"]}')
        print('Exception:', e)
153
154
        errors.append(row['patient'])

Nicole Bussola's avatar
Nicole Bussola committed
155
# %%
156
print(errors)