以下函数主要用于记录和保存,方便自己查阅。
---------持续更新
1. 3D 图像处理
def numpy2sitk(arr, sitk_ori_img):
# numpy转换为sitk
sitk_img = sitk.GetImageFromArray(arr)
sitk_img.SetOrigin(sitk_ori_img.GetOrigin())
sitk_img.SetSpacing(sitk_ori_img.GetSpacing())
sitk_img.SetDirection(sitk_ori_img.GetDirection())
return sitk_img
import dicom2nifti
import os
import pydicom as pdic
# dicom convert to nii file
# method 1
def dicom2nii(dicom_path, save_root):
for patient in os.listdir(dicom_path):
new_path = dicom_path + patient +'/'
os.makedirs(save_root + patient +'/',exist_ok=True)
dicom2nifti.convert_directory(new_path,save_root + patient +'/')
create_name = os.listdir(save_root + patient)[0]
print(create_name)
os.rename(save_root + patient +f'/{create_name}', save_root + patient + f'/{patient}.nii.gz')
# method 2
def dcm2nii(dcm_path, nii_path):
print(dcm_path)
dcm = pdic.read_file(dcm_path)
if 'ImagerPixelSpacing' in dcm:
sp_2d = dcm.ImagerPixelSpacing
else:
sp_2d = (0.278875, 0.278875)
arr = dcm.pixel_array
img = sitk.GetImageFromArray(arr)
img.SetSpacing((sp_2d[0], sp_2d[1], 1))
sitk.WriteImage(img, nii_path)
def resample_img(itk_image, out_spacing=[2.0, 2.0, 2.0], is_label=False):
# resample images to 2mm spacing with simple itk
original_spacing = itk_image.GetSpacing()
original_size = itk_image.GetSize()
out_size = [
int(np.round(original_size[0] * (original_spacing[0] / out_spacing[0]))),
int(np.round(original_size[1] * (original_spacing[1] / out_spacing[1]))),
int(np.round(original_size[2] * (original_spacing[2] / out_spacing[2])))]
resample = sitk.ResampleImageFilter()
resample.SetOutputSpacing(out_spacing)
resample.SetSize(out_size)
resample.SetOutputDirection(itk_image.GetDirection())
resample.SetOutputOrigin(itk_image.GetOrigin())
resample.SetTransform(sitk.Transform())
resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())
if is_label:
resample.SetInterpolator(sitk.sitkNearestNeighbor)
else:
resample.SetInterpolator(sitk.sitkBSpline)
return resample.Execute(itk_image)
def roi_extract(img_path, lab_path, patch_size):
"""从目标区域中心处扩展成patch_size大小的区域
Args:
img_path (str): 图像路径
lab_path (str): 标签路径
patch_size (int): 扩展的目标尺寸
Returns:
sitk image: 输出sitk数据
"""
img = sitk.ReadImage(img_path)
lab = sitk.ReadImage(lab_path)
img_arr = sitk.GetArrayFromImage(img)
lab_arr = sitk.GetArrayFromImage(lab)
expand_slice = 0
nonzero_list = []
for ind in [(1,2),(0,2),(0,1)]:
nonzero_ind = np.any(lab_arr, axis=ind)
start_slice, end_slice = np.where(nonzero_ind)[0][[0, -1]] # [ 0, 605]
nonzero_list.append([start_slice, end_slice])
for ind, slice in enumerate(nonzero_list):
slice_diff = slice[1]-slice[0]
expand_slice = (patch_size-slice_diff)//2
if slice[0] < expand_slice:
slice[0] = 0
slice[1] += 2*expand_slice
elif slice[1] + expand_slice >= lab_arr.shape[ind]:
slice[1] = lab_arr.shape[ind] - 1
slice[0] -= 2*expand_slice
else:
slice[0] -= expand_slice
slice[1] += expand_slice
# 补全
if slice[1]-slice[0] > patch_size:
slice[0] += 1
if slice[1]-slice[0] < patch_size:
slice[1] += 1
new_img_arr = img_arr[nonzero_list[0][0]:nonzero_list[0][1], nonzero_list[1][0]:nonzero_list[1][1], nonzero_list[2][0]:nonzero_list[2][1]]
new_lab_arr = lab_arr[nonzero_list[0][0]:nonzero_list[0][1], nonzero_list[1][0]:nonzero_list[1][1], nonzero_list[2][0]:nonzero_list[2][1]]
sitk_img = numpy2sitk(new_img_arr, img)
sitk_lab = numpy2sitk(new_lab_arr, lab)
return sitk_img, sitk_lab
def slice_slide_crop(img_path, lab_path, patch_num, patch_size):
"""朝3D数据的切片方向切patch_num+1个patch_size尺寸的Patch
Args:
img_path (str): 图像路径
lab_path (str): 标签路径
patch_num (int): patch数量
patch_size (int): patch尺寸
Returns:
sitk: 输出sitk数据并保存
"""
img = sitk.ReadImage(img_path)
lab = sitk.ReadImage(lab_path)
if patch_num*patch_size > img.GetSize()[2]:
stride = (img.GetSize()[2]-patch_size)//patch_num
start_slice, end_slice = 0, patch_size
for i in range(patch_num+1):
crop_slice_ind = [start_slice + i*stride, end_slice + i*stride]
new_img = img[:,:,crop_slice_ind[0]:crop_slice_ind[1]]
new_lab = lab[:,:,crop_slice_ind[0]:crop_slice_ind[1]]
sitk.WriteImage(new_img,f'0001_0001/crop_img/0001_0001_s{i}.nii.gz')
sitk.WriteImage(new_lab,f'0001_0001/crop_lab/0001_0001_s{i}.nii.gz')
return True
else:
print('patch_size/patch_num is too small')
return False
2. Pytorch相关处理
# 模型加载
from collections import OrderedDict
def load_distributed_model(model_path):
# 加载并行训练的模型参数
device = torch.device("cuda")
model = DGCNN().to(device) #自己的模型
state_dict = torch.load(model_path) #存放模型的位置
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
return model
# 多类分割dice损失
def generalized_dice_loss(pred, target):
"""compute the weighted dice_loss
Args:
pred (tensor): prediction after softmax, shape(bath_size, channels, height, width)
target (tensor): gt, shape(bath_size, channels, height, width)
Returns:
gldice_loss: loss value
"""
wei = torch.sum(target, axis=[0,2,3]) # (n_class,)
wei = 1/(wei**2+epsilon)
intersection = torch.sum(wei*torch.sum(pred * target, axis=[0,2,3]))
union = torch.sum(wei*torch.sum(pred + target, axis=[0,2,3]))
gldice_loss = 1 - (2. * intersection) / (union + epsilon)
return gldice_loss
3. python 相关处理
# sitk.ReadImage 报错 正交坐标系,重新写入sfrom/qfrom信息
def get_tips_bifs(skele):
"""get tips and bifs from skeleon image by the number of pixle's neighbor
Args:
skele (np.array): skeleton array of the binary image
Returns:
tips, bifs: list of tips and bifs
"""
kernel = np.ones((3, 3))
map = ndimage.convolve(skele, kernel)
map = map * skele
tips_idx = np.where(map == 2) # including the center of the pixel
tips = np.stack(tips_idx, axis=1) # !!! stack the coordinate of 'where' function
bifs = get_bifs_with_single_pixel(map)
return tips, bifs
def get_bifs_with_single_pixel(map):
"""one bif point may have more than one pixel, so choose the representative one as the bif pixel
Args:
map (np.array): pixle's neighbor
Returns:
bifs: get one pixel of the bif point
"""
map[map < 4] = 0 # >=4 is a bif point
map[map > 0] = 1
lab = label(map)
props = regionprops(lab)
bifs = []
for n in range(len(props)):
ipdb.set_trace()
temp = props[n].coords # get the coordinates of the region
x = np.argmax(np.bincount(temp[:, 0])) # !!! get the common coordinates(众数)
y = np.argmax(np.bincount(temp[:, 1]))
bifs.append([x, y])
return np.array(bifs)
4. 错误情况
# sitk.ReadImage 报错 正交坐标系,重新写入sfrom/qfrom信息
for patient in os.listdir(root):
new_path = root+patient
img = nib.load(new_path + f'/{patient}.nii.gz')
qform = img.get_qform()
img.set_qform(qform)
sfrom = img.get_sform()
img.set_sform(sfrom)
nib.save(img, new_path + f'/{patient}.nii.gz')
# 修改 spacing mm单位为cm单位
def spacing_adjust(nii_data):
ori_spacing = nii_data.GetSpacing()
ori_origin = nii_data.GetOrigin()
nii_data.SetSpacing([i/10. for i in ori_spacing])
nii_data.SetOrigin([i/10. for i in ori_origin])
return nii_data