374 lines
9.5 KiB
C++
374 lines
9.5 KiB
C++
|
#include <iostream>
|
||
|
|
#include <fstream>
|
||
|
|
#include <vector>
|
||
|
|
#include <cstdint>
|
||
|
|
#include <stdio.h>
|
||
|
|
#include <stdlib.h>
|
||
|
|
#include <math.h>
|
||
|
|
#include <unistd.h>
|
||
|
|
#include <time.h>
|
||
|
|
#include <math.h>
|
||
|
|
#include <fcntl.h>
|
||
|
|
#include <opencv2/opencv.hpp>
|
||
|
|
#include "yolov5_detect.h"
|
||
|
|
#include "rknn_api.h"
|
||
|
|
|
||
|
|
#include <sys/time.h>
|
||
|
|
|
||
|
|
using namespace std;
|
||
|
|
using namespace cv;
|
||
|
|
|
||
|
|
|
||
|
|
//unsigned char *model;
|
||
|
|
//detection* dets;
|
||
|
|
|
||
|
|
static void printRKNNTensor(rknn_tensor_attr *attr)
|
||
|
|
{
|
||
|
|
printf("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
|
||
|
|
"fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f\n",
|
||
|
|
attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
|
||
|
|
attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
|
||
|
|
attr->qnt_type, attr->fl, attr->zp, attr->scale);
|
||
|
|
}
|
||
|
|
|
||
|
|
static int letter_box(cv::Mat input_image, cv::Mat *output_image, int model_input_size)
|
||
|
|
{
|
||
|
|
int input_width, input_height;
|
||
|
|
|
||
|
|
input_width = input_image.cols;
|
||
|
|
input_height = input_image.rows;
|
||
|
|
float ratio;
|
||
|
|
ratio = min((float)model_input_size / input_width, (float)model_input_size / input_height);
|
||
|
|
|
||
|
|
int new_width, new_height;
|
||
|
|
new_width = round(ratio * input_width );
|
||
|
|
new_height = round(ratio * input_height);
|
||
|
|
|
||
|
|
|
||
|
|
int height_padding = 0;
|
||
|
|
int width_padding = 0;
|
||
|
|
int top = 0;
|
||
|
|
int bottom = 0;
|
||
|
|
int left = 0;
|
||
|
|
int right = 0;
|
||
|
|
if( new_width >= new_height)
|
||
|
|
{
|
||
|
|
height_padding = new_width - new_height;
|
||
|
|
if( (height_padding % 2) == 0 )
|
||
|
|
{
|
||
|
|
top = (int)((float)(height_padding/2));
|
||
|
|
bottom = (int)((float)(height_padding/2));
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
top = (int)((float)(height_padding/2));
|
||
|
|
bottom = (int)((float)(height_padding/2))+1;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
width_padding = new_height - new_width;
|
||
|
|
if( (width_padding % 2) == 0 )
|
||
|
|
{
|
||
|
|
left = (int)((float)(width_padding/2));
|
||
|
|
right = (int)((float)(width_padding/2));
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
left = (int)((float)(width_padding/2));
|
||
|
|
right = (int)((float)(width_padding/2))+1;
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat resize_img;
|
||
|
|
|
||
|
|
cv::resize(input_image, resize_img, cv::Size(new_width, new_height));
|
||
|
|
cv::copyMakeBorder(resize_img, *output_image, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
int yolov5_detect_init(rknn_context *ctx, const char * path)
|
||
|
|
{
|
||
|
|
int ret;
|
||
|
|
|
||
|
|
// Load model
|
||
|
|
FILE *fp = fopen(path, "rb");
|
||
|
|
if(fp == NULL)
|
||
|
|
{
|
||
|
|
printf("fopen %s fail!\n", path);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
fseek(fp, 0, SEEK_END); //fp指向end,fseek(FILE *stream, long offset, int fromwhere);
|
||
|
|
int model_len = ftell(fp); //相对文件首偏移
|
||
|
|
unsigned char *model_data = (unsigned char*)malloc(model_len);
|
||
|
|
|
||
|
|
fseek(fp, 0, SEEK_SET); //SEEK_SET为文件头
|
||
|
|
if(model_len != fread(model_data, 1, model_len, fp))
|
||
|
|
{
|
||
|
|
printf("fread %s fail!\n", path);
|
||
|
|
free(model_data);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
fclose(fp);
|
||
|
|
|
||
|
|
//init
|
||
|
|
ret = rknn_init(ctx, model_data, model_len, RKNN_FLAG_PRIOR_MEDIUM);
|
||
|
|
if(ret < 0)
|
||
|
|
{
|
||
|
|
printf("rknn_init fail! ret=%d\n", ret);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
|
||
|
|
free(model_data);
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
static int scale_coords(yolov5_detect_result_group_t *detect_result_group, int img_width, int img_height, int model_size)
|
||
|
|
{
|
||
|
|
for (int i = 0; i < detect_result_group->count; i++)
|
||
|
|
{
|
||
|
|
yolov5_detect_result_t *det_result = &(detect_result_group->results[i]);
|
||
|
|
|
||
|
|
|
||
|
|
int x1 = det_result->box.left;
|
||
|
|
int y1 = det_result->box.top;
|
||
|
|
int x2 = det_result->box.right;
|
||
|
|
int y2 = det_result->box.bottom;
|
||
|
|
|
||
|
|
|
||
|
|
if( img_width >= img_height )
|
||
|
|
{
|
||
|
|
int image_max_len = img_width;
|
||
|
|
float gain;
|
||
|
|
gain = (float)model_size / image_max_len;
|
||
|
|
int resized_height = img_height * gain;
|
||
|
|
int height_pading = (model_size - resized_height)/2;
|
||
|
|
y1 = (y1 - height_pading);
|
||
|
|
y2 = (y2 - height_pading);
|
||
|
|
x1 = int(x1 / gain);
|
||
|
|
y1 = int(y1 / gain);
|
||
|
|
x2 = int(x2 / gain);
|
||
|
|
y2 = int(y2 / gain);
|
||
|
|
|
||
|
|
det_result->box.left = x1;
|
||
|
|
det_result->box.top = y1;
|
||
|
|
det_result->box.right = x2;
|
||
|
|
det_result->box.bottom = y2;
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
int image_max_len = img_height;
|
||
|
|
float gain;
|
||
|
|
gain = (float)model_size / image_max_len;
|
||
|
|
int resized_width = img_width * gain;
|
||
|
|
int width_pading = (model_size - resized_width)/2;
|
||
|
|
x1 = (x1 - width_pading);
|
||
|
|
x2 = (x2 - width_pading);
|
||
|
|
x1 = int(x1 / gain);
|
||
|
|
y1 = int(y1 / gain);
|
||
|
|
x2 = int(x2 / gain);
|
||
|
|
y2 = int(y2 / gain);
|
||
|
|
|
||
|
|
det_result->box.left = x1;
|
||
|
|
det_result->box.top = y1;
|
||
|
|
det_result->box.right = x2;
|
||
|
|
det_result->box.bottom = y2;
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
int yolov5_detect_run(rknn_context ctx, cv::Mat input_image, yolov5_detect_result_group_t *detect_result_group)
|
||
|
|
{
|
||
|
|
int img_width = 0;
|
||
|
|
int img_height = 0;
|
||
|
|
int img_channel = 0;
|
||
|
|
|
||
|
|
size_t actual_size = 0;
|
||
|
|
const float vis_threshold = 0.1;
|
||
|
|
const float nms_threshold = 0.5;
|
||
|
|
const float conf_threshold = 0.2;
|
||
|
|
int ret;
|
||
|
|
|
||
|
|
img_width = input_image.cols;
|
||
|
|
img_height = input_image.rows;
|
||
|
|
|
||
|
|
rknn_sdk_version version;
|
||
|
|
ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version,
|
||
|
|
sizeof(rknn_sdk_version));
|
||
|
|
if (ret < 0)
|
||
|
|
{
|
||
|
|
printf("rknn_init error ret=%d\n", ret);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
/*
|
||
|
|
printf("sdk version: %s driver version: %s\n", version.api_version,
|
||
|
|
version.drv_version);
|
||
|
|
*/
|
||
|
|
|
||
|
|
rknn_input_output_num io_num;
|
||
|
|
ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
|
||
|
|
if (ret < 0)
|
||
|
|
{
|
||
|
|
printf("rknn_init error ret=%d\n", ret);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
/*
|
||
|
|
printf("model input num: %d, output num: %d\n", io_num.n_input,
|
||
|
|
io_num.n_output);
|
||
|
|
*/
|
||
|
|
|
||
|
|
rknn_tensor_attr input_attrs[io_num.n_input];
|
||
|
|
memset(input_attrs, 0, sizeof(input_attrs));
|
||
|
|
for (int i = 0; i < io_num.n_input; i++)
|
||
|
|
{
|
||
|
|
input_attrs[i].index = i;
|
||
|
|
ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]),
|
||
|
|
sizeof(rknn_tensor_attr));
|
||
|
|
if (ret < 0)
|
||
|
|
{
|
||
|
|
printf("rknn_init error ret=%d\n", ret);
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
//printRKNNTensor(&(input_attrs[i]));
|
||
|
|
}
|
||
|
|
|
||
|
|
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;
|
||
|
|
ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
|
||
|
|
sizeof(rknn_tensor_attr));
|
||
|
|
//printRKNNTensor(&(output_attrs[i]));
|
||
|
|
}
|
||
|
|
|
||
|
|
int input_channel = 3;
|
||
|
|
int input_width = 0;
|
||
|
|
int input_height = 0;
|
||
|
|
if (input_attrs[0].fmt == RKNN_TENSOR_NCHW)
|
||
|
|
{
|
||
|
|
//printf("model is NCHW input fmt\n");
|
||
|
|
input_width = input_attrs[0].dims[0];
|
||
|
|
input_height = input_attrs[0].dims[1];
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
//printf("model is NHWC input fmt\n");
|
||
|
|
input_width = input_attrs[0].dims[1];
|
||
|
|
input_height = input_attrs[0].dims[2];
|
||
|
|
}
|
||
|
|
|
||
|
|
/*
|
||
|
|
printf("model input height=%d, width=%d, channel=%d\n", height, width,
|
||
|
|
channel);
|
||
|
|
*/
|
||
|
|
|
||
|
|
/* Init input tensor */
|
||
|
|
rknn_input inputs[1];
|
||
|
|
memset(inputs, 0, sizeof(inputs));
|
||
|
|
inputs[0].index = 0;
|
||
|
|
inputs[0].type = RKNN_TENSOR_UINT8;
|
||
|
|
inputs[0].size = input_width * input_height * input_channel;
|
||
|
|
inputs[0].fmt = RKNN_TENSOR_NHWC;
|
||
|
|
inputs[0].pass_through = 0;
|
||
|
|
|
||
|
|
/* Init output tensor */
|
||
|
|
rknn_output outputs[io_num.n_output];
|
||
|
|
memset(outputs, 0, sizeof(outputs));
|
||
|
|
|
||
|
|
for (int i = 0; i < io_num.n_output; i++)
|
||
|
|
{
|
||
|
|
outputs[i].want_float = 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat letter_image;
|
||
|
|
letter_box(input_image, &letter_image, input_width);
|
||
|
|
inputs[0].buf = letter_image.data;
|
||
|
|
|
||
|
|
rknn_inputs_set(ctx, io_num.n_input, inputs);
|
||
|
|
ret = rknn_run(ctx, NULL);
|
||
|
|
ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
|
||
|
|
|
||
|
|
// Post process
|
||
|
|
std::vector<float> out_scales;
|
||
|
|
std::vector<uint8_t> out_zps;
|
||
|
|
for (int i = 0; i < io_num.n_output; ++i)
|
||
|
|
{
|
||
|
|
out_scales.push_back(output_attrs[i].scale);
|
||
|
|
out_zps.push_back(output_attrs[i].zp);
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
yolov5_post_process_u8((uint8_t *)outputs[0].buf, (uint8_t *)outputs[1].buf, (uint8_t *)outputs[2].buf, input_height, input_width,
|
||
|
|
conf_threshold, nms_threshold, out_zps, out_scales, detect_result_group);
|
||
|
|
|
||
|
|
|
||
|
|
/*
|
||
|
|
yolov5_post_process_fp((float *)outputs[0].buf, (float *)outputs[1].buf, (float *)outputs[2].buf, input_height, input_width,
|
||
|
|
conf_threshold, nms_threshold, &detect_result_group);
|
||
|
|
*/
|
||
|
|
|
||
|
|
rknn_outputs_release(ctx, io_num.n_output, outputs);
|
||
|
|
|
||
|
|
scale_coords(detect_result_group, img_width, img_height, input_width);
|
||
|
|
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
int yolov5_detect_release(rknn_context ctx)
|
||
|
|
{
|
||
|
|
rknn_destroy(ctx);
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
std::string base64_encode(unsigned char const* bytes_to_encode, unsigned int in_len) {
|
||
|
|
std::string ret;
|
||
|
|
int i = 0;
|
||
|
|
int j = 0;
|
||
|
|
unsigned char char_array_3[3];
|
||
|
|
unsigned char char_array_4[4];
|
||
|
|
|
||
|
|
while (in_len--) {
|
||
|
|
char_array_3[i++] = *(bytes_to_encode++);
|
||
|
|
if (i == 3) {
|
||
|
|
char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
|
||
|
|
char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
|
||
|
|
char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
|
||
|
|
char_array_4[3] = char_array_3[2] & 0x3f;
|
||
|
|
|
||
|
|
for(i = 0; (i <4) ; i++)
|
||
|
|
ret += base64_chars[char_array_4[i]];
|
||
|
|
i = 0;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if (i) {
|
||
|
|
for(j = i; j < 3; j++)
|
||
|
|
char_array_3[j] = '\0';
|
||
|
|
|
||
|
|
char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
|
||
|
|
char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
|
||
|
|
char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
|
||
|
|
char_array_4[3] = char_array_3[2] & 0x3f;
|
||
|
|
|
||
|
|
for (j = 0; (j < i + 1); j++)
|
||
|
|
ret += base64_chars[char_array_4[j]];
|
||
|
|
|
||
|
|
while((i++ < 3))
|
||
|
|
ret += '=';
|
||
|
|
}
|
||
|
|
|
||
|
|
return ret;
|
||
|
|
}
|
||
|
|
|
||
|
|
static inline bool is_base64(unsigned char c) {
|
||
|
|
return (isalnum(c) || (c == '+') || (c == '/'));
|
||
|
|
}
|