126 lines
4.1 KiB
C++
126 lines
4.1 KiB
C++
// Copyright (c) 2023 by Rockchip Electronics Co., Ltd. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
/*-------------------------------------------
|
|
Includes
|
|
-------------------------------------------*/
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#include "mobilenet.h"
|
|
#include "image_utils.h"
|
|
#include "file_utils.h"
|
|
|
|
#if defined(RV1106_1103)
|
|
#include "dma_alloc.cpp"
|
|
#endif
|
|
|
|
#define IMAGENET_CLASSES_FILE "./model/synset.txt"
|
|
|
|
/*-------------------------------------------
|
|
Main Function
|
|
-------------------------------------------*/
|
|
int main(int argc, char** argv)
|
|
{
|
|
if (argc != 3) {
|
|
printf("%s <model_path> <image_path>\n", argv[0]);
|
|
return -1;
|
|
}
|
|
|
|
const char* model_path = argv[1];
|
|
const char* image_path = argv[2];
|
|
|
|
int line_count;
|
|
const char *configured_labels = getenv("YOLO_LABELS_PATH");
|
|
const char *labels_path = configured_labels && configured_labels[0] ? configured_labels : IMAGENET_CLASSES_FILE;
|
|
char** lines = read_lines_from_file(labels_path, &line_count);
|
|
if (lines == NULL) {
|
|
line_count = 0;
|
|
printf("labels file unavailable; use class_<id> fallback: %s\n", labels_path);
|
|
}
|
|
|
|
int ret;
|
|
rknn_app_context_t rknn_app_ctx;
|
|
memset(&rknn_app_ctx, 0, sizeof(rknn_app_context_t));
|
|
|
|
ret = init_mobilenet_model(model_path, &rknn_app_ctx);
|
|
if (ret != 0) {
|
|
printf("init_mobilenet_model fail! ret=%d model_path=%s\n", ret, model_path);
|
|
return -1;
|
|
}
|
|
|
|
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
|
|
|
|
int topk = 5;
|
|
mobilenet_result result[topk];
|
|
|
|
ret = inference_mobilenet_model(&rknn_app_ctx, &src_image, result, topk);
|
|
if (ret != 0) {
|
|
printf("init_mobilenet_model fail! ret=%d\n", ret);
|
|
goto out;
|
|
}
|
|
|
|
for (int i = 0; i < topk; i++) {
|
|
if (result[i].cls < 0) continue;
|
|
const bool has_label = lines != NULL && result[i].cls >= 0 && result[i].cls < line_count;
|
|
if (has_label) {
|
|
printf("[%d] score=%.6f class=%s\n", result[i].cls, result[i].score, lines[result[i].cls]);
|
|
} else {
|
|
printf("[%d] score=%.6f class=class_%d\n", result[i].cls, result[i].score, result[i].cls);
|
|
}
|
|
}
|
|
write_image("out.png", &src_image);
|
|
|
|
out:
|
|
ret = release_mobilenet_model(&rknn_app_ctx);
|
|
if (ret != 0) {
|
|
printf("release_mobilenet_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
|
|
}
|
|
if (lines != NULL) {
|
|
free_lines(lines, line_count);
|
|
}
|
|
|
|
return 0;
|
|
}
|