support YOLOV3 & YOLOV5 model

Signed-off-by: zhangdanfeng <[email protected]>

Author: zhangdanfeng

Makefile-rv

@@ -112,9 +112,9 @@ segmentation/tflite_segmentation: segmentation.cc $(COMMON_SRC)
 	mkdir -p segmentation
 	$(CXX) segmentation.cc $(COMMON_SRC) -o segmentation/tflite_segmentation $(LDFLAGS) $(LIBS) $(CXXFLAGS) $(CCFLAGS) $(INCLUDES)
 
-detection/tflite_detection: detection.cc yolov5.cc $(COMMON_SRC)
+detection/tflite_detection: detection.cc yolov5.cc yolov3.cc $(COMMON_SRC)
 	mkdir -p detection
-	$(CXX) detection.cc yolov5.cc $(COMMON_SRC) -o detection/tflite_detection $(LDFLAGS) $(LIBS) $(CXXFLAGS) $(CCFLAGS) $(INCLUDES)
+	$(CXX) detection.cc yolov5.cc yolov3.cc $(COMMON_SRC) -o detection/tflite_detection $(LDFLAGS) $(LIBS) $(CXXFLAGS) $(CCFLAGS) $(INCLUDES)
 
 clean:
 	rm -rf classification/tflite_classification segmentation/tflite_segmentation detection/tflite_detection

detection.cc

@@ -39,11 +39,22 @@
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/model.h"
 #include "utils.h"
+#include "yolov3.h"
 #include "yolov5.h"
 
 using namespace cv;
 using namespace std;
 
+/*
+qemu-riscv64 detection/tflite_detection -m \
+detection/yolov5s_ultralytics_640_quantized.tflite -i detection/bus.jpg -l \
+detection/labels.txt -c 1 -b 0 -s 255 -t 1 -v 5
+
+qemu-riscv64 detection/tflite_detection -m \
+detection/yolov3_keras_416_quantized.tflite -i detection/bus.jpg -l \
+detection/labels.txt -c 1 -b 0 -s 255 -t 1 -v 3
+*/
+
 /*
  * Display command line usage
  */
@@ -60,6 +71,7 @@ void display_usage() {
        << "--input_std, -s: input standard deviation\n"
        << "--profiling, -p: [0|1], profiling or not\n"
        << "--threads, -t: number of threads\n"
+       << "--model-version, -v: yolo version\n"
        << "\n";
 }
 
