【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割10（测试推理篇）

对于直接将裁剪的patch，一个个的放到训练好的模型中进行预测，这部分代码可以直接参考前面的训练部分就行了。其实说白了，就是验证部分。不使用dataloader的方法，也只需要修改少部分代码即可。

但是，这种方法是不end to end的。我们接下来要做的，就是将一个CT数组作为输入，产生patch，最后得到预测的完整结果。这样一个初衷，就需要下面几个步骤：

1. 读取一个序列的CT数组；
2. OverLap的遍历所有位置，裁剪出一个个patch；
3. 一个个patch送进模型，进行预测；
4. 对预测结果，再一个个拼接起来，组成一个和CT数组一样大小的预测结果。

这里，就要不得不先补齐下对数组crop和merge操作的方法了，建议先去学习下本系列的文章，

链接：【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割9（patch 的 crop 和 merge 操作）。

对于本系列的其他文章，也可直达专栏：人工智能与医学影像专栏。后期可不知道什么时候就收费了，先Mark住。

那本篇主要是做什么呢？直接看导言。

一、导言

前面提到，其实推理过程，是可以参考训练过程的。那么，两者之间又存在着哪些的不同呢？

1. 训练是有标签金标准的，推理没有；
2. 训练是要梯度回归，更新模型的，推理没有；
3. 训练要保存模型，推理没有；
4. 训练要循环很多次epoch，推理 1 次就好。

除了上面的这些训练要做，而推理不需要的地方，推理最最重要的就是要把预测结果，给保存下来。可能是图像形式、类别形式、或者字典形式等等。

本文就将预测结果，保存成和输入数组一样大小的数组，便于后面查看和统计。

二、模型预测

说到将已经训练好的模型，给独立进行测试，需要经历哪些步骤呢？

1. 模型和保存参数加载；
2. 数据预处理；
3. 前向推理，进行预测；
4. 预测结果后处理；
5. 结果保存。

相比于训练过程，预测过程就简单了很多，最最关键的也就在于数据的前处理，和预测结果的后处理上面。

2.1、数据前处理

由于我们在本篇的开始，就定义了目标。就是输入的一个CT的完整数组，输出是一个和输入一样大小的预测结果。

但是呢，我们的模型输入，仅仅是一个固定大小的patch，比如48x96x96的大小。所以，这就需要将shape为320x265x252，或298x279x300大小的CT数组，裁剪成一个个小块，也就是一个个patch。

这里其实在独立的一篇文章，进行了单独详细的介绍。对于本文调用的函数，也是直接从那里引用的。这里就不过多的介绍了，简单说下就是：

1. 分别遍历z、y、x的长度；
2. 以overlap size的方式，有重叠的进行裁剪；
3. 一个个patch组成patches列表。

详细介绍的文章在这里，点击去看：【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割9（patch 的 crop 和 merge 操作）

此时单个patch的数组大小为48x96x96大小，但是输入模型的，需要是[b, 1, 48, 96, 96]，其中b为batch size。就还需要进行一次预处理，将一个patch，转成对应大小的数组，然后转成Tensor。

对于一些patch在边缘的，可能还存在裁剪来的数组大小不够的情况，这就需要进行pad操作。代码如下，对这部分进行了记录。

def data_preprocess(img_crop, crop_size):
    if img_crop.shape != crop_size:
        pad_width = [(0, crop_size[0] - img_crop.shape[0]),
                     (0, crop_size[1] - img_crop.shape[1]),
                     (0, crop_size[2] - img_crop.shape[2])]
        img_crop = np.pad(img_crop, pad_width, mode='constant', constant_values=0)

    img_crop = np.expand_dims(img_crop, 0)  # (1, 16, 96, 96)
    img_crop = torch.from_numpy(img_crop).float()
    return img_crop/255.0
1
2
3
4
5
6
7
8
9
10

2.2、结果后处理

后处理恰恰与前处理相反。他需要将预测数组shape为[b, 2, 48, 96, 96]大小的数组，转为大小为[48, 96, 96]的数组。然后多个patch，按照逆过程，再merge在一起，组成一个和输入CT一样大小的数组。

代码如下：

                for d in range(0, patches.shape[0], 1):
                    img_crop = patches[d, :, :, :]
                    data = data_preprocess(img_crop, Config.Crop_Size)

                    data = data.unsqueeze(0).to(DEVICE)
                    output = model(data)  # (output.shape) torch.Size([b, class_num, 16, 96, 96])

                    data_cpu = data.clone().cpu()
                    output_cpu = output.clone().cpu()

                    i=0
                    img_tensor = data_cpu[i][0]  # 16 * 96 * 96
                    res_tensor = torch.gt(output_cpu[i][1], output_cpu[i][0])  # 16 * 96 * 96

                    patch_res = res_tensor.numpy()
                    patches_res_list.append(patch_res)

                mask_ai = res_merge(patches_res_list, volume_size, Config.OverLap_Size)
                nrrd.write(os.path.join(mask_ai_dir,  name + '.nrrd'), mask_ai)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

