yolov8seg 瑞芯微RKNN部署

首先是把yolov8的onnx模型转成rknn模型，这里用的是yolov8n-seg.

转模型代码如下，这段是python代码：

if __name__ == '__main__':

    platform = 'rkXXXX' #写自己的型号
    exp = 'yolov8n_seg'
    Width = 640
    Height = 640
    MODEL_PATH = './onnx_models/yolov8n-seg.onnx'
    NEED_BUILD_MODEL = True
    # NEED_BUILD_MODEL = False
    im_file = './dog_bike_car_640x640.jpg'

    # Create RKNN object
    rknn = RKNN(verbose=False)

    OUT_DIR = "rknn_models"
    RKNN_MODEL_PATH = '{}/{}.rknn'.format(OUT_DIR,exp)
    if NEED_BUILD_MODEL:
        DATASET = './dataset.txt'
        rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform="rkXXXX")
        # Load model
        print('--> Loading model')
        ret = rknn.load_onnx(MODEL_PATH)
        if ret != 0:
            print('load model failed!')
            exit(ret)
        print('done')

        # Build model
        print('--> Building model')
        ret = rknn.build(do_quantization=True, dataset=DATASET)
        if ret != 0:
            print('build model failed.')
            exit(ret)
        print('done')

        # Export rknn model
        if not os.path.exists(OUT_DIR):
            os.mkdir(OUT_DIR)
        print('--> Export RKNN model: {}'.format(RKNN_MODEL_PATH))
        ret = rknn.export_rknn(RKNN_MODEL_PATH)
        if ret != 0:
            print('Export rknn model failed.')
            exit(ret)
        print('done')
    else:
        ret = rknn.load_rknn(RKNN_MODEL_PATH)

    rknn.release()

运行成功之后会得到yolov8n_seg.rknn模型。

往下是cpp代码。

运行rknn模型，不是零copy的话用这段代码。

//设置outputs数组保存结果
rknn_output outputs[io_num.n_output];  //长度为2的数组
memset(outputs, 0, sizeof(outputs));
for (int i = 0; i < io_num.n_output; i++) {
  outputs[i].want_float = 0;
}

ret = rknn_run(ctx, NULL); //运行rknn模型
//这里不是零copy方法，需要用到rknn_outputs_get函数
ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);

零copy方法用这段代码。

ret = rknn_run(ctx, nullptr);

memcpy(y0, output_mems[0]->virt_addr, output_attrs[0].n_elems * sizeof(char)); //outputs[0]copy到y0
memcpy(y1, output_mems[1]->virt_addr, output_attrs[1].n_elems * sizeof(char));

下面来看下rknn模型运行结果outputs的结构。

具体这个结果怎么来的以及详细含义的请参考此处。

outputs的size为2，其中outputs[0]的size是8400*176，8400是box个数，

176包含了3个部分，64是box坐标，80是类别prob, 32是mask coeff.

outputs[1]是proto_type, size为32 * 160 * 160.

want_float是这个结果是否输出为float，这里设0表示还是int输出，后面再反量化。

数据部分保存在buf里，它是(void*), 用的时候cast到(int8_t*).

现在需要写后处理，通过outputs得到目标框和mask.

主要思路如下：

1. 反量化

rknn模型输出的是量化过后的int8_t型，而计算时需要反量化成float型，

先看看rknn模型是怎么量化的。

  rknn_tensor_attr output_attrs[io_num.n_output];
  memset(output_attrs, 0, sizeof(output_attrs));
  for (int i = 0; i < io_num.n_output; i++) {
    output_attrs[i].index = i;
    //rknn_query:查询模型与sdk的相关信息
    //output_attrs:存放返回结果的结构体变量
    ret                   = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
    dump_tensor_attr(&(output_attrs[i]));
  }

dump结果如下，可以看到模型输出的type为int8, 量化用的是affine变换，

也就是 output / scale + zp.

那么反量化时就需要 (output - zp) * scale

//反量化计算
static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) { return ((float)qnt - (float)zp) * scale; }

2. box部分

8400 * 176, 其中8400是box个数，176包含了3个部分，64是box坐标，80是类别prob, 32是mask coeff.

8400个box部分取出每个box的80个类别prob, 选出最大的类别prob就是这个box的label,

类别prob > 阈值时说明有效，

每176的向量中前64个数据，每16个一组，一共4组，可得到box的(l,t,r,b).

用8400个anchor的中心点与(l,t,r,b)运算可得到box的(x0, y0, x1, y1).

另外把每个box的32维mask coeff保存起来。

注意判断类别prob是否 > 阈值时，要先把类别prob做反量化。

