医学图像裁剪为小块，保存为nii文件进行训练-Toy模板网

这篇具有很好参考价值的文章主要介绍了医学图像裁剪为小块，保存为nii文件进行训练。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

import numpy as np
import SimpleITK as sitk
import os
from tqdm import tqdm


def cal_board(label_array,index):
    """calculate borader
    0: depth
    1: hight
    2: width
    """
    D,H,W = label_array.shape    
    dots = np.argwhere(label_array==1)
    mins = np.min(dots,axis=0) # z,y,x // 这个一定是0
    maxs = np.max(dots,axis=0) # z,y,x // 这个一定是0
    # print(mins)
    # print(maxs)
    low_boarder, high_boarder = mins[index],maxs[index] 

    return low_boarder, high_boarder


def normalize_img(img:np.ndarray,min_val=-1000,max_val=600)->np.ndarray:
    """ 归一化 """
    # max_val = np.max(img)
    # min_val = np.min(img)
    value_range = max_val - min_val
    norm_0_1 = (img-min_val)/value_range
    img = np.clip(2*norm_0_1-1,-1,1)        
    return img

def crop_img(pure_name,
            image_array, label_array, boarder_depths, boarder_hight, boarder_width,
             target_size=[32,256,256],
             ):

    depth_low, depth_high = boarder_depths[0]-1, boarder_depths[1]+1
    hight_low, hight_high = boarder_hight[0]-1, boarder_hight[1]+1
    width_low, width_high = boarder_width[0]-1, boarder_width[1]+1

    d,h,w = target_size[0],target_size[1],target_size[2]

    cnt = 0

    for k in range(depth_low, depth_high-d//2+1, d//2):
        for j in range(hight_low, hight_high - h//2+1, h//2):
            for i in range(width_low, width_high-w//2 + 1, w//2):
                if(k+d >= depth_high):
                    k = depth_high - d
                if(j+h >= hight_high):
                    j = hight_high - h
                if(i+w >= width_high):
                    i = width_high - w

                crop_img = image_array[k:k+d,j:j+h,i:i+w]
                crop_lab = label_array[k:k+d,j:j+h,i:i+w]

                if np.sum(crop_lab) < 500:
                    continue

                print(f'd:[{k}-{k+d}), h:[{j}-{j+h}), w:[{i}-{i+w})')
                print(f'crop_img.shape={crop_img.shape}, crop_lab.shape={crop_lab.shape}')
                assert crop_img.shape==crop_lab.shape,f'crop_img.shape={crop_img.shape}, crop_lab.shape={crop_lab.shape}, they are not equal!'

                cropImage = sitk.GetImageFromArray(crop_img)
                cropLabel = sitk.GetImageFromArray(crop_lab)

                
                cropImageSaveDir = os.path.join(saveDir, "images")
                cropLabelSaveDir = os.path.join(saveDir, "labels")
                
                cropSaveName = f'{pure_name}_{cnt:04}.nii.gz'
                
                cropImageSavePath = os.path.join(cropImageSaveDir, cropSaveName)
                cropLabelSavePath = os.path.join(cropLabelSaveDir, cropSaveName)

                sitk.WriteImage(cropImage, cropImageSavePath)
                sitk.WriteImage(cropLabel, cropLabelSavePath)
                
                cnt +=1


if __name__ == '__main__':
    saveDir = "./head_dataset"


    image_path = "./test_data/86_image.nii.gz"
    label_path = "./test_data/86_label.nii.gz"

    file_name = os.path.basename(image_path)
    pureName = file_name.split('.')[0]
    # input(pureName)

    image = sitk.ReadImage(image_path)
    label = sitk.ReadImage(label_path)

    imageArray = sitk.GetArrayFromImage(image)
    labelArray = sitk.GetArrayFromImage(label)

    print(f'imageArray.shape={imageArray.shape}, labelArray.shape={labelArray.shape}')

    board_depth_1, board_depth_2 = cal_board(labelArray, 0)
    board_hight_1, board_hight_2 = cal_board(labelArray, 1)
    board_width_1, board_width_2 = cal_board(labelArray, 2)

    print(board_depth_1, board_depth_2)
    print(board_hight_1, board_hight_2)
    print(board_width_1, board_width_2)

    # normalize image
    imageArray = normalize_img(imageArray)

    # [32,256,256]
    crop_img(pureName,
        image_array=imageArray, label_array=labelArray,
              boarder_depths=[board_depth_1,board_depth_2],
              boarder_hight=[board_hight_1, board_hight_2],
              boarder_width=[board_width_1, board_width_2])