在本系列中，增加了一个背景类，于是就需要对包含背景类的channel与目标比类的channel做比较，得到最终的，包含目标的层。

torch.gt(Tensor1,Tensor2)
1

其中Tensor1和Tensor2为同维度的张量或者矩阵

含义：比较Tensor1和Tensor2的每一个元素,并返回一个0-1值。若Tensor1中的元素大于Tensor2中的元素,则结果取1，否则取0。

经过这样一个步骤，将包含背景类别channel=2的，变成channel为1的状态。比背景大，记为1，为前景；反着，则记为0，为背景。

2.3、预测代码

在这里，就完整的进行测试，将前面提到的需要经历所有步骤，统一进行了汇总。其中一些patch的crop和merge操作，你就参照上面提到的文章去看就可以了，调用的也是那个函数，比较好上手的。

看了那篇文章，即便没有本篇下面的代码，我相信你也能知道怎么搞了，没得担心的。

import os
import cv2
import numpy as np
import torch
import torch.utils.data
import matplotlib.pyplot as plt
import shutil
import nrrd
from tqdm import tqdm

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 没gpu就用cpu
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 屏蔽通知和警告信息
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 使用gpu0

def test():
    Config = Configuration()
    Config.display()

    results_dir = './results'
    mask_ai_dir = os.path.join(results_dir, 'pred_nrrd')

    if os.path.exists(results_dir):
        shutil.rmtree(results_dir)
    os.mkdir(results_dir)
    os.makedirs(mask_ai_dir, exist_ok=True)

    from models.unet3d_bn_activate import UNet3D
    model = UNet3D(num_out_classes=2, input_channels=1, init_feat_channels=32, testing=True)

    model_ckpt = torch.load(Config.model_path + "/best_model.pth")
    model.load_state_dict(model_ckpt)
    model = model.to(DEVICE)  # 模型部署到gpu或cpu里
    model.eval()

    with torch.no_grad():
        for idx, file in enumerate(tqdm(os.listdir(Config.valid_path))):
            if '_clean.nrrd' in file:
                name = file.split('_clean.nrrd')[0]
                nrrdClean_path = os.path.join(Config.valid_path, file)
                imgs, volume_size = load_img(nrrdClean_path)
                print('volume_size:', volume_size)

                # crop
                patches = crop_volume(imgs, Config.Crop_Size, Config.OverLap_Size)
                print('patches shape:', patches.shape)

                print(patches.shape)
                patches_res_list = []
                for d in range(0, patches.shape[0], 1):
                    img_crop = patches[d, :, :, :]
                    data = data_preprocess(img_crop, Config.Crop_Size)

                    data = data.unsqueeze(0).to(DEVICE)
                    output = model(data)  # (output.shape) torch.Size([b, class_num, 16, 96, 96])

                    data_cpu = data.clone().cpu()
                    output_cpu = output.clone().cpu()

                    i=0
                    img_tensor = data_cpu[i][0]  # 16 * 96 * 96
                    res_tensor = torch.gt(output_cpu[i][1], output_cpu[i][0])  # 16 * 96 * 96

                    patch_res = res_tensor.numpy()
                    patches_res_list.append(patch_res)

                mask_ai = res_merge(patches_res_list, volume_size, Config.OverLap_Size)
                nrrd.write(os.path.join(mask_ai_dir,  name + '.nrrd'), mask_ai)

class Configuration(object):
    valid_path = r"./database/valid"
    model_path = r'./checkpoints'

    Crop_Size = (48, 96, 96)
    OverLap_Size = [4, 8, 8]
    Num_Workers = 0

    def display(self):
        """Display Configuration values."""
        print("\nConfigurations:")
        print("")
        for a in dir(self):
            if not a.startswith("__") and not callable(getattr(self, a)):
                print("{:30} {}".format(a, getattr(self, a)))
        print("\n")

if __name__=='__main__':
    test()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

最终，我们输入的是一个CT的.nrrd的数组文件，最终预测结果也存储在了.nrrd的文件内。这个文件可以和标注文件做比较，进而对预测结果进行评估，也可以将预测结果打印出来，更加直观的在二维slice层面上进行查看。

三、结果可视化

现在假定，你是有了这批数据的CT数据、标注数据，和在二章节里面的预测结果，你可以有下面两种方式进行查看。

