UNet基于deepstream部署

0. sample背景介绍

deepstream-segmentation-test这个例子讲解了如何在deepstream中部署语义分割算法U-Net。根据网络的输出类别，分为semantic segmentation（输出四个类别）和industrial segmentation（单类别预测）。

源码路径：/opt/nvidia/deepstream/deepstream-6.0/sources/apps/sample_apps/deepstream-segmentation-test

1. 代码讲解

1.1 输入命令行解析

该sample中，执行推理运行的命令行格式如下：

./deepstream-segmentation-app [-t infer-type] config_file <file1> [file2] ... [fileN]

 第一个参数" -t "和第二个参数" infer-type "用于指定推理的网络类型[infer, inferserver]，默认情况下是[infer]状态。第三个参数是nvinfer推理需要指定的config文件路径，后续便是待处理的数据路径，因此常用的命令行如下：

// for semantic segmentation
./deepstream-segmentation-app dstest_segmentation_config_semantic.txt sample_720p.mjpeg sample_720p.mjpeg

// for industrial segmentation
./deepstream-segmentation-app dstest_segmentation_config_industrial.txt sample_industrial.jpg

上述输入命令行解析功能的实现，主要涉及如下的代码 

  /* Parse infer type and file names */
  for (gint k = 1; k < argc;)
  {
    if (!strcmp("-t", argv[k]))   // 第一个参数应该是 -t，用来指明推理的类别【infer, infer-server】
    {
      if (k + 1 >= argc)  
      {
        usage(argv[0]);
        return -1;
      }
      if (!strcmp("infer", argv[k+1])) 
      {
        is_nvinfer_server = FALSE;
      } 
      else if (!strcmp("inferserver", argv[k+1])) 
      {
        is_nvinfer_server = TRUE;
      } else 
      {
        usage(argv[0]);
        return -1;
      }
      k += 2;
    } 
    else if (!infer_config_file) // nvinfer的config文件路径
    {
      infer_config_file = argv[k];
      k++;
    } 
    else // 待处理数据的路径
    {
      files = g_list_append (files, argv[k]);
      num_sources++;
      k++;
    }
  }

1.2 Pipeline构建

该sample中，pipeline添加链接的插件流程如下：

filesrc -> jpegparse -> nvv4l2decoder -> nvstreammux -> nvinfer/nvinferserver (segmentation)

nvsegvidsual -> nvmultistreamtiler -> (nvegltransform) -> nveglglessink. 

下面对pipeline中一些主要的elements进行介绍：

1. jpegparse element

由于在pipeline中采用jpegparse解析器，因此该sample仅仅支持jpg图片和mjpeg格式的视频输入。在DeepStream中如何使用image作为输入进行推理，SDK中也提供了一个example，deepstream-image-decode-test，具体使用方法可以参考博客：【deepstream-image-decode-test解析】。

2. nvinfer element 

在之前的应用程序中，nvinfer 插件常用来进行目标检测和分类模型的推理，而该sample中模型的任务是语义分割，因此，在配置文件中一些参数与之前的有所区别，重要参数如下：

 3. nvsegvidsual element 

在语义分割任务中，nvsegvidsual插件用来替代nvosd，来渲染最终mask的输出。此时，输出mask结果的大小与网络输入大小相同。

上图为nvsegvidsual的渲染输出，同时，我们需要注意的是nvsegvidsual插件的输入中包含NvDsInferSegmentationMeta对象，描述了最后推理输出mask的信息。

2. 网络的输出

对于语义分割任务，网络的输出维度一般是【N, C, H, W】，其中N表示批处理的大小，C表示数据集的类别数，H和W为输出图像的高宽，通常情况下与网络输入大小相等，其内容描述了对于每个像素属于各个类别的概率。

因此，后续的处理流程是：首先给定一个阈值判断该像素点属于前景还是背景，如果是前景的话，找到概率值最大的位置，表示其类别。然后就能得到一个class map，描述了每个像素点的类别。最后，我们通过一个color map为每个类别设置不同的颜色就能得到最终的mask图像。

 上图描述了在deepstream中，语义分割模型在推理过程中，各阶段工作完成相应的位置。因此，我们在部署语义分割模型时，需要保证模型的输出维度应该【N, C, H, W】，表示每个像素属于所有类别的概率。然后通过config文件中配置的segmentation-threshold，得到相应的class map.

3. 获取输出信息

通过该sample给出的pipeline，我们能得到一张经过渲染的mask图像作为最终的结果，但有时我们并不想只是得到mask图像就完事了，更多的情况下我们想要对输出的概率和类别信息进行进一步的分析，那么如何在deepstream中获取语义分割网络的输出信息呢？

