Main Content

使用 U-Net 的语义分割网络的代码生成

此示例说明如何为使用深度学习的图像分割应用程序生成代码。它使用 codegen 命令生成一个基于 U-Net 的 DAG 网络对象执行预测的 MEX 函数。U-Net 是一种用于图像分割的深度学习网络。

有关使用 U-Net 进行图像分割但不使用 codegen 命令的类似示例,请参阅使用深度学习对多光谱图像进行语义分割 (Image Processing Toolbox)

第三方前提条件

必需

此示例生成 CUDA MEX,并具有以下第三方要求。

  • CUDA® 支持 NVIDIA® GPU 和兼容驱动程序。

可选

对于非 MEX 编译,如静态、动态库或可执行文件,此示例有以下附加要求。

验证 GPU 环境

使用 coder.checkGpuInstall (GPU Coder) 函数验证运行此示例所需的编译器和库是否已正确设置。

envCfg = coder.gpuEnvConfig('host');
envCfg.DeepLibTarget = 'cudnn';
envCfg.DeepCodegen = 1;
envCfg.Quiet = 1;
coder.checkGpuInstall(envCfg);

分割网络

U-Net [1] 是一种卷积神经网络 (CNN),专为语义图像分割而设计。U-Net 中的初始卷积层序列与最大池化层交叠,从而逐步降低输入图像的分辨率。这些层后跟一系列使用上采样算子散布的卷积层,从而会连续增加输入图像的分辨率。将这两个系列路径组合在一起将形成一个 U 形图。该网络最初是为生物医学图像分割应用执行预测而训练和使用的。本示例中演示的是该网络跟踪森林植被随时间变化的能力。环保机构跟踪森林采伐,以评估和鉴定地区的环境和生态健康。

基于深度学习的语义分割可以通过高分辨率航拍照片精确测量植被覆盖度。一个挑战是区分具有相似视觉特性的类,例如尝试将绿色像素分类为草、灌木或树。为了提高分类准确度,一些数据集包含多光谱图像,这些图像提供关于每个像素的附加信息。例如,Hamlin Beach 国家公园数据集用近红外通道为彩色图像补充信息,以提供更清晰的分类。

本示例使用 Hamlin Beach 国家公园数据 [2] 以及预训练 U-Net 网络,以便正确地对每个像素进行分类。

使用的 U-Net 经过训练,用于对属于以下 18 个类的像素进行分割:

0. Other Class/Image Border      7. Picnic Table         14. Grass
1. Road Markings                 8. Black Wood Panel     15. Sand
2. Tree                          9. White Wood Panel     16. Water (Lake)
3. Building                     10. Orange Landing Pad   17. Water (Pond)
4. Vehicle (Car, Truck, or Bus) 11. Water Buoy           18. Asphalt (Parking Lot/Walkway)
5. Person                       12. Rocks
6. Lifeguard Chair              13. Other Vegetation

segmentImageUnet 入口函数

segmentImageUnet.m 入口函数使用 multispectralUnet.mat 文件中的 multispectralUnet 网络对输入图像执行逐块语义分割。该函数将 multispectralUnet.mat 文件中的网络对象加载到持久变量 mynet 中,并在后续的预测调用中重用该持久变量。

type('segmentImageUnet.m')
function out = segmentImageUnet(im,patchSize,trainedNet)  
%  OUT = segmentImageUnet(IM,patchSize,trainedNet) returns a semantically
%  segmented image, segmented using the multi-spectral Unet specified in
%  trainedNet. The segmentation is performed over each patch of size
%  patchSize.
%
% Copyright 2019-2022 The MathWorks, Inc.

%#codegen
persistent mynet;

if isempty(mynet)
    mynet = coder.loadDeepLearningNetwork(trainedNet);
end

[height, width, nChannel] = size(im);
patch = coder.nullcopy(zeros([patchSize, nChannel-1]));

% Pad image to have dimensions as multiples of patchSize
padSize = zeros(1,2);
padSize(1) = patchSize(1) - mod(height, patchSize(1));
padSize(2) = patchSize(2) - mod(width, patchSize(2));

im_pad = padarray (im, padSize, 0, 'post');
[height_pad, width_pad, ~] = size(im_pad);

out = zeros([size(im_pad,1), size(im_pad,2)], 'uint8');

for i = 1:patchSize(1):height_pad    
    for j =1:patchSize(2):width_pad        
        for p = 1:nChannel-1              
            patch(:,:,p) = squeeze( im_pad( i:i+patchSize(1)-1,...
                                            j:j+patchSize(2)-1,...
                                            p));            
        end
         
        % Pass in input
        segmentedLabels = activations(mynet, patch, 'Segmentation-Layer');
        
        % Takes the max of each channel (6 total at this point)
        [~,L] = max(segmentedLabels,[],3);
        patch_seg = uint8(L);
        
        % Populate section of output
        out(i:i+patchSize(1)-1, j:j+patchSize(2)-1) = patch_seg;
       
    end
end

% Remove the padding
out = out(1:height, 1:width);

获取预训练的 U-Net 网络

此示例使用包含预训练 U-Net 网络的 multispectralUnet MAT 文件。此文件大小约为 117 MB。从 MathWorks 网站下载该文件。

trainedUnetFile = matlab.internal.examples.downloadSupportFile('vision/data','multispectralUnet.mat');

U-Net 是一种包含 58 个层的 DAG 网络,其中包括卷积层、最大池化层、深度串联层和像素分类输出层。

load(trainedUnetFile);
disp(net)
  DAGNetwork with properties:

         Layers: [58×1 nnet.cnn.layer.Layer]
    Connections: [61×2 table]
     InputNames: {'ImageInputLayer'}
    OutputNames: {'Segmentation-Layer'}