1. itk-snap软件直接查看nrrd文件，但是他一次只能查看CT数据和标注数据，或者CT数据、预测结果，同时打开两个窗口，实现联动也是可以的，如下面这样；

2. 也可以将标注数据和预测结果都以slice层的形式，绘制到一起，存储到本地，这样一层一层的查看。如下面这样：

第一种方式就不赘述了，直接下载itk-snap打开即可，这部分的资料比较多。

对于第二种，将标注和预测结果，按照slice都绘制到图像上面，，就稍微展开介绍下，将代码给到大家，自己可以使用。

1. 读取CT数组，标注数组和预测结果数组，都是nrrd文件；
2. 获取单层的slice，包括了上面三种类型数据的；
3. 分别将标注内容，预测内容，绘制到图像上；
4. 最后就是以图片的形式，存储到本地。

完整的代码如下：

import numpy as np
import nrrd
import os, cv2

def load_img(path_to_img):
    if path_to_img.startswith('LKDS'):
        img = np.load(path_to_img)
    else:
        img, _ = nrrd.read(path_to_img)

    return img, img.shape
    
def load_mask(path_to_mask):
    mask, _ = nrrd.read(path_to_mask)
    mask[mask > 1] = 1
    return mask, mask.shape

def getContours(output):
    img_seged = output.copy()
    img_seged = img_seged * 255

    # ---- Predict bounding box results with txt ----
    kernel = np.ones((5, 5), np.uint8)
    img_seged = cv2.dilate(img_seged, kernel=kernel)
    _, img_seged_p = cv2.threshold(img_seged, 127, 255, cv2.THRESH_BINARY)
    try:
        _, contours, _ = cv2.findContours(np.uint8(img_seged_p), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    except:
        contours, _ = cv2.findContours(np.uint8(img_seged_p), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    return contours

def drawImg(img, mask_ai, isAI=True):
    pred_oneImg = np.expand_dims(mask_ai, axis=2)
    contours = getContours(pred_oneImg)
    print(contours)
    if isAI:
        color = (0, 0, 255)
    else:
        color = (0, 255, 0)
    if len(contours) != 0:
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            xmin, ymin, xmax, ymax = x, y, x + w, y + h
            print('contouts:', xmin, ymin, xmax, ymax)
            # if isAI:
            #     cv2.drawContours(img, contour, -1, color, 2)
            # else:
            cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, thickness=2)

    return img

if __name__ == '__main__':
  ai_dir = r'./results/pred_nrrd'
  gt_dir = r'./database_nodule/valid'
  save_dir = r'./results/img_ai_gt'
  os.makedirs(save_dir, exist_ok=True)

  for filename in os.listdir(ai_dir):
      name = os.path.splitext(filename)[0]
      print(name, filename)
      ai_path = os.path.join(ai_dir, filename)
      gt_path = os.path.join(gt_dir, name + '_mask.nrrd')
      clean_path = os.path.join(gt_dir, name + '_clean.nrrd')

      imgs, volume_size = load_img(clean_path)

      masks_ai, masks_ai_shape = load_mask(ai_path)
      masks_gt, masks_gt_shape = load_mask(gt_path)

      assert volume_size==masks_ai_shape==masks_gt_shape

      for i in range(volume_size[0]):
          img = imgs[i, :, :]  # 获得第i张的单一数组
          mask_ai = masks_ai[i, :, :]
          mask_gt = masks_gt[i, :, :]

          print(np.max(img), np.min(img))

          img = np.expand_dims(img, axis=2)
          img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

          img = drawImg(img, mask_ai, isAI=True)
          img = drawImg(img, mask_gt, isAI=False)

          save_path = os.path.join(save_dir, name)
          os.makedirs(save_path, exist_ok=True)
          cv2.imwrite(os.path.join(save_path, '{}_img.png'.format(i)), img)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

四、总结

本文是继训练之后，对训练的模型进行独立的推理，实现对单个CT图像，经过patch操作，进行预测后，恢复成与原始CT一样大小的预测结果。并对预测结果和标注结果进行可视化对比，可以直观的看到对于单个检查，哪些结节是很容易的被识别到，而哪些比较的困难，哪些又是假阳性。

后面，就是对多个检查进行预测结果的评估，包括了结节级别的敏感度、特异度。在这样一个评估下，我们可以知道这个训练好的模型，究竟整体的性能怎么样，需不需要进一步的提高，从哪些角度进行提高？敬请期待。

---

最后，如果您觉得本篇文章对你有帮助，欢迎点赞 👍，让更多人看到，这是对我继续写下去的鼓励。如果能再点击下方的红包打赏，给博主来一杯咖啡，那就太好了。💪