preprocessing/stage2.py at main · SynthRAD2025/preprocessing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import preprocessing_utils as utils
import os
import SimpleITK as sitk
import logging
import sys

# Set this to true if you want to skip already pre-processsed patients (checks if output files already exists)
skip_existing = True

if __name__ == "__main__":
    ## set up logging to console and file
    log_file = 'stage2.log'
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    file_handler = logging.FileHandler(log_file,mode = 'a')
    file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    logger.info('Starting stage 2 preprocessing')

    ## load pre-processing configuration from .csv file (sys.argv[1])
    file = sys.argv[1]
    if file.endswith('.csv'):
        patient_dict = utils.csv_to_dict(file)
    else:
        logger.error('Input file must be a csv file')
        sys.exit(1)
    for i in patient_dict:
        patient = patient_dict[i]
        # check if output files already exist and skip if flag is set
        if skip_existing:
            if patient['task'] == 1:
                if (os.path.isfile(os.path.join(patient['output_dir'],'mr_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'ct_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'ct_s2_def.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'mask_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],f'{patient["ID"]}.png'))):
                    logger.info(f'Patient {i} already pre-processed. Skipping...')
                    continue
            elif patient['task'] == 2:
                if (os.path.isfile(os.path.join(patient['output_dir'],'cbct_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'ct_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'ct_s2_def.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],'mask_s2.nii.gz')) and
                    os.path.isfile(os.path.join(patient['output_dir'],f'{patient["ID"]}.png'))):
                    logger.info(f'Patient {i} already pre-processed. Skipping...')
                    continue

        # log patient details
        logger.info(f'''Processing case {i}:
                            Output_dir: {patient['output_dir']}''')

        #Read Files
        if patient['task'] == 1:
            input = utils.read_image(os.path.join(patient['output_dir'],'mr_s1.nii.gz'),log=logger)
        elif patient['task'] == 2:
            input = utils.read_image(os.path.join(patient['output_dir'],'cbct_s1.nii.gz'),log=logger)
        else:
            logger.error('Task not valid')
            sys.exit(1)
        ct = utils.read_image(os.path.join(patient['output_dir'],'ct_s1.nii.gz'),log=logger)
        fov_s1 = utils.read_image(os.path.join(patient['output_dir'],'fov_s1.nii.gz'),log=logger)
        if patient['defacing_correction'] == True:
            face = utils.read_image(os.path.join(patient['output_dir'],'defacing_mask.nii.gz'),log=logger)

        # Clip CT to valid HU range
        ct = utils.clip_image(ct,-1024,3071)

        # Perform cone correction for fov mask if task2
        if patient['task'] == 2:
            if patient['cone_correction']:
                fov_s1 = utils.cone_correction(fov_s1,log=logger)

        #Generate patient outline and postprocess it
        mask = utils.segment_outline(input,patient['mask_thresh'],log=logger)
        if patient['defacing_correction']:
            defacing_correction = os.path.join(patient['output_dir'],'defacing_mask.nii.gz')
            defacing_correction = utils.read_image(defacing_correction,log=logger)
        else:
            defacing_correction = None
        if patient['IS_correction']:
            IS_correction = 10
        else:
            IS_correction = None
        mask = utils.postprocess_outline(mask,
                                         fov_s1,
                                         defacing_correction=defacing_correction,
                                         IS_correction=IS_correction,
                                         log=logger)

        #Crop images using mask generated above
        input = utils.crop_image(input,fov_s1)
        ct = utils.crop_image(ct,fov_s1)
        mask = utils.crop_image(mask,fov_s1)
        fov = utils.crop_image(fov_s1,fov_s1)

        #deform CT to match input
        ct_deformed, transform = utils.deformable_registration(input,ct,patient['parameter_def'],mask=mask,log=logger)
        sitk.WriteParameterFile(transform, os.path.join(patient['output_dir'],'transform_def.txt'))

        # #deform defacing mask if necessary and apply to fov
        # if patient['defacing_correction']:
        #     face_deformed = utils.warp_structure(face,transform)
        #     fov[face_deformed == 1] = 0

        #apply fov mask to all images
        if patient['task'] == 1:
            mask_value = 0
        if patient['task'] == 2:
            mask_value = -1000
        ct_deformed = utils.mask_image(ct_deformed,fov,-1000)
        ct = utils.mask_image(ct,fov,-1000)
        input = utils.mask_image(input,fov,mask_value)
        mask = utils.mask_image(mask,fov,0)

        #preprocess structures
        logger.info('Preprocessing and warping structures...')
        rigid_reg = sitk.ReadTransform(os.path.join(patient['output_dir'],'transform.tfm'))
        ct_s1 = utils.read_image(os.path.join(patient['output_dir'],'ct_s1.nii.gz'),log=logger)
        #utils.preprocess_structures(patient,input,ct_s1,fov_s1,fov,rigid_reg,transform,mask,log=logger)

        # Stitch CT_def to CT_s1 for planning (structures are stitched above)
        ct_deformed_stitched = utils.stitch_image(ct_deformed, ct_s1, mask)

        #Save cropped images and transform
        if patient['task'] == 1:
            utils.save_image(input,os.path.join(patient['output_dir'],'mr_s2.nii.gz'),dtype='int16')
        if patient['task'] == 2:
            utils.save_image(input,os.path.join(patient['output_dir'],'cbct_s2.nii.gz'),dtype='int16')
        utils.save_image(ct,os.path.join(patient['output_dir'],'ct_s2.nii.gz'),dtype='int16')
        utils.save_image(mask,os.path.join(patient['output_dir'],'mask_s2.nii.gz'))
        utils.save_image(fov,os.path.join(patient['output_dir'],'fov_s2.nii.gz'))
        utils.save_image(ct_deformed,os.path.join(patient['output_dir'],'ct_s2_def.nii.gz'),dtype='int16')
        utils.save_image(ct_deformed_stitched, os.path.join(patient['output_dir'],'ct_s2_def_stitched.nii.gz'),dtype='int16')

        #Generate png overviews
        utils.generate_overview_png(ct,input,mask,patient)
        utils.generate_overview_planning(ct,input,ct_deformed,mask,patient)