230922-金属伪影MAR方法中LIMAR和NMAR

https://github.com/hongwang01/InDuDoNet/blob/main/CLINIC_metal/preprocess_clinic/preprocessing_clinic.py
提供了LI的代码。
NMAR代码：
https://github.com/yanbozhang007/CNN-MAR/blob/ba525c75efdcf9d35889c749e66f49e4445b187b/evaluation/nmar.m
借助ChatGPT改写为python代码

# Given clinical Xma, generate data,including: XLI, M, Sma, SLI, Tr for infering InDuDoNet
import nibabel
import numpy as np
import os
from scipy.ndimage import gaussian_filter
from sklearn.cluster import KMeans
from scipy.interpolate import interp1d
from .utils import get_config
from .build_gemotry import initialization, imaging_geo
import PIL
from PIL import Image
config = get_config('CLINIC_metal/preprocess_clinic/dataset_py_640geo.yaml')
CTpara = config['CTpara']  # CT imaging parameters
mask_thre = 2500 / 1000 * 0.192 + 0.192  # taking 2500HU as a thresholding to segment the metal region
param = initialization()
ray_trafo, FBPOper = imaging_geo(param)  # CT imaging geometry, ray_trafo is fp, FBPoper is fbp
allXma = []
allXLI = []
allM = []
allSma = []
allSLI = []
allTr = []
allaffine = []
allfilename=[]
miuWater = 0.192
miuAir = 0
starpoint = np.array([miuAir, miuWater, 2 * miuWater])
# process and save all the to-the-tested volumes
def clinic_input_data(test_path):
    for file_name in os.listdir(test_path):
        file_path = test_path+'/'+file_name
        img = nibabel.load(file_path)
        imag = img.get_fdata()  # imag with pixel as HU unit
        affine = img.affine
        allaffine.append(affine)
        num_s = imag.shape[2]
        M  = np.zeros((CTpara['imPixNum'], CTpara['imPixNum'], num_s), dtype='float32')
        Xma = np.zeros_like(M)
        XLI = np.zeros_like(M)
        Tr = np.zeros((CTpara['sinogram_size_x'], CTpara['sinogram_size_y'], num_s), dtype='float32')
        Sma = np.zeros_like(Tr)
        SLI =  np.zeros_like(Tr)
        for i in range(70, num_s):
            image = np.array(Image.fromarray(imag[:,:,i]).resize((CTpara['imPixNum'], CTpara['imPixNum']), PIL.Image.BILINEAR))
            image[image < -1000] = -1000
            image = image / 1000 * 0.192 + 0.192
            Xma[...,i] = image
            [rowindex, colindex] = np.where(image > mask_thre)
            M[rowindex, colindex, i] = 1
            Pmetal_kev = np.asarray(ray_trafo(M[:,:,i]))
            Tr[...,i] = Pmetal_kev > 0
            Sma[...,i] = np.asarray(ray_trafo(image))
            SLI[...,i] = interpolate_projection(Sma[...,i], Tr[...,i])
            try:
                XLI[..., i] = np.asarray(FBPOper(SLI[..., i]))
            except:
                pass
            XLI[...,i] = np.asarray(FBPOper(SLI[...,i]))

            # smFilter = gaussian_filter(size=(5, 5), sigma=1)
            # NMAR method 1: Raw
            imRaw = Xma[...,i]
            metalBW = np.array(M[:,:,i],dtype=bool)
            metalTrace = Tr[...,i]
            proj = Sma[...,i]
            imRaw[metalBW] = miuWater
            imPriorNMAR1 = nmar_prior(imRaw, metalBW)
            projPrior1 = np.asarray(ray_trafo(imPriorNMAR1))
            PNMAR1 = nmar_proj(proj, projPrior1, metalTrace)
            imNMAR1 = np.asarray(FBPOper(PNMAR1))

            # NMAR method 2: LI
            imLI = XLI[..., i]
            imLI[metalBW] = miuWater
            imPriorNMAR2 = nmar_prior(imLI, metalBW)
            projPrior2 = np.asarray(ray_trafo(imPriorNMAR2))
            PNMAR2 = nmar_proj(proj, projPrior2, metalTrace)
            imNMAR2 = np.asarray(FBPOper(PNMAR2))

        allXma.append(Xma)
        allXLI.append(XLI)
        allM.append(M)
        allSma.append(Sma)
        allSLI.append(SLI)
        allTr.append(Tr)
        allfilename.append(file_name)
    return allXma, allXLI, allM, allSma, allSLI, allTr, allaffine, allfilename


def interpolate_projection(proj, metalTrace):
    # projection linear interpolation
    # Input:
    # proj:         uncorrected projection
    # metalTrace:   metal trace in projection domain (binary image)
    # Output:
    # Pinterp:      linear interpolation corrected projection
    Pinterp = proj.copy()
    for i in range(Pinterp.shape[0]):
        mslice = metalTrace[i]
        pslice = Pinterp[i]

        metalpos = np.nonzero(mslice==1)[0]
        nonmetalpos = np.nonzero(mslice==0)[0]
        pnonmetal = pslice[nonmetalpos]
        pslice[metalpos] = interp1d(nonmetalpos,pnonmetal)(metalpos)
        Pinterp[i] = pslice

    return Pinterp


def nmar_prior(XLI, M):
    XLI[M == 1] = miuWater
    h, w = XLI.shape
    im1d = XLI.reshape(h * w)
    kmeans = KMeans(n_clusters=3, init=starpoint.reshape(-1, 1), max_iter=300)
    labels = kmeans.fit_predict(im1d.reshape(-1, 1))
    threshBone2 = np.min(im1d[labels == 2])
    threshBone2 = np.max([threshBone2, 1.2 * miuWater])
    threshWater2 = np.min(im1d[labels == 1])
    # imPriorNMAR = nmarprior(XLI, threshWater2, threshBone2, miuAir, miuWater, smFilter)
    imPriorNMAR = nmarprior(XLI, threshWater2, threshBone2, miuAir, miuWater)
    return imPriorNMAR

def nmarprior(im, threshWater, threshBone, miuAir, miuWater):
    imSm = gaussian_filter(im, sigma=1)
    # smFilter = sio.loadmat('deeplesion/gaussianfilter.mat')['smFilter']
    # imSm = scipy.ndimage.filters.convolve(im, smFilter, mode='nearest')
    priorimgHU = imSm.copy()
    priorimgHU[imSm <= threshWater] = miuAir
    priorimgHU[(imSm > threshWater) & (imSm < threshBone)] = miuWater
    return priorimgHU

def nmar_proj(proj, Pprior, metalTrace):
    Pprior[Pprior < 0] = 0
    eps = 1e-6
    Pprior += eps
    Pnorm = proj / Pprior
    Pnorminterp = interpolate_projection(Pnorm, metalTrace)
    Pnmar = Pnorminterp * Pprior
    nonmetalpos = (metalTrace == 0)
    Pnmar[nonmetalpos] = proj[nonmetalpos]
    return Pnmar