关于输出解码的一些问题

[rmsuhang]

您好，我是视觉组一名成员，负责今年的自瞄，有幸看到您的代码，想问一下你们的代码中推理之后的output_buffer是（1，15120，20）吗，如果是的话，20的含义是四个点的坐标+class+confidence+标注的类别吗，如能回复不胜感激

[Harry-hhj]

你好，有关的代码位于autoaim/detector/TRTModule.cpp的void TRTModule::cache_engine(const std::string &cache_file)中，关键代码如下：

for (int i = 0; i < TOPK_NUM; i++) {
        auto *box_buffer = output_buffer + i * 20;  // 20->23
        if (box_buffer[8] < inv_sigmoid(KEEP_THRES)) break;
        if (removed[i]) continue;
        rst.emplace_back();
        auto &box = rst.back();
        memcpy(&box.pts, box_buffer, 8 * sizeof(float));
        for (auto &pt : box.pts) pt.x *= fx, pt.y *= fy;
        box.confidence = sigmoid(box_buffer[8]);
        box.color_id = argmax(box_buffer + 9, 4);
        box.tag_id = argmax(box_buffer + 13, 7);
        for (int j = i + 1; j < TOPK_NUM; j++) {
            auto *box2_buffer = output_buffer + j * 20;
            if (box2_buffer[8] < inv_sigmoid(KEEP_THRES)) break;
            if (removed[j]) continue;
            if (is_overlap(box_buffer, box2_buffer)) removed[j] = true;
        }
    }

模型的输出为8个归一化角点坐标、1个目标检测置信度、4种color_id的预测值、7个兵种类别的预测值。标签类别顺序及角点顺序如https://github.com/xinyang-go/LabelRoboMaster中定义。希望能解决你的问题。