Merge pull request #8 from james20902/yolov8

.vscode/c_cpp_properties.json

@@ -0,0 +1,17 @@
+{
+    "configurations": [
+        {
+            "name": "Win32",
+            "includePath": [
+                "${workspaceFolder}/**"
+            ],
+            "defines": [
+                "_DEBUG",
+                "UNICODE",
+                "_UNICODE"
+            ],
+            "compileCommands": "${workspaceFolder}/opencv_test/compile_commands.json"
+        }
+    ],
+    "version": 4
+}

src/main/java/org/photonvision/rknn/RknnJNI.java

@@ -21,6 +21,11 @@
 import org.opencv.core.Rect2d;
 
 public class RknnJNI {
+    public static enum ModelVersion {
+        YOLO_V5,
+        YOLO_V8
+    }
+
     public static class RknnResult {
         public RknnResult(
             int left, int top, int right, int bottom, float conf, int class_id
@@ -75,9 +80,10 @@ public boolean equals(Object obj) {
      * Create a RKNN detector. Returns valid pointer on success, or NULL on error
      * @param modelPath Absolute path to the model on disk
      * @param numClasses How many classes. MUST MATCH or native code segfaults
+     * @param modelVer Which model is being used. Detections will be incorrect if not set to corrresponding model.
      * @return
      */
-    public static native long create(String modelPath, int numClasses);
+    public static native long create(String modelPath, int numClasses, int modelVer);
     public static native long destroy(long ptr);
 
     /**

src/main/native/cpp/rknn_jni.cpp

@@ -62,12 +62,12 @@ static jobject MakeJObject(JNIEnv *env, const detect_result_t &result) {
  */
 JNIEXPORT jlong JNICALL
 Java_org_photonvision_rknn_RknnJNI_create
-  (JNIEnv *env, jclass, jstring javaString, jint numClasses)
+  (JNIEnv *env, jclass, jstring javaString, jint numClasses, jint modelVer)
 {
   const char *nativeString = env->GetStringUTFChars(javaString, 0);
   std::printf("Creating for %s\n", nativeString);
 
-  auto ret = new RknnWrapper(nativeString, numClasses);
+  auto ret = new RknnWrapper(nativeString, numClasses, reinterpret_cast<ModelVersion>(modelVer));
   env->ReleaseStringUTFChars(javaString, nativeString);
   return reinterpret_cast<jlong>(ret);
 }

src/main/native/cpp/yolov8/postprocess.cc

@@ -23,6 +23,7 @@
 
 #include <set>
 #include <vector>
+#include "include/postprocess.h"
 #define LABEL_NALE_TXT_PATH "./model/coco_80_labels_list.txt"
 
 static char *labels[OBJ_CLASS_NUM];
@@ -351,7 +352,7 @@ static int process_fp32(float *box_tensor, float *score_tensor, float *score_sum
 }
 
 
-int post_process(rknn_app_context_t *app_ctx, rknn_output *outputs, letterbox_t *letter_box, float conf_threshold, float nms_threshold, object_detect_result_list *od_results)
+int post_process(rknn_app_context_t *app_ctx, rknn_output *outputs, BOX_RECT *padding, float conf_threshold, float nms_threshold, detect_result_group_t *od_results)
 {
     std::vector<float> filterBoxes;
     std::vector<float> objProbs;
@@ -363,7 +364,7 @@ int post_process(rknn_app_context_t *app_ctx, rknn_output *outputs, letterbox_t
     int model_in_w = app_ctx->model_width;
     int model_in_h = app_ctx->model_height;
 
-    memset(od_results, 0, sizeof(object_detect_result_list));
+    memset(od_results, 0, sizeof(detect_result_group_t));
 
     // default 3 branch
     int dfl_len = app_ctx->output_attrs[0].dims[1] /4;
@@ -425,6 +426,10 @@ int post_process(rknn_app_context_t *app_ctx, rknn_output *outputs, letterbox_t
     int last_count = 0;
     od_results->count = 0;
 
+    int width = padding->right - padding->left;
+    int height = padding->bottom - padding->top;
+    float rect_scale = static_cast<float>(width) / static_cast<float>(height);
+
     /* box valid detect target */
     for (int i = 0; i < validCount; ++i)
     {
@@ -434,19 +439,21 @@ int post_process(rknn_app_context_t *app_ctx, rknn_output *outputs, letterbox_t
         }
         int n = indexArray[i];
 
-        float x1 = filterBoxes[n * 4 + 0] - letter_box->x_pad;
-        float y1 = filterBoxes[n * 4 + 1] - letter_box->y_pad;
+        // if we run letterboxing, it gives us a BOX_RECT struct representing the box
+        // so we can get our offset by using the top left (0,0) by opencv coord frame
+        float x1 = filterBoxes[n * 4 + 0] - padding->left;
+        float y1 = filterBoxes[n * 4 + 1] - padding->top;
         float x2 = x1 + filterBoxes[n * 4 + 2];
         float y2 = y1 + filterBoxes[n * 4 + 3];
         int id = classId[n];
         float obj_conf = objProbs[i];
 
-        od_results->results[last_count].box.left = (int)(clamp(x1, 0, model_in_w) / letter_box->scale);
-        od_results->results[last_count].box.top = (int)(clamp(y1, 0, model_in_h) / letter_box->scale);
-        od_results->results[last_count].box.right = (int)(clamp(x2, 0, model_in_w) / letter_box->scale);
-        od_results->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / letter_box->scale);
-        od_results->results[last_count].prop = obj_conf;
-        od_results->results[last_count].cls_id = id;
+        od_results->results[last_count].box.left = (int)(clamp(x1, 0, model_in_w) / rect_scale);
+        od_results->results[last_count].box.top = (int)(clamp(y1, 0, model_in_h) / rect_scale);
+        od_results->results[last_count].box.right = (int)(clamp(x2, 0, model_in_w) / rect_scale);
+        od_results->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / rect_scale);
+        od_results->results[last_count].obj_conf = obj_conf;
+        od_results->results[last_count].id = id;
         last_count++;
     }
     od_results->count = last_count;

src/main/native/cpp/yolov8/yolov8.cc

@@ -21,6 +21,12 @@
 #include "common.h"
 #include "file_utils.h"
 #include "image_utils.h"
+#include "include/preprocess.h"
+
+#include "opencv2/core/core.hpp"
+#include "opencv2/highgui/highgui.hpp"
+#include "opencv2/imgproc/imgproc.hpp"
+
 
 static void dump_tensor_attr(rknn_tensor_attr *attr)
 {
@@ -155,63 +161,80 @@ int release_yolov8_model(rknn_app_context_t *app_ctx)
     return 0;
 }
 
-int inference_yolov8_model(rknn_app_context_t *app_ctx, image_buffer_t *img, object_detect_result_list *od_results)
+int inference_yolov8_model(rknn_app_context_t *app_ctx, cv::Mat &orig_img, detect_result_group_t *od_results, DetectionFilterParams params)
 {
     int ret;
-    image_buffer_t dst_img;
-    letterbox_t letter_box;
+    cv::Mat img;
+    BOX_RECT padding;
+    memset(&padding, 0, sizeof(BOX_RECT));
+
     rknn_input inputs[app_ctx->io_num.n_input];
     rknn_output outputs[app_ctx->io_num.n_output];
-    const float nms_threshold = NMS_THRESH;      // 默认的NMS阈值
-    const float box_conf_threshold = BOX_THRESH; // 默认的置信度阈值
+    const float nms_threshold = params.nms_thresh;
+    const float box_conf_threshold = params.box_thresh;
     int bg_color = 114;
 
-    if ((!app_ctx) || !(img) || (!od_results))
+    if ((!app_ctx) || !(orig_img) || (!od_results))
     {
         return -1;
     }
 
-    memset(od_results, 0x00, sizeof(*od_results));
-    memset(&letter_box, 0, sizeof(letterbox_t));
-    memset(&dst_img, 0, sizeof(image_buffer_t));
-    memset(inputs, 0, sizeof(inputs));
-    memset(outputs, 0, sizeof(outputs));
+    // get model desired width/height
+    width = app_ctx->model_width;
+    height = app_ctx->model_height;
 
-    // Pre Process
-    dst_img.width = app_ctx->model_width;
-    dst_img.height = app_ctx->model_height;
-    dst_img.format = IMAGE_FORMAT_RGB888;
-    dst_img.size = get_image_size(&dst_img);
-    dst_img.virt_addr = (unsigned char *)malloc(dst_img.size);
-    if (dst_img.virt_addr == NULL)
+    // rknn wants RGB, cameras usually give BGR
+    cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB);
+    // get current width/height
+    img_width = img.cols;
+    img_height = img.rows;
+
+    cv::Size target_size(width, height);
+    // create 8 bit 3 channel buffer
+    cv::Mat resized_img(target_size.height, target_size.width, CV_8UC3);
+
+    // Calculate scaling ratio
+    float scale_w = (float)target_size.width / img.cols;
+    float scale_h = (float)target_size.height / img.rows;
+
+    // if input doesn't match model desired
+    if (img_width != width || img_height != height)
     {
-        printf("malloc buffer size:%d fail!\n", dst_img.size);
-        return -1;
+        // rga
+        rga_buffer_t src;
+        rga_buffer_t dst;
+        memset(&src, 0, sizeof(src));
+        memset(&dst, 0, sizeof(dst));
+        ret = resize_rga(src, dst, img, resized_img, target_size);
+        if (ret != 0)
+        {
+            fprintf(stderr, "resize with rga error\n");
+        }
+        /*********
+        // opencv
+        // ripped from v5, why do we not letterbox???
+        float min_scale = std::min(scale_w, scale_h);
+        scale_w = min_scale;
+        scale_h = min_scale;
+        letterbox(img, resized_img, padding, min_scale, target_size);
+        *********/
+        inputs[0].buf = resized_img.data;
     }
-
-    // letterbox
-    ret = convert_image_with_letterbox(img, &dst_img, &letter_box, bg_color);
-    if (ret < 0)
+    else
     {
-        printf("convert_image_with_letterbox fail! ret=%d\n", ret);
-        return -1;
+        // otherwise just send the image as is
+        inputs[0].buf = img.data;
     }
 
-    // Set Input Data
-    inputs[0].index = 0;
-    inputs[0].type = RKNN_TENSOR_UINT8;
-    inputs[0].fmt = RKNN_TENSOR_NHWC;
-    inputs[0].size = app_ctx->model_width * app_ctx->model_height * app_ctx->model_channel;
-    inputs[0].buf = dst_img.virt_addr;
-
+    // load our image into npu
     ret = rknn_inputs_set(app_ctx->rknn_ctx, app_ctx->io_num.n_input, inputs);
     if (ret < 0)
     {
         printf("rknn_input_set fail! ret=%d\n", ret);
         return -1;
     }
 
-    // Run
+    // send it
     printf("rknn_run\n");
     ret = rknn_run(app_ctx->rknn_ctx, nullptr);
     if (ret < 0)
@@ -220,13 +243,15 @@ int inference_yolov8_model(rknn_app_context_t *app_ctx, image_buffer_t *img, obj
         return -1;
     }
 
-    // Get Output
+    // prep output storage
     memset(outputs, 0, sizeof(outputs));
     for (int i = 0; i < app_ctx->io_num.n_output; i++)
     {
         outputs[i].index = i;
         outputs[i].want_float = (!app_ctx->is_quant);
     }
+
+    // read outputs from npu
     ret = rknn_outputs_get(app_ctx->rknn_ctx, app_ctx->io_num.n_output, outputs, NULL);
     if (ret < 0)
     {
@@ -235,9 +260,9 @@ int inference_yolov8_model(rknn_app_context_t *app_ctx, image_buffer_t *img, obj
     }
 
     // Post Process
-    post_process(app_ctx, outputs, &letter_box, box_conf_threshold, nms_threshold, od_results);
+    post_process(app_ctx, outputs, &padding, box_conf_threshold, nms_threshold, od_results);
 
-    // Remeber to release rknn output
+    // release rknn output
     rknn_outputs_release(app_ctx->rknn_ctx, app_ctx->io_num.n_output, outputs);
 
 out:

src/main/native/cpp/yolov8/yolov8.h

@@ -18,8 +18,13 @@
 
 #include "rknn_api.h"
 #include "common.h"
+#include "postprocess.h"
+#include "include/postprocess.h"
 
-typedef struct {
+class rkYolov8s
+{
+private:
+ typedef struct {
     rknn_context rknn_ctx;
     rknn_input_output_num io_num;
     rknn_tensor_attr* input_attrs;
@@ -29,14 +34,34 @@ typedef struct {
     int model_height;
     bool is_quant;
 } rknn_app_context_t;
+rknn_app_context_t context;
 
-#include "postprocess.h"
+public:
+    // rkYolov5s(const std::string &model_path, int numClasses);
+
+    int init_yolov8_model(const char* model_path, rknn_app_context_t* app_ctx);
+
+    int init(const char* model_path) {
+        return init_yolov8_model(model_path, &context)
+    }
+
+    // rknn_context *get_pctx();
 
+    /**
+     * Run forward inference only, returning resulting detections
+    */
+    // int inferOnly(cv::Mat &orig_img, detect_result_group_t *outReults, DetectionFilterParams params);
 
-int init_yolov8_model(const char* model_path, rknn_app_context_t* app_ctx);
+    int inference_yolov8_model(rknn_app_context_t *app_ctx, cv::Mat &orig_img, detect_result_group_t *od_results, DetectionFilterParams params);
+    int infer(cv::Mat &orig_img, detect_result_group_t *od_results, DetectionFilterParams params) {
+        return inference_yolov8_model(&context, orig_img, od_results, params);
+    }
 
-int release_yolov8_model(rknn_app_context_t* app_ctx);
+    int release_yolov8_model(rknn_app_context_t* app_ctx);
 
-int inference_yolov8_model(rknn_app_context_t* app_ctx, image_buffer_t* img, object_detect_result_list* od_results);
+    ~rkYolov8s() {
+        release_yolov8_model(&context);
+    }
+};
 
 #endif //_RKNN_DEMO_YOLOV8_H_

src/main/native/include/rknn_wrapper.h

@@ -7,26 +7,51 @@
 #include "opencv2/imgproc/imgproc.hpp"
 #include "rkYolov5s.hpp"
 #include "rknnPool.hpp"
+#include "yolov8/yolov8.h"
 #include <optional>
 
+enum ModelVersion {
+    YOLO_V5,
+    YOLO_V8
+};
+
 class RknnWrapper
 {
 private:
-    rkYolov5s yolo;
+    rkYolov5s yolov5;
+    rkYolov8s yolov8;
     int m_numClasses;
 
 public:
-    RknnWrapper(const char *model_name, int numClasses) : yolo(model_name, numClasses)
+    RknnWrapper(const char *model_name, int numClasses, ModelVersion model_ver)
     { 
-        yolo.init(yolo.get_pctx(), false);
+        switch (model_ver) {
+            case YOLO_V5:
+                yolov5 = rkYolov5s(model_name, numClasses);
+                yolov5.init(yolov5.get_pctx(), false);
+            break;
+            case YOLO_V8:
+                yolov8 = rkYolov8s();
+                yolov8.init(model_name);
+            break;
+        }
     }
 
     detect_result_group_t forward(cv::Mat &img, DetectionFilterParams params) {
         detect_result_group_t ret;
-        int code = yolo.inferOnly(img, &ret, params);
+        int code;
+        switch (model_ver) {
+            case YOLO_V5:
+                code = yolov5.inferOnly(img, &ret, params);
+            break;
+            case YOLO_V8:
+                code = yolov8.infer(img, &ret, params);
+            break;
+        }
         return ret;
     }
 
     // Let the rkYolov5s dtor take care of cleanup
+    // how to destroy yolov8 processor?????
     ~RknnWrapper() = default;
 };