pix2pix处理dicom数据集

Doer491

已于 2022-09-02 19:29:16 修改

阅读量425

点赞数

文章标签： python 人工智能机器学习

于 2022-08-12 14:51:24 首次发布

本文链接：https://blog.csdn.net/qq_60742456/article/details/126304120

版权

import os
import numpy as np
import cv2
import argparse
import SimpleITK as sitk
from multiprocessing import Pool


def image_write(path_A, path_B, path_AB):
    im_A = sitk.ReadImage(path_A)
    im_B = sitk.ReadImage(path_B)
    im_A_npy = sitk.GetArrayFromImage(im_A)
    im_B_npy = sitk.GetArrayFromImage(im_B)
    im_AB_npy = np.concatenate((im_A_npy, im_B_npy), axis=2)
    im_AB = sitk.GetImageFromArray(im_AB_npy)
    sitk.WriteImage(im_AB, path_AB + ".mha")


if __name__ == '__main__':

    def image_write(path_A, path_B, path_AB):
        im_A = sitk.ReadImage(path_A)
        im_B = sitk.ReadImage(path_B)
        im_A_npy = sitk.GetArrayFromImage(im_A)
        im_B_npy = sitk.GetArrayFromImage(im_B)
        im_AB_npy = np.concatenate((im_A_npy, im_B_npy), axis=2)
        im_AB = sitk.GetImageFromArray(im_AB_npy)
        sitk.WriteImage(im_AB, path_AB + ".mha")


    parser = argparse.ArgumentParser('create image pairs')
    parser.add_argument('--fold_A', dest='fold_A', help='input directory for image A', type=str, default='./datasets/A')
    parser.add_argument('--fold_B', dest='fold_B', help='input directory for image B', type=str, default='./datasets/B')
    parser.add_argument('--fold_AB', dest='fold_AB', help='output directory', type=str, default='./datasets/AB')
    parser.add_argument('--num_imgs', dest='num_imgs', help='number of images', type=int, default=1000000)
    parser.add_argument('--use_AB', dest='use_AB', help='if true: (0001_A, 0001_B) to (0001_AB)', action='store_true')
    parser.add_argument('--no_multiprocessing', dest='no_multiprocessing',
                        help='If used, chooses single CPU execution instead of parallel execution', action='store_true',
                        default=False)
    args = parser.parse_args()

    for arg in vars(args):
        print('[%s] = ' % arg, getattr(args, arg))

    splits = os.listdir(args.fold_A)

    if not args.no_multiprocessing:
        pool = Pool()

    for sp in splits:
        img_fold_A = os.path.join(args.fold_A, sp)
        img_fold_B = os.path.join(args.fold_B, sp)
        img_list = os.listdir(img_fold_A)
        if args.use_AB:
            img_list = [img_path for img_path in img_list if '_A.' in img_path]

        num_imgs = min(args.num_imgs, len(img_list))
        print('split = %s, use %d/%d images' % (sp, num_imgs, len(img_list)))
        img_fold_AB = os.path.join(args.fold_AB, sp)
        if not os.path.isdir(img_fold_AB):
            os.makedirs(img_fold_AB)
        print('split = %s, number of images = %d' % (sp, num_imgs))
        for n in range(num_imgs):
            name_A = img_list[n]
            path_A = os.path.join(img_fold_A, name_A)
            if args.use_AB:
                name_B = name_A.replace('_A.', '_B.')
            else:
                name_B = name_A
            path_B = os.path.join(img_fold_B, name_B)
            if os.path.isfile(path_A) and os.path.isfile(path_B):
                name_AB = name_A
                if args.use_AB:
                    name_AB = name_AB.replace('_A.', '.')  # remove _A
                path_AB = os.path.join(img_fold_AB, name_AB)
                if not args.no_multiprocessing:
                    pool.apply_async(image_write, args=(path_A, path_B, path_AB))
                else:
                    im_A = sitk.GetArrayFromImage(path_A)  # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
                    im_B = sitk.GetArrayFromImage(path_B)  # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR
                    im_AB = np.concatenate([im_A, im_B], 1)
                    sitk.WriteImage(im_AB, path_AB)
    if not args.no_multiprocessing:
        pool.close()
        pool.join()

sitk读image，转np(数组)，再转image，保存image（.mha或.dcm）


import SimpleITK as sitk
import numpy as np
# 文件路径
img_path = r"E:\data\data\1\DYN_10s_DCE_701\54_1.dcm"
mask_path = r"E:\data\data\1\DYN_10s_DCE_701\54_2.dcm"
# 读文件
im_A = sitk.ReadImage(img_path)
im_B = sitk.ReadImage(mask_path)
im_A_npy = np.array(sitk.GetArrayFromImage(im_A))
im_B_npy = np.array(sitk.GetArrayFromImage(im_B))
new_im_A = sitk.GetImageFromArray(im_A_npy.astype('int16'))
new_im_B = sitk.GetImageFromArray(im_B_npy.astype('int16'))
im_AB_npy = np.concatenate((im_A_npy, im_B_npy), axis=2)
im_AB = sitk.GetImageFromArray(im_AB_npy) #.astype('int16')
# print(im_AB_npy)
# print(im_AB)
sitk.WriteImage(new_im_A, './1.dcm')
sitk.WriteImage(new_im_B, './2.dcm')
sitk.WriteImage(im_AB, './3.mha')
# if __name__ == '__main__':
# 	# sitk.WriteImage('./', im_AB)
# 	print(im_A)
# 	print(im_A_npy)

dicom2tensor

import SimpleITK as sitk
import torchvision

AB_path = "./datasets/AB/train/3900.dcm.mha"
AB = sitk.ReadImage(AB_path)

w = AB.GetWidth()
h = AB.GetHeight()

w2 = int(w / 2)

A = sitk.Crop(AB, lowerBoundaryCropSize=[0, 0, 0], upperBoundaryCropSize=[w2, 0, 0])
B = sitk.Crop(AB, lowerBoundaryCropSize=[w2, 0, 0], upperBoundaryCropSize=[0, 0, 0])
AW = A.GetWidth()
BW = B.GetWidth()
AH = A.GetHeight()
BH = B.GetHeight()
A = sitk.GetArrayFromImage(A)
B = sitk.GetArrayFromImage(B)
A = sitk.GetArrayFromImage(A)
B = sitk.GetArrayFromImage(B)
A = torch.tensor(A)
B = torch.tensor(B)

print(A.shape)

numpy_resize之pad，默认添0。

import numpy as np
import SimpleITK as sitk
img_path = r"C:\1\1\datasets\AB\train\1.dcm.mha"
im_A = sitk.ReadImage(img_path)
im_A_npy = sitk.GetArrayFromImage(im_A)
print(im_A_npy.shape)
print(im_A_npy)
ndarray = np.pad(im_A_npy, ((0, 0), (4, 4), (8, 8)))
print(ndarray.shape)
print(ndarray)

加载并显示。

import numpy as np
import SimpleITK as sitk
import matplotlib.pyplot as plt
AB_path = "./datasets/AB/train/20.dcm.mha"
# transforms.ToTensor()
A = sitk.ReadImage(AB_path)# 读取图像
# print(A.GetSize())
A = sitk.GetArrayFromImage(A)
# print(A.shape)
print(A.shape)
A = np.transpose(A, (1, 2, 0))
plt.imshow(A, cmap=plt.cm.bone)
plt.show()