Python解析nii文件为png图像

最近在做一个医学影像的项目,Java在解析nii文件比较繁琐,最终采用Python来解析nii文件.然后由Java调用Python来实现这一功能. 闲话没有上代码:

import os
import nibabel
import imageio
import numpy
import json
import sys
from PIL import Image

'''
脚本说明: 
    解析单个nii文件为png图片格式(支持解析标注)
    nii_file: nii文件地址
    output_dir: png输出地址 
    prefix_png_name: 输出图片名前缀
    isTumor: 是否为标注文件
Example: 
    nii_to_png("G:/projects/testpy/20180803001499/T1SC_tumor.nii", "G:/projects/testpy/images/T1SC_tumor", "T1SC_tumor", "true")
'''
def nii_to_png(nii_file, output_dir, prefix_png_name, isTumor = "false"):
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 加载nii文件
    img = nibabel.load(nii_file)
    data = img.get_fdata()

    # 获取图像的维度
    dims = data.shape

    # 所有解析完成png_name的合集
    file_names = []

    # 遍历每个切片并将其保存为png图像
    for i in range(dims[2]):
        # 切片数据
        slice_data = data[:, :, i]

        # 切片对用png文件名
        file_name = f"{prefix_png_name}_{i}"

        # png输出地址
        output_file = os.path.join(output_dir, f"{file_name}.png")

        min_value = numpy.min(slice_data)
        max_value = numpy.max(slice_data)
        
        # 将图像转换为灰度格式
        # gray_slice_data = (slice_data - numpy.min(slice_data)) / (numpy.max(slice_data) - numpy.min(slice_data)) * 255
        if max_value != min_value:
            gray_slice_data = (slice_data - min_value) / (max_value - min_value) * 255
        else:
            gray_slice_data = slice_data * 0 + 255  # 如果最大值等于最小值，将灰度值设置为255（黑色）

        # 字符集
        gray_slice_data = gray_slice_data.astype(numpy.uint8)

        # 翻转90度
        flipped_slice_data = numpy.rot90(gray_slice_data, k=1)

        # 保存为PNG格式
        imageio.imwrite(output_file, flipped_slice_data)

        # 如果是标注文件对png背景进行处理
        if bool(isTumor):
            # 处理背景
            make_transparent(output_file)

        # 把png文件名放到集合中
        file_names.append(file_name)

    # 将文件名数组转换为JSON字符串并输出
    json_file_names = json.dumps(file_names)

    # 输出文件名
    print(json_file_names)

# 把png图片的黑色背景改为透明背景
def make_transparent(image_path):
    # 打开图片并转换为RGBA模式
    img = Image.open(image_path).convert("RGBA")

    # 获取图片的宽度和高度
    width, height = img.size

    # 遍历图片的每个像素
    for x in range(width):
        for y in range(height):
            # 获取当前像素的RGB值
            r, g, b, a = img.getpixel((x, y))

            # 如果像素是黑色（R=0, G=0, B=0），则将其设置为透明（A=0）
            if r == 0 and g == 0 and b == 0:
                img.putpixel((x, y), (0, 0, 0, 0))

    # 保存修改后的图片
    img.save(image_path, "PNG")

if __name__ == "__main__":
    nii_file_path = sys.argv[1]
    png_output_dir = sys.argv[2]
    png_prefix_name = sys.argv[3]
    is_tumor = sys.argv[4]
    nii_to_png(nii_file_path, png_output_dir, png_prefix_name, is_tumor)