python获取nii中不同标签分别输出为对应nii文件,并统计对应标签数量

import nibabel as nib
import numpy as np
import time

# 读取nifti文件
starttime = time.time()
nii_path = 'd:/test.nii'
nii_img = nib.load(nii_path)

# 获取原始数据和头部信息
data = nii_img.get_fdata()
header = nii_img.header

# 输出标签为1的文件
label_1_data = np.zeros_like(data)
label_1_data[data == 1] = 1
label_1_nii_img = nib.Nifti1Image(label_1_data, nii_img.affine, header)
output_path_label_1 = 'd:/test1.nii'
nib.save(label_1_nii_img, output_path_label_1)

# # 输出标签为2的文件
# label_2_data = np.zeros_like(data)
# label_2_data[data == 2] = 2
# label_2_nii_img = nib.Nifti1Image(label_2_data, nii_img.affine, header)
# output_path_label_2 = 'd:/test2.nii.gz'
# nib.save(label_2_nii_img, output_path_label_2)

# 输出标签为3的文件
label_3_data = np.zeros_like(data)
label_3_data[data == 3] = 3
label_3_nii_img = nib.Nifti1Image(label_3_data, nii_img.affine, header)
output_path_label_3 = 'd:/test3.nii'
nib.save(label_3_nii_img, output_path_label_3)

# 输出标签为4的文件
label_4_data = np.zeros_like(data)
label_4_data[data == 4] = 4
label_4_nii_img = nib.Nifti1Image(label_4_data, nii_img.affine, header)
output_path_label_4 = 'd:/test4.nii'
nib.save(label_4_nii_img, output_path_label_4)

endtime = time.time()
print(endtime - starttime)

# 统计每个标签的像素数量
label_counts = {
    'Label 1': np.sum(label_1_data == 1),
    'Label 2': np.sum(label_2_data == 2),
    'Label 3': np.sum(label_3_data == 3),
    'Label 4': np.sum(label_4_data == 4)
}

print("Label Counts:")
for label, count in label_counts.items():
    print(label, ": ", count)