GRFB UNet——基于多尺度注意网络盲道检测算法实现与模型C++部署

1. 概述

盲道是视障人士安全出行的重要辅助设施。识别盲道的形状和位置，对于增强视障人士的自主移动能力至关重要，而视觉分割技术正是应对这一挑战的有效工具。为了显著提升盲道分割的精确度和稳定性，本文提出了一种创新的分割方法，该方法融合了UNet网络与多尺度特征提取技术。本方法在UNet架构中引入了组感受野块（GRFB）的设计，用以捕获盲道的多级视觉信息。通过应用组卷积，该方法有效降低了计算的复杂度。此外，在每个组卷积之后引入了小尺度卷积，以促进不同通道间的信息交流和融合，进而提取更为丰富和高层次的特征。

在本研究中，我们构建并标注了一个包含多种环境条件下盲道的数据集，用以进行实验评估。我们还对本方法与现有的典型网络结构和模块进行了详尽的比较分析。实验结果表明，我们提出的网络在盲道分割任务上的表现超越了其他对比网络，为盲道的检测提供了一个有力的参考，这不仅证明了本方法的有效性，也为视障人士的导航辅助技术的发展做出了贡献。

训练代码：https://github.com/Chon2020/GRFB-Unet

2. 网络架构

GRFB-UNet网络的结构是指在传统的UNet网络基础上增加了组感受野块（Group Receptive Field Block，简称GRFB）的改进型网络。GRFB的设计旨在通过组卷积来捕获图像中的多尺度特征，从而提高网络对不同尺度目标的识别能力。
（1）. 输入层：网络接收输入图像，并开始进行特征提取。
（2）. GRFB结构：在UNet的多个阶段中引入GRFB，每个GRFB由多个组卷积层组成，这些层可以并行地从不同尺度捕获图像特征。
（3）. 组卷积：GRFB中的组卷积允许网络在每个组内独立地学习特征，这有助于网络专注于不同的空间尺度。
（4）. 跨组卷积：在组卷积之后，使用小尺度的卷积核进行跨组卷积，以实现组间的信息整合。
（5）. 特征融合：UNet网络的上采样和跳跃连接有助于将低层的高分辨率特征与高层的抽象特征进行融合，增强了特征的表达能力。
（6）. 多尺度特征提取：通过GRFB结构，网络能够同时提取不同尺度的特征，这对于理解图像中的局部和全局上下文非常重要。
（7）. 输出层：网络的最终输出是一个分割图，它将输入图像中的每个像素分类为属于或不属于目标类别（例如，盲道）。

GRFB-UNet网络的设计特别适合于需要精确定位和多尺度特征提取的图像分割任务，如触觉铺路的分割。通过这种结构，网络能够更好地理解和处理图像中的复杂结构，从而提高分割的准确性和鲁棒性。

3.环境安装

3.1 环境安装

conda create -n py python=3.7
conda activate py
conda install pytorch==1.5.1 torchvision==0.6.1 cudatoolkit=10.1 -c pytorch
pip install -r requirements.txt

3.2 模型转换

import torch
import numpy as np
import onnx
from models.unet import GRFBUNet

def to_onnx(model_unet):
    with torch.no_grad():
        img_unet = torch.randn(1,3,640,640)
        outputs_unet = model_unet(img_unet)

        # 导出ONNX文件
        torch.onnx.export(
            model_unet,
            img_unet,
            'grfb_unet.onnx',
            opset_version=11,
            input_names=['input'],
            output_names=['output']
        )

        # prediction = outputs_unet['out'].argmax(1)
        
    return 


def main():
    # 加载unet模型
    model_unet = GRFBUNet(in_channels=3, num_classes=2, base_c=32)
    model_unet.load_state_dict(torch.load('./weights/grfb-unet.pth', map_location='cpu')['model'])
    model_unet.eval() 
    to_onnx(model_unet)


if __name__ == "__main__":
    main()

4. 模型推理

#include <provider_options.h>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>


