下载
中文
注册

自定义模型后处理开发步骤(tensorinfer框架)

步骤介绍

图1 已支持模型的后处理UML类图

  1. 根据任务类型,选择Vision SDK已经支持的后处理基类去派生一个新的子类,后处理基类分别为目标检测、分类任务、语义分割、文本生成,这些基类都继承自相同的父类PostProcessBase。

    在Init函数中,需要调取父类的Init()接口读取配置参数,然后再调用父类对象的configData_.GetFileValue()接口读取子类所需要的配置参数。

    在Process函数中,需要以模型的输出张量Tensors作为输入,以对应的数据结构类对象作为输出。

  2. 调用父类的CheckAndMoveTensors()接口,对Tensors的形状进行校验并将内存移动至Host侧,然后再进行相应操作获得结果。
  3. 后处理开发完成后,增加一个对外的接口如GetObjectInstance(),以便于让业务流中的后处理插件动态加载此后处理so。

    如果当前后处理基类所采用的数据结构无法满足需求,可以新增后处理基类继承PostProcessBase,并写入新的数据结构。

    // 1.目标检测数据结构:
    class ObjectInfo {
    public:
        float x0;
        float y0;
        float x1;
        float y1;
        float confidence;
        float classId;
        std::string className;
        std::vector<std::vector <int>> mask; // 用于实例分割
    };
    
    // 2.分类任务数据结构:
    class ClassInfo {
    public:
        int classId;
        float confidence;
        std::string className;
    };
    
    // 3.语义分割任务数据结构:
    class SemanticSegInfo {
    public:
        std::vector<std::vector<int>> pixels;
        std::vector<std::string> labelMap;
    };
    
    // 4.文本生成(机器翻译,文字识别,语音识别等)数据结构:
    class TextsInfo {
    public:
        std::vector<std::string> text;
    };
    
    // 5.文本框检测数据结构:
    class TextObjectInfo {
    public:
        float x0;
        float y0;
        float x1;
        float y1;
        float x2;
        float y2;
        float x3;
        float y3;
        float confidence;
        std::string result;
    };