通过联通成分分析，移除小区域

# 通过连通成分分析，移除小区域
import SimpleITK as sitk
import os
import argparse
from pathlib import Path
 
 
def RemoveSmallConnectedCompont(sitk_maskimg, rate=0.5):
    '''
    two steps:
        step 1: Connected Component analysis: 将输入图像分成 N 个连通域
        step 2: 假如第 N 个连通域的体素小于最大连通域 * rate，则被移除
    :param sitk_maskimg: input binary image 使用 sitk.ReadImage(path, sitk.sitkUInt8) 读取，
                        其中sitk.sitkUInt8必须注明，否则使用 sitk.ConnectedComponent 报错
    :param rate: 移除率，默认为0.5， 小于 1/2最大连通域体素的连通域被移除
    :return:  binary image， 移除了小连通域的图像
    '''
 
    # step 1 Connected Component analysis
    cc = sitk.ConnectedComponent(sitk_maskimg)
    stats = sitk.LabelIntensityStatisticsImageFilter()
    stats.Execute(cc, sitk_maskimg)
    maxlabel = 0   # 获取最大连通域的索引
    maxsize = 0    # 获取最大连通域的体素大小
 
    # 遍历每一个连通域， 获取最大连通域的体素大小和索引
    for l in stats.GetLabels():  # stats.GetLabels()  (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)
        size = stats.GetPhysicalSize(l)   # stats.GetPhysicalSize(5)=75  表示第5个连通域的体素有75个
        if maxsize < size:
            maxlabel = l
            maxsize = size
 
    # step 2 获取每个连通域的大小，保留 size >= maxsize * rate 的连通域
    not_remove = []
    for l in stats.GetLabels():
        size = stats.GetPhysicalSize(l)
        if size >= maxsize * rate:
            not_remove.append(l)
 
    labelmaskimage = sitk.GetArrayFromImage(cc)
    outmask = labelmaskimage.copy()
    outmask[labelmaskimage != maxlabel] = 0
    for i in range(len(not_remove)):
        outmask[labelmaskimage == not_remove[i]] = 1
  # 保存图像
    outmask = outmask.astype('float32')
 
    out = sitk.GetImageFromArray(outmask)
    out.SetDirection(sitk_maskimg.GetDirection())
    out.SetSpacing(sitk_maskimg.GetSpacing())
    out.SetOrigin(sitk_maskimg.GetOrigin())   # 使 out 的层厚等信息同输入一样
 
    return out  # to save image: sitk.WriteImage(out, 'largecc.nii.gz')
 
 
if __name__ == '__main__':
#     parser = argparse.ArgumentParser(description="remove small connected domains")
#     parser.add_argument('--input', type=str, default="output.nii")
#     parser.add_argument("--output", type=str, default='output.nii')
#     args = parser.parse_args()
    #args,unknow =parser.parse_know_args()
 
    # for single image
  
    sitk_maskimg = sitk.ReadImage('output_1.nii', sitk.sitkUInt8)
    out = RemoveSmallConnectedCompont(sitk_maskimg, rate=0.3)  # 可以设置不同的比率
    sitk.WriteImage(out, 'test_1.nii')