Luckfox Forums

Code: Select all

/*-------------------------------------------
                Includes
-------------------------------------------*/
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "yolov5.h"
#include "image_utils.h"
#include "file_utils.h"

#if defined(RV1106_1103)
    #include "dma_alloc.cpp"
#endif

void print_output_results(output_result *results, int num_results) {
    printf("Predictions:\n");
    for (int i = 0; i < num_results; i++) {
        printf("Class: %d, Score: %.4f\n", results[i].cls, results[i].score);
    }
    // 打开文件
    FILE *file = fopen("/tmp/output", "w");
    if (file == NULL) {
        printf("无法打开文件。\n");
    }

    // 将整数写入文件
    fprintf(file, "%d", results[0].cls);

    // 关闭文件
    fclose(file);

    printf("The prediction of %d has been save into /tmp/output.\n", results[0].cls);
}

/*-------------------------------------------
                  Main Function
-------------------------------------------*/
int main(int argc, char **argv)
{
    if (argc != 3 && argc != 4)
    {
        printf("%s <model_path> <image_path>\n", argv[0]);
        return -1;
    }

    const char *model_path = argv[1];
    const char *image_path = argv[2];

    int ret;
    rknn_app_context_t rknn_app_ctx;
    memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t));

    int topk = 2;
    output_result result[topk];

    ret = init_yolov5_model(model_path, &rknn_app_ctx);
    if (ret != 0)
    {
        printf("init_yolov5_model fail! ret=%d model_path=%s\n", ret, model_path);
        goto out;
    }

    image_buffer_t src_image;
    memset(&src_image, 0, sizeof(image_buffer_t));
    ret = read_image(image_path, &src_image);
    if (ret != 0) {
        printf("read image fail! ret=%d image_path=%s\n", ret, image_path);
        return -1;
    }

#if defined(RV1106_1103)
    //RV1106 rga requires that input and output bufs are memory allocated by dma
    ret = dma_buf_alloc(RV1106_CMA_HEAP_PATH, src_image.size, &rknn_app_ctx.img_dma_buf.dma_buf_fd, 
                       (void **) & (rknn_app_ctx.img_dma_buf.dma_buf_virt_addr));
    memcpy(rknn_app_ctx.img_dma_buf.dma_buf_virt_addr, src_image.virt_addr, src_image.size);
    dma_sync_cpu_to_device(rknn_app_ctx.img_dma_buf.dma_buf_fd);
    free(src_image.virt_addr);
    src_image.virt_addr = (unsigned char *)rknn_app_ctx.img_dma_buf.dma_buf_virt_addr;
    src_image.fd = rknn_app_ctx.img_dma_buf.dma_buf_fd;
    rknn_app_ctx.img_dma_buf.size = src_image.size;
#endif

    printf(">>> going to inference_yolov5_model\n");

    ret = inference_yolov5_model(&rknn_app_ctx, &src_image, result, topk);
    if (ret != 0) {
        printf("inference_yolov5_model fail! ret=%d\n", ret);
        goto out;
    }

    print_output_results(result, topk);

out:
    printf(">>> going to release_yolov5_model\n");
    ret = release_yolov5_model(&rknn_app_ctx);
    if (ret != 0)
    {
        printf("release_yolov5_model fail! ret=%d\n", ret);
    }

    if (src_image.virt_addr != NULL) {
#if defined(RV1106_1103)
        dma_buf_free(rknn_app_ctx.img_dma_buf.size, &rknn_app_ctx.img_dma_buf.dma_buf_fd,
                rknn_app_ctx.img_dma_buf.dma_buf_virt_addr);
#else
        free(src_image.virt_addr);
#endif
    }

    return 0;
}

Code: Select all