@@ -72,6 +84,7 @@ int main(int argc, char **argv) {
   std::string label_path = "labels.txt";
   std::string input_path = "grace_hopper.bmp";
   eInputType input_source = INPUT_Image;
+  int yolo_version = 5;
   int frame_cnt = 1;
   int num_threads = 1;
   float input_mean = 0.f;
@@ -90,12 +103,13 @@ int main(int argc, char **argv) {
         {"threads", required_argument, nullptr, 't'},
         {"input_mean", required_argument, nullptr, 'b'},
         {"input_std", required_argument, nullptr, 's'},
+        {"--model-version", required_argument, nullptr, 'v'},
         {nullptr, 0, nullptr, 0}};
 
     /* getopt_long stores the option index here. */
     int option_index = 0;
 
-    c = getopt_long(argc, argv, "b:c:i:m:p:r:s:t:h", long_options,
+    c = getopt_long(argc, argv, "b:c:i:m:l:p:r:s:t:v:h", long_options,
                     &option_index);
 
     /* Detect the end of the options. */
@@ -130,6 +144,9 @@ int main(int argc, char **argv) {
     case 't':
       num_threads = strtol(optarg, nullptr, 10);
       break;
+    case 'v':
+      yolo_version = strtol(optarg, nullptr, 10);
+      break;
     case 'h':
       display_usage();
       exit(-1);
@@ -138,16 +155,25 @@ int main(int argc, char **argv) {
     }
   }
 
-  YOLOV5 model;
+  YOLOV5 *model = NULL;
+  if (yolo_version == 3) {
+    model = new YOLOV3;
+  } else {
+    model = new YOLOV5;
+  }
+  if (!model) {
+    exit(-1);
+  }
+
   Prediction out_pred;
   std::vector<std::string> labelNames;
 
   std::cout << "Loading model... " << std::endl;
 
   // Read model.
-  model.loadModel(model_path);
+  model->loadModel(model_path);
 
-  model.getLabelsName(label_path, labelNames);
+  model->getLabelsName(label_path, labelNames);
   std::cout << "\nLabel Count: " << labelNames.size() << "\n" << std::endl;
 
   // Setup input
@@ -163,7 +189,7 @@ int main(int argc, char **argv) {
   // Predict on the input image
   cv::Mat show_image;
   input_image.copyTo(show_image);
-  model.run(input_image, out_pred);
+  model->run(input_image, out_pred);
   auto stop = std::chrono::high_resolution_clock::now();
   auto duration =
       std::chrono::duration_cast<std::chrono::milliseconds>(stop - start);
@@ -188,6 +214,6 @@ int main(int argc, char **argv) {
   cv::imwrite("out.png", show_image);
 
   std::cout << "detection completes! " << std::endl;
-
+  delete model;
   return 0;
 }

scripts/tf/build-linux-riscv64.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# tf v2.11.1 tested
+
+set -e
+rm -rf build_tflite_2_11_x
+mkdir -p build_tflite_2_11_x
+
+CMAKE_ARGS=()
+
+# CMake-level configuration
+CMAKE_ARGS+=("-DCMAKE_MAKE_PROGRAM=make")
+CMAKE_ARGS+=("-DCMAKE_TOOLCHAIN_FILE=$(pwd)/scripts/tf/cmake/riscv64.JDSK.toolchain.cmake")
+CMAKE_ARGS+=("-DCMAKE_BUILD_TYPE=RelWithDebInfo")
+
+# Cross-compilation options for Google Benchmark
+CMAKE_ARGS+=("-DHAVE_POSIX_REGEX=0")
+CMAKE_ARGS+=("-DHAVE_STEADY_CLOCK=0")
+CMAKE_ARGS+=("-DHAVE_STD_REGEX=0")
+
+# xnnpack
+CMAKE_ARGS+=("-DTFLITE_ENABLE_XNNPACK=OFF")
+
+# flatbuffer
+# 防止交叉编译出的elf直接运行报错
+CMAKE_ARGS+=("-DFLATBUFFERS_BUILD_FLATC=OFF")
+# 最好有预装
+CMAKE_ARGS+=("-DFLATBUFFERS_FLATC_EXECUTABLE=flatc")
+
+# Use-specified CMake arguments go last to allow overridding defaults
+CMAKE_ARGS+=($@)
+
+pushd build_tflite_2_11_x
+
+# 需要手动修改tensorflow的地址
+cmake ../../tensorflow/tensorflow/lite "${CMAKE_ARGS[@]}"
+make -j4
+
+popd

scripts/tf/cmake/riscv64.JDSK.toolchain.cmake

@@ -10,6 +10,7 @@
 
 set(CMAKE_SYSTEM_NAME Linux)
 set(CMAKE_SYSTEM_PROCESSOR riscv64)
+set(CMAKE_CROSSCOMPILING TRUE)
 
 if(DEFINED ENV{RISCV_ROOT_PATH})
     file(TO_CMAKE_PATH $ENV{RISCV_ROOT_PATH} RISCV_ROOT_PATH)

yolov3.cc

@@ -7,7 +7,6 @@ std::vector<std::vector<float>> YOLOV3::tensorToVector2D() {
     std::cout << __FILE__ << ": " << __LINE__ << std::endl;
     exit(-1);
   }
-  const size_t num_anchors = 3;
   int masks[3][3] = {
       {6, 7, 8},
       {3, 4, 5},
@@ -30,15 +29,16 @@ std::vector<std::vector<float>> YOLOV3::tensorToVector2D() {
     if (pOutputTensor->type == kTfLiteFloat32) {
       for (size_t i = 0; i < _out_row; i++) {
         for (size_t j = 0; j < _out_colum; j++) {
-          for (size_t k = 0; k < num_anchors; k++) {
+          for (size_t k = 0; k < _num_anchors; k++) {
             std::vector<float> vtem;
-            for (int l = 0; l < _out_channel / num_anchors; l++) {
+            for (int l = 0; l < _out_channel / _num_anchors; l++) {
               float val_float =
                   pOutputTensor->data
                       .f[i * _out_colum * _out_channel + j * _out_channel +
-                         k * _out_channel / num_anchors + l];
+                         k * _out_channel / _num_anchors + l];
               if (l != 2 && l != 3) {
-                val_float = 1. / (1. + exp(-val_float)); // logistic
+                val_float =
+                    1. / (1. + exp(-val_float)); // logistic (redundancy)
               }
               vtem.push_back(val_float);
             }
@@ -78,7 +78,7 @@ void YOLOV3::nonMaximumSupprition(std::vector<std::vector<float>> &predV,
       int w = predV[i][2] * _img_width;
       int h = predV[i][3] * _img_height;
 
-      for (int j = 5; j < 85; j++) {
+      for (int j = 5; j < _out_channel / _num_anchors; j++) {
         // # conf = obj_conf * cls_conf
         scores.push_back(predV[i][j] * predV[i][4]);
       }

yolov3.h

@@ -20,7 +20,7 @@
 
 class YOLOV3 : public YOLOV5 {
 public:
-  // private:
+  const size_t _num_anchors{3};
   virtual std::vector<std::vector<float>> tensorToVector2D() override;
   virtual void nonMaximumSupprition(std::vector<std::vector<float>> &predV,
                                     std::vector<cv::Rect> &boxes,

yolov5.cc

@@ -163,12 +163,7 @@ void YOLOV5::nonMaximumSupprition(std::vector<std::vector<float>> &predV,
   cv::Point classId;
 
   for (int i = 0; i < _out_row; i++) {
-    if (predV[i][4] > 0.5) {
-      printf("rowid:%d score:%f\n", i, predV[i][4]);
-    }
-
     if (predV[i][4] > _conf_threshold) {
-      // height--> image.rows,  width--> image.cols;
       int left = (predV[i][0] - predV[i][2] / 2) * _img_width;
       int top = (predV[i][1] - predV[i][3] / 2) * _img_height;
       int w = predV[i][2] * _img_width;

yolov5.h

@@ -57,7 +57,6 @@ class YOLOV5 {
   // number of threads
   const int _n_threads = 1;
 
-  // private:
   // model's
   std::unique_ptr<tflite::FlatBufferModel> _model;
   std::unique_ptr<tflite::Interpreter> _interpreter;