283 lines
12 KiB
C++
283 lines
12 KiB
C++
|
#include "yolo_runtime_backend.h"
|
||
|
|
#include "yolo_mask_utils.h"
|
||
|
|
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
#include "yolov8_seg.h"
|
||
|
|
#elif defined(YOLO_BACKEND_CLS)
|
||
|
|
#include "mobilenet.h"
|
||
|
|
#elif defined(YOLO_BACKEND_POSE)
|
||
|
|
#include "yolov8-pose.h"
|
||
|
|
#elif defined(YOLO_BACKEND_OBB)
|
||
|
|
#include "yolov8-obb.h"
|
||
|
|
#else
|
||
|
|
#include "yolov8.h"
|
||
|
|
#endif
|
||
|
|
|
||
|
|
#include <algorithm>
|
||
|
|
#include <chrono>
|
||
|
|
#include <cstdio>
|
||
|
|
#include <cstdlib>
|
||
|
|
#include <cstring>
|
||
|
|
#include <mutex>
|
||
|
|
#include <new>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
extern "C" void yolo_adapter_set_thresholds(float confidence, float nms);
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
extern "C" void yolo_adapter_set_class_agnostic_nms(int enabled);
|
||
|
|
#endif
|
||
|
|
extern "C" void yolo_adapter_get_performance(
|
||
|
|
double* preprocess_ms,
|
||
|
|
double* inference_ms,
|
||
|
|
double* postprocess_ms);
|
||
|
|
|
||
|
|
struct yolo_runtime_backend {
|
||
|
|
rknn_app_context_t context{};
|
||
|
|
std::mutex context_mutex;
|
||
|
|
bool initialized = false;
|
||
|
|
float confidence_threshold = 0.25f;
|
||
|
|
float nms_threshold = 0.45f;
|
||
|
|
bool class_agnostic_nms = false;
|
||
|
|
std::vector<std::vector<unsigned char>> segmentation_masks;
|
||
|
|
};
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
|
||
|
|
int packed_row_bytes(const yolo_image_view_v1_t& image) {
|
||
|
|
switch (image.pixel_format) {
|
||
|
|
case YOLO_PIXEL_FORMAT_GRAY8:
|
||
|
|
case YOLO_PIXEL_FORMAT_NV12:
|
||
|
|
case YOLO_PIXEL_FORMAT_NV21: return image.width;
|
||
|
|
case YOLO_PIXEL_FORMAT_RGB888:
|
||
|
|
case YOLO_PIXEL_FORMAT_BGR888: return image.width * 3;
|
||
|
|
case YOLO_PIXEL_FORMAT_RGBA8888:
|
||
|
|
case YOLO_PIXEL_FORMAT_BGRA8888: return image.width * 4;
|
||
|
|
default: return 0;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
bool prepare_image(const yolo_image_view_v1_t& source, image_buffer_t& destination, std::vector<unsigned char>& owned) {
|
||
|
|
const int packed_stride = packed_row_bytes(source);
|
||
|
|
if (packed_stride <= 0) return false;
|
||
|
|
const int source_stride = source.row_stride_bytes > 0 ? source.row_stride_bytes : packed_stride;
|
||
|
|
if (source_stride < packed_stride) return false;
|
||
|
|
int rows = source.height;
|
||
|
|
if (source.pixel_format == YOLO_PIXEL_FORMAT_NV12 || source.pixel_format == YOLO_PIXEL_FORMAT_NV21) rows = source.height * 3 / 2;
|
||
|
|
const bool swap_red_blue = source.pixel_format == YOLO_PIXEL_FORMAT_BGR888 || source.pixel_format == YOLO_PIXEL_FORMAT_BGRA8888;
|
||
|
|
const bool copy_rows = source_stride != packed_stride;
|
||
|
|
const unsigned char* pixels = source.data;
|
||
|
|
if (swap_red_blue || copy_rows) {
|
||
|
|
owned.resize(static_cast<size_t>(packed_stride) * rows);
|
||
|
|
for (int row = 0; row < rows; ++row) {
|
||
|
|
const unsigned char* input = source.data + static_cast<size_t>(row) * source_stride;
|
||
|
|
unsigned char* output = owned.data() + static_cast<size_t>(row) * packed_stride;
|
||
|
|
std::memcpy(output, input, static_cast<size_t>(packed_stride));
|
||
|
|
if (swap_red_blue && row < source.height) {
|
||
|
|
const int channels = source.pixel_format == YOLO_PIXEL_FORMAT_BGR888 ? 3 : 4;
|
||
|
|
for (int column = 0; column < source.width; ++column) std::swap(output[column * channels], output[column * channels + 2]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
pixels = owned.data();
|
||
|
|
}
|
||
|
|
destination.width = source.width;
|
||
|
|
destination.height = source.height;
|
||
|
|
destination.width_stride = source.width;
|
||
|
|
destination.height_stride = source.height;
|
||
|
|
destination.virt_addr = const_cast<unsigned char*>(pixels);
|
||
|
|
destination.size = packed_stride * rows;
|
||
|
|
destination.fd = 0;
|
||
|
|
switch (source.pixel_format) {
|
||
|
|
case YOLO_PIXEL_FORMAT_GRAY8: destination.format = IMAGE_FORMAT_GRAY8; break;
|
||
|
|
case YOLO_PIXEL_FORMAT_RGB888:
|
||
|
|
case YOLO_PIXEL_FORMAT_BGR888: destination.format = IMAGE_FORMAT_RGB888; break;
|
||
|
|
case YOLO_PIXEL_FORMAT_RGBA8888:
|
||
|
|
case YOLO_PIXEL_FORMAT_BGRA8888: destination.format = IMAGE_FORMAT_RGBA8888; break;
|
||
|
|
case YOLO_PIXEL_FORMAT_NV12: destination.format = IMAGE_FORMAT_YUV420SP_NV12; break;
|
||
|
|
case YOLO_PIXEL_FORMAT_NV21: destination.format = IMAGE_FORMAT_YUV420SP_NV21; break;
|
||
|
|
default: return false;
|
||
|
|
}
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
bool valid_image(const yolo_image_view_v1_t* image) {
|
||
|
|
return image && image->struct_size >= sizeof(yolo_image_view_v1_t) && image->version == YOLO_IMAGE_VIEW_ABI_VERSION && image->data && image->width > 0 && image->height > 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
#if defined(YOLO_BACKEND_OBB)
|
||
|
|
void set_box(yolo_backend_result_t& target, const image_obb_box_t& box) {
|
||
|
|
target.bbox.x = static_cast<float>(box.x);
|
||
|
|
target.bbox.y = static_cast<float>(box.y);
|
||
|
|
target.bbox.width = static_cast<float>(box.w);
|
||
|
|
target.bbox.height = static_cast<float>(box.h);
|
||
|
|
target.bbox.angle = box.angle;
|
||
|
|
target.bbox.has_angle = 1;
|
||
|
|
}
|
||
|
|
#else
|
||
|
|
void set_box(yolo_backend_result_t& target, const image_rect_t& box) {
|
||
|
|
target.bbox.x = static_cast<float>(box.left);
|
||
|
|
target.bbox.y = static_cast<float>(box.top);
|
||
|
|
target.bbox.width = static_cast<float>(box.right - box.left);
|
||
|
|
target.bbox.height = static_cast<float>(box.bottom - box.top);
|
||
|
|
}
|
||
|
|
#endif
|
||
|
|
|
||
|
|
} // namespace
|
||
|
|
|
||
|
|
extern "C" int yolo_backend_get_abi_version(void) { return YOLO_RUNTIME_BACKEND_ABI_VERSION; }
|
||
|
|
|
||
|
|
extern "C" int yolo_backend_create(
|
||
|
|
const char* model_path,
|
||
|
|
const char* config_path,
|
||
|
|
float confidence_threshold,
|
||
|
|
float nms_threshold,
|
||
|
|
int class_agnostic_nms,
|
||
|
|
yolo_runtime_backend_t** output) {
|
||
|
|
if (!model_path || !output) return YOLO_STATUS_INVALID_ARGUMENT;
|
||
|
|
(void)config_path;
|
||
|
|
auto* backend = new (std::nothrow) yolo_runtime_backend();
|
||
|
|
if (!backend) return YOLO_STATUS_INITIALIZATION_FAILED;
|
||
|
|
backend->confidence_threshold = confidence_threshold;
|
||
|
|
backend->nms_threshold = nms_threshold;
|
||
|
|
backend->class_agnostic_nms = class_agnostic_nms != 0;
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
const int status = init_yolov8_seg_model(model_path, &backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_CLS)
|
||
|
|
const int status = init_mobilenet_model(model_path, &backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_POSE)
|
||
|
|
const int status = init_yolov8_pose_model(model_path, &backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_OBB)
|
||
|
|
const int status = init_yolov8_obb_model(model_path, &backend->context);
|
||
|
|
#else
|
||
|
|
const int status = init_yolov8_model(model_path, &backend->context);
|
||
|
|
#endif
|
||
|
|
if (status != 0) { delete backend; return YOLO_STATUS_INITIALIZATION_FAILED; }
|
||
|
|
backend->initialized = true;
|
||
|
|
*output = backend;
|
||
|
|
return YOLO_STATUS_OK;
|
||
|
|
}
|
||
|
|
|
||
|
|
extern "C" int yolo_backend_infer(yolo_runtime_backend_t* backend, const yolo_image_view_v1_t* image, yolo_backend_result_list_t* results) {
|
||
|
|
if (!backend || !backend->initialized || !valid_image(image) || !results) return YOLO_STATUS_INVALID_ARGUMENT;
|
||
|
|
std::lock_guard<std::mutex> lock(backend->context_mutex);
|
||
|
|
std::memset(results, 0, sizeof(*results));
|
||
|
|
image_buffer_t input{};
|
||
|
|
std::vector<unsigned char> owned_pixels;
|
||
|
|
if (!prepare_image(*image, input, owned_pixels)) return YOLO_STATUS_INVALID_ARGUMENT;
|
||
|
|
const auto start = std::chrono::steady_clock::now();
|
||
|
|
#if defined(YOLO_BACKEND_CLS)
|
||
|
|
mobilenet_result classifications[5]{};
|
||
|
|
const int status = inference_mobilenet_model(&backend->context, &input, classifications, 5);
|
||
|
|
results->task = YOLO_TASK_CLS;
|
||
|
|
if (status == 0) {
|
||
|
|
for (const auto& classification : classifications) {
|
||
|
|
if (classification.cls < 0 || results->count >= YOLO_BACKEND_MAX_RESULTS) continue;
|
||
|
|
auto& result = results->results[results->count++];
|
||
|
|
result.task = YOLO_TASK_CLS;
|
||
|
|
result.class_id = classification.cls;
|
||
|
|
result.score = classification.score;
|
||
|
|
std::snprintf(result.label, sizeof(result.label), "class_%d", result.class_id);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
#else
|
||
|
|
object_detect_result_list detections{};
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
|
||
|
|
yolo_adapter_set_class_agnostic_nms(backend->class_agnostic_nms ? 1 : 0);
|
||
|
|
const int status = inference_yolov8_seg_model(&backend->context, &input, &detections);
|
||
|
|
results->task = YOLO_TASK_SEG;
|
||
|
|
#elif defined(YOLO_BACKEND_POSE)
|
||
|
|
yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
|
||
|
|
const int status = inference_yolov8_pose_model(&backend->context, &input, &detections);
|
||
|
|
results->task = YOLO_TASK_POSE;
|
||
|
|
#elif defined(YOLO_BACKEND_OBB)
|
||
|
|
yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
|
||
|
|
const int status = inference_yolov8_obb_model(&backend->context, &input, &detections);
|
||
|
|
results->task = YOLO_TASK_OBB;
|
||
|
|
#else
|
||
|
|
yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
|
||
|
|
const int status = inference_yolov8_model(&backend->context, &input, &detections);
|
||
|
|
results->task = YOLO_TASK_DETECT;
|
||
|
|
#endif
|
||
|
|
if (status == 0) {
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
backend->segmentation_masks.clear();
|
||
|
|
backend->segmentation_masks.reserve(static_cast<size_t>(std::min(detections.count, YOLO_BACKEND_MAX_RESULTS)));
|
||
|
|
#else
|
||
|
|
results->count = std::min(detections.count, YOLO_BACKEND_MAX_RESULTS);
|
||
|
|
#endif
|
||
|
|
const int source_count = std::min(detections.count, YOLO_BACKEND_MAX_RESULTS);
|
||
|
|
for (int index = 0; index < source_count; ++index) {
|
||
|
|
const auto& source = detections.results[index];
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
unsigned char* full_mask = detections.results_seg[index].seg_mask;
|
||
|
|
const auto bounds = yolo_find_mask_bounds(full_mask, image->width, image->height, image->width);
|
||
|
|
auto mask_roi = yolo_copy_binary_mask_roi(full_mask, image->width, bounds);
|
||
|
|
std::free(full_mask);
|
||
|
|
detections.results_seg[index].seg_mask = nullptr;
|
||
|
|
if (bounds.empty() || mask_roi.empty()) continue;
|
||
|
|
backend->segmentation_masks.push_back(std::move(mask_roi));
|
||
|
|
auto& result = results->results[results->count++];
|
||
|
|
#else
|
||
|
|
auto& result = results->results[index];
|
||
|
|
#endif
|
||
|
|
result.task = results->task;
|
||
|
|
result.class_id = source.cls_id;
|
||
|
|
result.score = source.prop;
|
||
|
|
std::snprintf(result.label, sizeof(result.label), "class_%d", source.cls_id);
|
||
|
|
set_box(result, source.box);
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
auto& mask = backend->segmentation_masks.back();
|
||
|
|
result.mask_bbox.x = static_cast<float>(bounds.x);
|
||
|
|
result.mask_bbox.y = static_cast<float>(bounds.y);
|
||
|
|
result.mask_bbox.width = static_cast<float>(bounds.width);
|
||
|
|
result.mask_bbox.height = static_cast<float>(bounds.height);
|
||
|
|
result.mask_data = mask.data();
|
||
|
|
result.mask_width = bounds.width;
|
||
|
|
result.mask_height = bounds.height;
|
||
|
|
result.mask_stride_bytes = bounds.width;
|
||
|
|
result.mask_origin_x = bounds.x;
|
||
|
|
result.mask_origin_y = bounds.y;
|
||
|
|
#elif defined(YOLO_BACKEND_POSE)
|
||
|
|
result.keypoint_count = 17;
|
||
|
|
for (int keypoint = 0; keypoint < result.keypoint_count; ++keypoint) {
|
||
|
|
result.keypoints[keypoint] = { source.keypoints[keypoint][0], source.keypoints[keypoint][1], source.keypoints[keypoint][2] };
|
||
|
|
}
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
for (int index = source_count; index < detections.count; ++index) {
|
||
|
|
std::free(detections.results_seg[index].seg_mask);
|
||
|
|
detections.results_seg[index].seg_mask = nullptr;
|
||
|
|
}
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
#endif
|
||
|
|
const double elapsed = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
|
||
|
|
results->performance.total_ms = elapsed;
|
||
|
|
yolo_adapter_get_performance(&results->performance.preprocess_ms, &results->performance.inference_ms, &results->performance.postprocess_ms);
|
||
|
|
return status == 0 ? YOLO_STATUS_OK : YOLO_STATUS_BACKEND_FAILED;
|
||
|
|
}
|
||
|
|
|
||
|
|
extern "C" const char* yolo_backend_last_error(const yolo_runtime_backend_t* backend) {
|
||
|
|
(void)backend;
|
||
|
|
return "RKNN backend did not provide detailed error information";
|
||
|
|
}
|
||
|
|
|
||
|
|
extern "C" void yolo_backend_destroy(yolo_runtime_backend_t* backend) {
|
||
|
|
if (!backend) return;
|
||
|
|
if (backend->initialized) {
|
||
|
|
#if defined(YOLO_BACKEND_SEG)
|
||
|
|
release_yolov8_seg_model(&backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_CLS)
|
||
|
|
release_mobilenet_model(&backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_POSE)
|
||
|
|
release_yolov8_pose_model(&backend->context);
|
||
|
|
#elif defined(YOLO_BACKEND_OBB)
|
||
|
|
release_yolov8_obb_model(&backend->context);
|
||
|
|
#else
|
||
|
|
release_yolov8_model(&backend->context);
|
||
|
|
#endif
|
||
|
|
}
|
||
|
|
delete backend;
|
||
|
|
}
|