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

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

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

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

这里，就要不得不先补齐下对数组crop和merge操作的方法了，建议先去学习下本系列的文章，链接：【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割9（patch 的 crop 和 merge 操作）。对于本系列的其他文章，也可直达专栏：人工智能与医学影像专栏。后期可不知道什么时候就收费了，先Mark住。

一、导言

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

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

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

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

二、模型预测

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

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

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

2.1、数据前处理

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

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

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

分别遍历z、y、x的长度；
以overlap size的方式，有重叠的进行裁剪；
一个个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

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=0img_tensor = data_cpu[i][0]  # 16 * 96 * 96res_tensor = torch.gt(output_cpu[i][1], output_cpu[i][0])  # 16 * 96 * 96patch_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)

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

torch.gt(Tensor1,Tensor2)

其中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 tqdmDEVICE = 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"  # 使用gpu0def 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 UNet3Dmodel = 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)# croppatches = 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=0img_tensor = data_cpu[i][0]  # 16 * 96 * 96res_tensor = torch.gt(output_cpu[i][1], output_cpu[i][0])  # 16 * 96 * 96patch_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 = 0def 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()

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

三、结果可视化

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

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

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

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

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

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

完整的代码如下：

import numpy as np
import nrrd
import os, cv2def 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.shapedef load_mask(path_to_mask):mask, _ = nrrd.read(path_to_mask)mask[mask > 1] = 1return mask, mask.shapedef 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 contoursdef 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 + hprint('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 imgif __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_shapefor 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)