下图为nvinfer插件的输出描述，因此我们知道语义分割网络的输出信息应该保存在NvDsInferSegmentationMeta中。

 下图为NvDsInferSegmentationMeta的介绍，需要重点注意的是，该Meta是以NvDsUserMeta 的形式链接到frame_meta的frame_user_meta_list 或者obj_meta的object_user_meta_list 中，且其meta_type被定义为NVDSINFER_SEGMENTATION_META。

 此外，还想对NvDsInferSegmentationMeta中class_map和class_probabilities_map两个重要参数的含义进行解释，class_map是二维数组，表示每个像素值预测的类别，其中-1代表背景，0-n代表n个类别数，需要重点注意的是deepstream在做语义分割网络推理过程中的类别数是实际的类别数，即不包含背景。class_probabilities_map应该就是engine的输出，[c, h, w]维度，表示每个像素属于所有类别的概率。

因此，我们可以通过如下的代码链接到网络的输出：

static GstPadProbeReturn
tiler_src_pad_buffer_probe (GstPad * pad, GstPadProbeInfo * info,
    gpointer u_data)
{
    GstBuffer *buf = (GstBuffer *) info->data;
    NvDsMetaList * l_frame = NULL;
    NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta (buf);

    for (l_frame = batch_meta->frame_meta_list; l_frame != NULL;
      l_frame = l_frame->next) {
        NvDsFrameMeta *frame_meta = (NvDsFrameMeta *) (l_frame->data);
        NvDsUserMetaList *usrMetaList = frame_meta->frame_user_meta_list;
         while (usrMetaList != NULL) 
         {
            NvDsUserMeta *usrMetaData = (NvDsUserMeta *) usrMetaList->data;
            if (usrMetaData->base_meta.meta_type == NVDSINFER_SEGMENTATION_META) 
            {
              NvDsInferSegmentationMeta* segMetaData = (NvDsInferSegmentationMeta*)(usrMetaData->user_meta_data);
              int classes = segMetaData->classes;
              int width = segMetaData->width;
              int height = segMetaData->height;
              int* class_ids = segMetaData->class_map;
              float* scores = segMetaData->class_probabilities_map;
              g_print("segmentation. classes: %d width: %d, height: %d\n", classes, width, height);
              usrMetaList = usrMetaList->next;
            } 
            else 
            {
              g_print("others.\n");
              usrMetaList = usrMetaList->next;
          }
        }
    }
    return GST_PAD_PROBE_OK;
}

4. 部署自定义的UNet模型

 在博客U-Net基于TensorRT部署中，介绍了如何在tensorrt中部署unet模型，同样基于这篇博客生成的engine文件可用于deepstream中，来实现在deepstream上部署自定义的UNet模型。

当使用Pytorch训练好UNet模型，然后得到相应的onnx模型，由于UNet网络中不包含某些复杂操作，一种简便得到TensorRT engine的方法就是采用自带的trtexec工具。

./trtexec --onnx=path-to-onnx-file  --workspace=4096 --saveEngine=path-to-save-engine

 修改deepstream_segmentation_app.cpp代码中，关于图像和模型维度的设置，重新编译。

// 图像统一resize维度
#define MUXER_OUTPUT_WIDTH 1280
#define MUXER_OUTPUT_HEIGHT 720


// network输入维度
#define TILED_OUTPUT_WIDTH 960
#define TILED_OUTPUT_HEIGHT 640

对于UNet网络，nvinfer的配置文件内容应该如下：

# dstest_unet_config_industrial.txt
[property]
gpu-id=0
net-scale-factor=0.003921568627451
maintain-aspect-ratio=0
scaling-filter=1
scaling-compute-hw=0
# Integer 0: RGB 1: BGR 2: GRAY
model-color-format=0

# onnx-file=unet.onnx
model-engine-file=unet.trt

infer-dims=3;640;960

force-implicit-batch-dim=1
batch-size=1
## 0=FP32, 1=INT8, 2=FP16 mode
network-mode=0

num-detected-classes=1

interval=0

gie-unique-id=1

network-type=2
process-mode=1
segmentation-threshold=0.5


[class-attrs-all]
pre-cluster-threshold=0.5
roi-top-offset=0
roi-bottom-offset=0
detected-min-w=0
detected-min-h=0
detected-max-w=0
detected-max-h=0

 运行如下命令进行推理（只接收jpep和mjpeg格式数据输入）：

./deepstream-segmentation-app dstest_unet_config_industrial.txt image.jpg

5. 有用的链接

U-Net基于TensorRT部署_hello_dear_you的博客-CSDN博客_unet部署

https://forums.developer.nvidia.com/t/steps-to-use-my-custom-segmentation-model-in-deepstream/159999/11

fredlsousa/deepstream-test1-segmentation: Modified deepstream-test1 sample app to accept segmentation models and output the masks. There's also a CMake file for those who like it better than Makefiles. (github.com)