int main()
{
    int gpu_index = 0;
    int gpu_ram = 4;
    int num_thread = 4;
    std::string model_path = "grfb_unet.onnx";
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    std::vector<std::string> available_providers = Ort::GetAvailableProviders();
    auto cuda_available = std::find(available_providers.begin(),
        available_providers.end(), "CUDAExecutionProvider");

    Ort::SessionOptions session_options = Ort::SessionOptions();

    session_options.SetInterOpNumThreads(num_thread);
    session_options.SetIntraOpNumThreads(num_thread);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    if (gpu_index >= 0 && (cuda_available != available_providers.end()))
    {
        OrtCUDAProviderOptions cuda_options;
        cuda_options.device_id = gpu_index;
        cuda_options.arena_extend_strategy = 0;

        if (gpu_ram == -1)
        {
            cuda_options.gpu_mem_limit = ~0ULL;
        }
        else
        {
            cuda_options.gpu_mem_limit = gpu_ram * 1024 * 1024 * 1024;
        }

        cuda_options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::OrtCudnnConvAlgoSearchExhaustive;
        cuda_options.do_copy_in_default_stream = 1;

        session_options.AppendExecutionProvider_CUDA(cuda_options);
    }

    float mean[3] = { 0.709, 0.381, 0.224 };
    float std[3] = { 0.127, 0.079, 0.043 };
    try 
    {
        // 加载模型并创建环境空间
        std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "grfb_unet");
        Ort::Session session(env, widestr.c_str(), session_options);
        Ort::AllocatorWithDefaultOptions allocator;

        char input[] = "input";
        char output[] = "output";

        std::vector<char*> input_node_names;
        std::vector<char*> output_node_names;

        input_node_names.push_back(input);
        output_node_names.push_back(output);

      
        std::string path = "images";
        std::vector<std::string> filenames;
        cv::glob(path, filenames, false);

        for (auto imgpath : filenames)
        {
            cv::Mat original_image = cv::imread(imgpath, cv::IMREAD_COLOR);
            cv::Mat resized_image;
            cv::resize(original_image, resized_image, cv::Size(640, 640));


            // 确定输入数据维度
            std::vector<int64_t> input_node_dims = { 1,3,640,640 };
            size_t input_tensor_size = 1 * 3 * 640 * 640;

            // 填充数据输入
            std::vector<float> input_tensor_values(input_tensor_size);
            for (int h = 0; h < 640; ++h)
            {
                for (int w = 0; w < 640; ++w)
                {
                    for (int c = 0; c < 3; ++c)
                    {
                        // 均一化像素值
                        float pix = resized_image.at<cv::Vec3b>(h, w)[c];
                        pix = pix / 255.0f;
                        pix = (pix - mean[c]) / std[c];
                        input_tensor_values[640 * 640 * c + h * 640 + w] = pix;
                    }
                }
            }

            Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
                memory_info,
                input_tensor_values.data(),
                input_tensor_size,
                input_node_dims.data(),
                input_node_dims.size()
                );

            assert(input_tensor.IsTensor());

            std::vector<Ort::Value> ort_inputs;
            ort_inputs.push_back(std::move(input_tensor));


            // 启动模型预测并获取输出张量
            auto output_tensors = session.Run(
                Ort::RunOptions{ nullptr },
                input_node_names.data(),
                ort_inputs.data(),
                ort_inputs.size(),
                output_node_names.data(),
                1
            );

            // 解析输出张量
            Ort::Value& output_tensor = output_tensors[0];
            const float* output_data = output_tensor.GetTensorData<float>();
            std::vector<int64_t> output_dims = output_tensor.GetTensorTypeAndShapeInfo().GetShape();

            // 存储输出图像
            cv::Mat result_image(640, 640, CV_8UC1);

            // 对输出的2通道图像进行二分类预测
            for (int h = 0; h < 640; ++h) {
                for (int w = 0; w < 640; ++w) {
                    int index_max = output_data[w + h * 640] > output_data[w + h * 640 + 640 * 640] ? 0 : 1;
                    result_image.at<uchar>(h, w) = 255 * index_max;
                }
            }

            cv::imshow("Resized Image", resized_image);
            // 显示结果
            cv::imshow("Result Image", result_image);
            cv::waitKey(0);
        }

    }
    catch (const Ort::Exception& e) 
    {
        // 打印异常
        std::cerr << "Caught Ort::Exception: " << std::string(e.what()) << std::endl;
        size_t pos = std::string(e.what()).find("ErrorCode: ");
        if (pos != std::string::npos) {
            std::string error_code_str = std::string(e.what()).substr(pos + 12);
            int error_code = std::stoi(error_code_str);
            std::cerr << "Error Code: " << error_code << std::endl;
        }
        return -1;
    }

    return 0;
}