这里我们也只针对 ‘car’ 的类别进行图像分割测试，运行代码后测试场景如图所示 。左边为原图，中间为标注的图像分割场景，右边为测试的图像分歌结果，可以看出测试的结果和标注的图像分割场景还是很相似的，说明训练结果还不错 。
3.3 UNet++图像分割检测
通过测试后，我们可以使用我们训练好后UNet++模型，进行图像分割的实际操作 。检测代码如下：
import osos.environ['CUDA_VISIBLE_DEVICES'] = '0'import numpy as npimport cv2import matplotlib.pyplot as pltimport albumentations as albuimport torchimport segmentation_models_pytorch as smpfrom torch.utils.data import Dataset as BaseDatasetimport imageio# ---------------------------------------------------------------### Dataloaderclass Dataset(BaseDataset):"""CamVid数据集 。进行图像读取，图像增强增强和图像预处理.Args:images_dir (str): 图像文件夹所在路径masks_dir (str): 图像分割的标签图像所在路径class_values (list): 用于图像分割的所有类别数augmentation (albumentations.Compose): 数据传输管道preprocessing (albumentations.Compose): 数据预处理"""# CamVid数据集中用于图像分割的所有标签类别CLASSES = ['sky', 'building', 'pole', 'road', 'pavement','tree', 'signsymbol', 'fence', 'car','pedestrian', 'bicyclist', 'unlabelled']def __init__(self,images_dir,# masks_dir,classes=None,augmentation=None,preprocessing=None,):self.ids = os.listdir(images_dir)self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]# convert str names to class values on masksself.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]self.augmentation = augmentationself.preprocessing = preprocessingdef __getitem__(self, i):# read dataimage = cv2.imread(self.images_fps[i])image = cv2.resize(image, (480, 384))# 改变图片分辨率image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# 图像增强应用if self.augmentation:sample = self.augmentation(image=image)image = sample['image']# 图像预处理应用if self.preprocessing:sample = self.preprocessing(image=image)image = sample['image']return imagedef __len__(self):return len(self.ids)# ---------------------------------------------------------------def get_validation_augmentation():"""调整图像使得图片的分辨率长宽能被32整除"""test_transform = [albu.PadIfNeeded(384, 480)]return albu.Compose(test_transform)def to_tensor(x, **kwargs):return x.transpose(2, 0, 1).astype('float32')def get_preprocessing(preprocessing_fn):"""进行图像预处理操作Args:preprocessing_fn (callbale): 数据规范化的函数(针对每种预训练的神经网络)Return:transform: albumentations.Compose"""_transform = [albu.Lambda(image=preprocessing_fn),albu.Lambda(image=to_tensor),]return albu.Compose(_transform)# 图像分割结果的可视化展示def visualize(**images):"""PLot images in one row."""n = len(images)plt.figure(figsize=(16, 5))for i, (name, image) in enumerate(images.items()):plt.subplot(1, n, i + 1)plt.xticks([])plt.yticks([])plt.title(' '.join(name.split('_')).title())plt.imshow(image)plt.show()# ---------------------------------------------------------------if __name__ == '__main__':DATA_DIR = './data/CamVid/'x_test_dir = os.path.join(DATA_DIR, 'abc')img_test = cv2.imread('data/CamVid/abc/car_test.jpg')height = img_test.shape[0]weight = img_test.shape[1]ENCODER = 'se_resnext50_32x4d'ENCODER_WEIGHTS = 'imagenet'CLASSES = ['car']ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multiclass segmentationDEVICE = 'cuda'# 按照权重预训练的相同方法准备数据preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)# 加载最佳模型best_model = torch.load('./best_model.pth')# 创建检测数据集predict_dataset = Dataset(x_test_dir,augmentation=get_validation_augmentation(),preprocessing=get_preprocessing(preprocessing_fn),classes=CLASSES,)# 对检测图像进行图像分割并进行图像可视化展示predict_dataset_vis = Dataset(x_test_dir,classes=CLASSES,)for i in range(len(predict_dataset)):# 原始图像image_visimage_vis = predict_dataset_vis[i].astype('uint8')image = predict_dataset[i]# 通过图像分割得到的0-1图像pr_maskx_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)pr_mask = best_model.predict(x_tensor)pr_mask = (pr_mask.squeeze().cpu().numpy().round())print(pr_mask.shape)# 恢复图片原来的分辨率image_vis = cv2.resize(image_vis, (weight, height))pr_mask = cv2.resize(pr_mask, (weight, height))# 保存图像分割后的黑白结果图像imageio.imwrite('car_test_out.png', pr_mask)# 原始图像和图像分割结果的可视化展示visualize(image=image_vis,predicted_mask=pr_mask)