int inference_yolov5_model(rknn_app_context_t *app_ctx, image_buffer_t *src_img, output_result *out_result, int topk)
{
    int ret;
    image_buffer_t img;
    memset(&img, 0, sizeof(image_buffer_t));

    // Pre Process
    img.width = app_ctx->model_width;
    img.height = app_ctx->model_height;
    img.format = IMAGE_FORMAT_RGB888;
    img.size = get_image_size(&img);
    img.virt_addr = (unsigned char *)app_ctx->input_mems[0]->virt_addr;
    if (img.virt_addr == NULL) {
        printf("malloc buffer size:%d fail!\n", img.size);
        return -1;
    }

    ret = convert_image(src_img, &img, NULL, NULL, 0);
//    ret = convert_image_gray(src_img, &img, NULL, NULL, 0);
    if (ret < 0) {
        printf("convert_image fail! ret=%d\n", ret);
        return -1;
    }

    // Run
    printf("rknn_run...\n");
    ret = rknn_run(app_ctx->rknn_ctx, nullptr);
    if (ret < 0) {
        printf("rknn_run fail! ret=%d\n", ret);
        return -1;
    }

    float outputs_float[app_ctx->output_attrs[0].size_with_stride];
    if (app_ctx->output_attrs[0].fmt == RKNN_TENSOR_NC1HWC2) {
        printf("RKNN_TENSOR_NC1HWC2\n");
        int size = app_ctx->output_attrs[0].size_with_stride;
        int zp = app_ctx->output_attrs[0].zp;
        float scale = app_ctx->output_attrs[0].scale;
        int channel = app_ctx->output_attrs[0].dims[1];
        int h = app_ctx->output_attrs[0].n_dims > 2 ? app_ctx->output_attrs[0].dims[2] : 1;
        int w = app_ctx->output_attrs[0].n_dims > 3 ? app_ctx->output_attrs[0].dims[3] : 1;
        if (app_ctx->output_attrs[0].type == RKNN_TENSOR_INT8) {
            NC1HWC2_int8_to_NCHW_float((int8_t *)app_ctx->output_mems[0]->virt_addr, outputs_float,
                                       (int *)app_ctx->output_attrs[0].dims, channel, h, w, zp, scale);
        } else {
            printf("dtype: %s cannot convert!", get_type_string(app_ctx->output_attrs[0].type));
        }
    } else {
        if (app_ctx->output_attrs[0].type == RKNN_TENSOR_INT8) {
            printf("RKNN_TENSOR_INT8\n");
            int8_t *src = (int8_t *)app_ctx->output_mems[0]->virt_addr;
            for (int index = 0; index < app_ctx->output_attrs[0].n_elems; index++) {
                outputs_float[index] = ((float)src[index] - app_ctx->output_attrs[0].zp) * app_ctx->output_attrs[0].scale;
            }
            // 输出转换前的数组
            printf("Original array:\n");
            for (int index = 0; index < app_ctx->output_attrs[0].n_elems; index++) {
                printf("%d ", src[index]);
            }
            printf("\n");
        } else if(app_ctx->output_attrs[0].type == RKNN_TENSOR_UINT8) {
            printf("RKNN_TENSOR_UINT8\n");
            uint8_t *src = (uint8_t *)app_ctx->output_mems[0]->virt_addr;
            for (int index = 0; index < app_ctx->output_attrs[0].n_elems; index++) {
                outputs_float[index] = ((float)src[index] - app_ctx->output_attrs[0].zp) * app_ctx->output_attrs[0].scale;
            }
            // 输出转换前的数组
            printf("Original array:\n");
            for (int index = 0; index < app_ctx->output_attrs[0].n_elems; index++) {
                printf("%d ", src[index]);
            }
            printf("\n");
        }
    }

    // Post Process
    printf("softmax with num_classes %d...\n", app_ctx->output_attrs[0].n_elems);
    softmax(outputs_float, app_ctx->output_attrs[0].n_elems);

    printf("get_topk_with_indices...\n");
    get_topk_with_indices(outputs_float, app_ctx->output_attrs[0].n_elems, topk, out_result);
    printf("inference done...\n");

out:
    return ret;
}