for (int i = 0; i < box_num; i++) { //遍历8400个box
    ...
    // find label with max score
	int label = -1;
	float score = -FLT_MAX;
	//找到最大score和对应的label
	for (int k = 0; k < num_class; k++)  //80个类别prob
	{
	    float confidence = deqnt_affine_to_f32(score_ptr[k], zp, scale);  //反量化,int转float
	    if (confidence > score)
	    {
	        label = k;
	        score = confidence;
	    }
	}
	float box_prob = sigmoid(score); 
	       
	if (box_prob >= prob_threshold){
        ... 
    	softmax(bbox_pred);  //对4x16中的每16个一组求softmax,因为作了softmax,就不再需要对每个数据作反量化
    	...
		//anchor中心点
		float pb_cx = (grid_strides[i].grid0 + 0.5f) * grid_strides[i].stride;
		float pb_cy = (grid_strides[i].grid1 + 0.5f) * grid_strides[i].stride;
		
		float x0 = pb_cx - pred_ltrb[0]; //center_x - l
		float y0 = pb_cy - pred_ltrb[1]; //center_y - t
		float x1 = pb_cx + pred_ltrb[2]; //center_x + r
		float y1 = pb_cy + pred_ltrb[3]; //center_y + b
		
		Object obj;
		obj.rect.x = x0;
		obj.rect.y = y0;
		obj.rect.width = x1 - x0;
		obj.rect.height = y1 - y0;
		obj.label = label;
		obj.prob = box_prob;
		obj.mask_feat.resize(32);
		
		std::copy(bbox_ptr + 64 + num_class, bbox_ptr + 64 + num_class + 32, obj.mask_feat.begin());
		objects.push_back(obj);
	}
}

上面是用类别prob的阈值过滤了第一波box,

然后还要用nms再过滤一波。

nms_sorted_bboxes(proposals, picked, nms_threshold); //picked里面保存的是proposals的下标

3. mask部分

mask由mask_coeff 和 proto_type进行矩阵乘得到。

mask_coeff保存在上面objects中每个obj的mask_feat里面。

注意上面的mask_feat还没有作反量化。

outputs[1]是数据开头的指针，

为了方便后面的矩阵乘，resize等操作，需要把它转为cv::Mat.

注意如果是int8型的Mat，它会自动把数值限制到[0,255],

而显然mask_feat里面是有负数的，所以要用float型的Mat.

先把nms挑选出来的目标的mask_coeff转为Mat。然后作反量化计算。

//mask_feat里面保存每个目标的mask coeff
cv::Mat mask_feat_qnt(count, 32, CV_32FC1); //float型的Mat, count为nms后目标的个数
for (int i = 0; i < count; i++) {
    float* mask_feat_ptr = mask_feat_qnt.ptr<float>(i);  //指向第i行的指针
    std::memcpy(mask_feat_ptr, proposals[picked[i]].mask_feat.data(), sizeof(float) * proposals[picked[i]].mask_feat.size());
}

//因为mask_feat从outputs[0]中得到，所以用zps[0],scales[0]作反量化
cv::Mat mask_feat = (mask_feat_qnt - out_zps[0]) * out_scales[0];

把proto_type转为cv::Mat. proto_type中的值也要进行反量化计算。

cv::Mat proto = cv::Mat(32, 25600, CV_32FC1, proto_arr);

mask_coeff 和 proto_type 矩阵乘，得到目标的mask.

cv::Mat mulResMask = (mask_feat * proto).t(); //(n,32)*(32,25600)=n*25600,必须加转置，不加的话mask会很奇怪
cv::Mat masks = mulResMask.reshape(count, {160, 160}); //每个目标有一个160*160的mask 

std::vector<cv::Mat> maskChannels;
cv::split(masks, maskChannels); //把count个mask分割开，每个目标的mask保存为vector的一个元素

把每个目标的box, mask保存到结果objects中。

std::vector<Object> objects;
objects.resize(count);  //保存结果的object
 for (int i = 0; i < count; i++)
 {
     objects[i] = proposals[picked[i]];

     // adjust offset to original unpadded
     float x0 = (objects[i].rect.x - (wpad / 2)) / scale;
     float y0 = (objects[i].rect.y - (hpad / 2)) / scale;
     float x1 = (objects[i].rect.x + objects[i].rect.width - (wpad / 2)) / scale;
     float y1 = (objects[i].rect.y + objects[i].rect.height - (hpad / 2)) / scale;

     // clip
     x0 = std::max(std::min(x0, (float)(img_width - 1)), 0.f);
     y0 = std::max(std::min(y0, (float)(img_height - 1)), 0.f);
     x1 = std::max(std::min(x1, (float)(img_width - 1)), 0.f);
     y1 = std::max(std::min(y1, (float)(img_height - 1)), 0.f);

     objects[i].rect.x = x0;
     objects[i].rect.y = y0;
     objects[i].rect.width = x1 - x0;
     objects[i].rect.height = y1 - y0;

     //计算mask并保存到object
     cv::Mat dest, mask;
     cv::exp(-maskChannels[i], dest); //每个mat单独计算sigmoid
     dest = 1.0 / (1.0 + dest);  //160x160的mask
     dest = dest(roi);  //取rect区域

     cv::resize(dest, mask, cv::Size(img_width, img_height), cv::INTER_NEAREST);

     //crop
     cv::Rect tmp_rect = objects[i].rect;
     mask = mask(tmp_rect) > 0.5; //mask阈值设为0.5
     objects[i].mask = mask;
 }

效果如下：

完整版代码放在github的semantic。