DCM文件转换为NII文件

import random
import math
import numbers
import collections
import numpy as np
import torch
import cv2
import scipy.ndimage
from PIL import Image, ImageOps
import SimpleITK as sitk
import os
import pydicom
from torchvision.transforms import Normalize, Compose
class ImageProcess(object):
    def __init__(self, save_dir=r"F:\data\nii_2"):
        self.save_dir = save_dir

    def write_imge(self, arry, path, name):
        arry_ = sitk.GetImageFromArray(arry)
        path = os.path.join(path, name)
        sitk.WriteImage(arry_, path)

    def read_nii(self, path):
        return sitk.GetArrayFromImage(sitk.ReadImage(path))

    # 形状统一
    def transform_shape(self, file_, save_test, save_name, target_shape=(455, 512, 512), path="null"):
        # data = read_nii(file_)
        data = file_
        current_shape = data.shape
        if current_shape != target_shape:
            stat_ = current_shape[1:]
            if stat_ == tuple((512, 512)):
                padded_data = np.zeros(target_shape, dtype=data.dtype)
                padded_data[:current_shape[0], :current_shape[1], :current_shape[2]] = data
                data = padded_data
                data = np.expand_dims(data, axis=-1)
                data = np.repeat(data, 3, axis=-1)
                self.write_imge(data, save_test, save_name)
            else:
                print(f"分辨率不是512*512.({path})")
        return data

    # 保存杂项
    def save_jpg(self, number_dir2, out_file):
        for file in os.listdir(number_dir2):
            if file.endswith('.dcm') or file.endswith('.DCM'):
                ds = pydicom.dcmread(os.path.join(number_dir2, file))
                # 可以根据窗宽窗位信息对图像进行处理
                image = ds.pixel_array
                # 将图像数据转换为PIL Image对象
                try:
                    image = Image.fromarray(image)
                except TypeError:
                    pass
                else:
                    # 设置输出文件路径和文件名
                    output_file = os.path.join(out_file, os.path.splitext(file)[0] + '.jpeg')
                    # 保存图像为JPEG文件
                    image = image.convert('RGB')
                    if not os.path.exists(output_file):
                        image.save(output_file, 'JPEG')

    # img静态化
    def dcm2nii(self, dicom_dir_, save_dir_, name, coverage=True):
        save_dir_ = os.path.join(save_dir_, name)
        if not os.path.exists(save_dir_):
            os.mkdir(save_dir_)
        series_ids = sitk.ImageSeriesReader.GetGDCMSeriesIDs(dicom_dir_)
        print(len(series_ids))
        filename = f'{name}.nii'
        save_path = os.path.join(save_dir_, filename)
        if coverage:
            for idx_series_ids in range(len(series_ids)):
                series_file_names = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(dicom_dir_,
                                                                                  series_ids[idx_series_ids])
                try:
                    series_reader = sitk.ImageSeriesReader()
                    series_reader.SetFileNames(series_file_names)
                    image = series_reader.Execute()
                except RuntimeError as e:
                    print("图像为非医学图像".center(60, "-"))
                    if not os.path.exists(save_dir_):
                        os.mkdir(save_dir_)
                    self.save_jpg(number_dir2=dicom_dir_, out_file=save_dir_)
                else:
                    print(f'{idx_series_ids} spacing: {image.GetSpacing()}')
                    data = sitk.GetArrayFromImage(image)
                    data = self.transform_shape(data, save_dir_, filename)
                    # sitk.WriteImage(image, save_path)

            print('转换结束'.center(60, '='))
        else:
            if not os.path.exists(save_path):
                for idx_series_ids in range(len(series_ids)):
                    series_file_names = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(dicom_dir_,
                                                                                      series_ids[idx_series_ids])
                    try:
                        series_reader = sitk.ImageSeriesReader()
                        series_reader.SetFileNames(series_file_names)
                        image = series_reader.Execute()
                    except RuntimeError as e:
                        print(e)
                        if not os.path.exists(save_dir_):
                            os.mkdir(save_dir_)
                        self.save_jpg(number_dir2=dicom_dir_, out_file=save_dir_)
                    else:
                        print(f'{idx_series_ids} spacing: {image.GetSpacing()}')
                        data = sitk.GetArrayFromImage(image)
                        data = self.transform_shape(data, save_dir_, filename)
                        # sitk.WriteImage(image, save_path)

                print('转换结束'.center(60, '='))
            else:
                print(f"{save_path}文件已存在，未覆盖")

    def __call__(self, dicom_dir=r"F:\data\test"):
        #dcm转nii文件并实现数据形状统一
        for num_, dir_name in enumerate(os.listdir(dicom_dir)):
            path = os.path.join(dicom_dir, dir_name)
            self.dcm2nii(dicom_dir_=path, save_dir_=self.save_dir, name=dir_name, coverage=True)
        return self.save_dir
    
DataTransform=Compose([ImageProcess(save_dir="save_dir")])
save_dir=DataTransform("dicom_dir")