要查看网络架构,请使用 analyzeNetwork 函数。

analyzeNetwork(net);

准备数据

此示例使用来自 [2] 的高分辨率多光谱数据。图像集是用无人机在纽约 Hamlin Beach 州立公园拍摄的。数据包含已标注的训练集、验证集和测试集,有 18 个对象类标签。数据文件的大小约为 3.0 GB。

使用 downloadHamlinBeachMSIData 辅助函数下载数据集的 MAT 文件版本。此函数作为支持文件包含在本示例中。

if ~exist(fullfile(pwd,'data'),'dir')
    url = 'https://home.cis.rit.edu/~cnspci/other/data/rit18_data.mat';
    downloadHamlinBeachMSIData(url,pwd+"/data/");
end

在 MATLAB 中加载并检查数据。

load(fullfile(pwd,'data','rit18_data','rit18_data.mat'));

% Examine data
whos test_data
  Name           Size                         Bytes  Class     Attributes

  test_data      7x12446x7654            1333663576  uint16              

该图像有七个通道。RGB 颜色通道是第三个、第二个和第一个图像通道。接下来的三个通道对应于近红外波段,并基于其热能特征突出显示图像的不同分量。通道 7 是指示有效分割区域的掩膜。

多光谱图像数据排列为通道数×宽度×高度数组。在 MATLAB 中,多通道图像排列为宽×高×通道数数组。要重构数据以使通道处于第三个维度中,请使用辅助函数 switchChannelsToThirdPlane

test_data  = switchChannelsToThirdPlane(test_data);

% Confirm data has the correct structure (channels last).
whos test_data
  Name               Size                     Bytes  Class     Attributes

  test_data      12446x7654x7            1333663576  uint16              

运行 MEX 代码生成

要为 segmentImageUnet.m 入口函数生成 CUDA 代码,请为 MEX 目标创建一个 GPU 代码配置对象,并将目标语言设置为 C++。使用 coder.DeepLearningConfig (GPU Coder) 函数创建一个 CuDNN 深度学习配置对象,并将其赋给 GPU 代码配置对象的 DeepLearningConfig 属性。运行 codegen 命令,指定输入大小为 12446×7654×7,补片大小为 1024×1024。这些值对应于整个 test_data 大小。较小的补片大小可加速推断。要了解如何计算补片,请参阅 segmentImageUnet 入口函数。

cfg = coder.gpuConfig('mex');
cfg.ConstantInputs = 'Remove';
cfg.TargetLang = 'C++';
cfg.DeepLearningConfig = coder.DeepLearningConfig('cudnn');
inputArgs = {ones(size(test_data),'uint16'),...
    coder.Constant([1024 1024]),coder.Constant(trainedUnetFile)};

codegen -config cfg segmentImageUnet -args inputArgs -report
Code generation successful: View report

运行生成的 MEX 来预测 test_data 的结果

segmentImageUnet 函数接受要测试的数据 (test_data) 和一个向量,向量中包含要使用的补片的大小。获取图像的补片,预测特定补片中的像素,然后将所有补片组合在一起。由于测试数据的大小为 12446×7654×7,以补片形式处理如此大的图像更容易。

segmentedImage = segmentImageUnet_mex(test_data);

为了只提取分割的有效部分,将分割的图像乘以测试数据的封装通道。

segmentedImage = uint8(test_data(:,:,7)~=0) .* segmentedImage;

由于语义分割的输出含有噪声,请使用 medfilt2 函数去除噪声和杂散像素。

segmentedImage = medfilt2(segmentedImage,[5,5]);

显示 U-Net 分割后的 test_data

以下代码行创建类名的向量。

classNames = [ "RoadMarkings","Tree","Building","Vehicle","Person", ...
               "LifeguardChair","PicnicTable","BlackWoodPanel",...
               "WhiteWoodPanel","OrangeLandingPad","Buoy","Rocks",...
               "LowLevelVegetation","Grass_Lawn","Sand_Beach",...
               "Water_Lake","Water_Pond","Asphalt"];

在分割的 RGB 测试图像上叠加标签,并向分割图像添加颜色栏。

cmap = jet(numel(classNames));
B = labeloverlay(imadjust(test_data(:,:,[3,2,1]),[0 0.6],[0.1 0.9],0.55),...
    segmentedImage,'Transparency',0.8,'Colormap',cmap);
figure
imshow(B)

N = numel(classNames);
ticks = 1/(N*2):1/N:1;
colorbar('TickLabels',cellstr(classNames),'Ticks',ticks,'TickLength',0,...
    'TickLabelInterpreter','none');
colormap(cmap)
title('Segmented Image');

Figure contains an axes object. The axes object with title Segmented Image contains an object of type image.

参考资料

[1] Ronneberger, Olaf, Philipp Fischer, and Thomas Brox."U-Net:Convolutional Networks for Biomedical Image Segmentation." arXiv preprint arXiv:1505.04597, 2015.

[2] Kemker, R., C. Salvaggio, and C. Kanan."High-Resolution Multispectral Dataset for Semantic Segmentation."CoRR, abs/1703.01918, 2017.

[3] Kemker, Ronald, Carl Salvaggio, and Christopher Kanan."Algorithms for Semantic Segmentation of Multispectral Remote Sensing Imagery Using Deep Learning."ISPRS Journal of Photogrammetry and Remote Sensing, Deep Learning RS Data, 145 (November 1, 2018):60-77. https://doi.org/10.1016/j.isprsjprs.2018.04.014.