preprocess.py 文件源码

def preprocess(inputfile, outputfile, order=0, df=None, input_key=None, output_key=None):
    img = nib.load(inputfile)
    data = img.get_data()
    affine = img.affine
    zoom = img.header.get_zooms()[:3]
    data, affine = reslice(data, affine, zoom, (1., 1., 1.), order)
    data = np.squeeze(data)
    data = np.pad(data, [(0, 256 - len_) for len_ in data.shape], "constant")
    if order == 0:
        if df is not None:
            tmp = np.zeros_like(data)
            for target, source in zip(df[output_key], df[input_key]):
                tmp[np.where(data == source)] = target
            data = tmp
        data = np.int32(data)
        assert data.ndim == 3, data.ndim
    else:
        data_sub = data - gaussian_filter(data, sigma=1)
        img = sitk.GetImageFromArray(np.copy(data_sub))
        img = sitk.AdaptiveHistogramEqualization(img)
        data_clahe = sitk.GetArrayFromImage(img)[:, :, :, None]
        data = np.concatenate((data_clahe, data[:, :, :, None]), 3)
        data = (data - np.mean(data, (0, 1, 2))) / np.std(data, (0, 1, 2))
        assert data.ndim == 4, data.ndim
        assert np.allclose(np.mean(data, (0, 1, 2)), 0.), np.mean(data, (0, 1, 2))
        assert np.allclose(np.std(data, (0, 1, 2)), 1.), np.std(data, (0, 1, 2))
        data = np.float32(data)
    img = nib.Nifti1Image(data, affine)
    nib.save(img, outputfile)