Utils.py 文件源码

def export_data(prediction_dir, nii_image_dir, tfrecords_dir, export_dir, transformation=None):

    for file_path in os.listdir(prediction_dir):
        name, prediction, probability = read_prediction_file(os.path.join(prediction_dir, file_path))

        if transformation:
            image, ground_truth = get_original_image(os.path.join(tfrecords_dir, name + '.tfrecord'), False)
            prediction = transformation.transform_image(prediction, probability, image)

        # build cv_predictions .nii image
        img = nib.Nifti1Image(prediction, np.eye(4))
        img.set_data_dtype(dtype=np.uint8)

        path = os.path.join(nii_image_dir, name)

        adc_name = next(l for l in os.listdir(path) if 'MR_ADC' in l)
        export_image = nib.load(os.path.join(nii_image_dir, name, adc_name, adc_name + '.nii'))

        i = export_image.get_data()
        i[:] = img.get_data()

        # set name to specification and export
        _id = next(l for l in os.listdir(path) if 'MR_MTT' in l).split('.')[-1]
        export_path = os.path.join(export_dir, 'SMIR.' + name + '.' + _id + '.nii')
        nib.save(export_image, os.path.join(export_path))