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)
maxs = np.max(dots,axis=0)
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:
""" 归一化 """
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]
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)
imageArray = normalize_img(imageArray)
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])