yolov3-spp update

Author: wang-xinyu

yolov3-spp/CMakeLists.txt

@@ -13,8 +13,16 @@ find_package(CUDA REQUIRED)
 set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)
 
 include_directories(${PROJECT_SOURCE_DIR}/include)
-include_directories(/usr/local/cuda/targets/aarch64-linux/include)
-link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
+    message("embed_platform on")
+    include_directories(/usr/local/cuda/targets/aarch64-linux/include)
+    link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
+else()
+    message("embed_platform off")
+    include_directories(/usr/local/cuda/include)
+    link_directories(/usr/local/cuda/lib64)
+endif()
+
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
 

yolov3-spp/YoloConfigs.h

@@ -1,4 +1,3 @@
-#include "YoloConfigs.h"
 #include "yololayer.h"
 
 using namespace Yolo;
@@ -205,7 +204,7 @@ namespace nvinfer1
                 }
             }
             float box_prob = Logist(input[input_col + k * info_len_i * total_grid + 4 * total_grid]);
-            if (max_cls_prob < 0.1 || box_prob < 0.1) continue;
+            if (max_cls_prob < IGNORE_THRESH || box_prob < IGNORE_THRESH) continue;
 
             float *res_count = output;
             int count = (int)atomicAdd(res_count, 1);

yolov3-spp/yololayer.cu

@@ -12,7 +12,34 @@
 
 namespace Yolo
 {
-    struct YoloKernel;
+    static constexpr int CHECK_COUNT = 3;
+    static constexpr float IGNORE_THRESH = 0.1f;
+    static constexpr int CLASS_NUM = 80;
+    static constexpr int INPUT_H = 608;
+    static constexpr int INPUT_W = 608;
+
+    struct YoloKernel
+    {
+        int width;
+        int height;
+        float anchors[CHECK_COUNT*2];
+    };
+
+    static YoloKernel yolo1 = {
+        INPUT_W / 32,
+        INPUT_H / 32,
+        {116,90,  156,198,  373,326}
+    };
+    static YoloKernel yolo2 = {
+        INPUT_W / 16,
+        INPUT_H / 16,
+        {30,61,  62,45,  59,119}
+    };
+    static YoloKernel yolo3 = {
+        INPUT_W / 8,
+        INPUT_H / 8,
+        {10,13,  16,30,  33,23}
+    };
 
     static constexpr int LOCATIONS = 4;
     struct alignas(float) Detection{

yolov3-spp/yololayer.h

@@ -17,11 +17,10 @@
 #define DEVICE 0  // GPU id
 
 using namespace nvinfer1;
-using namespace Yolo;
 
 // stuff we know about the network and the input/output blobs
-static const int INPUT_H = 256;
-static const int INPUT_W = 416;
+static const int INPUT_H = Yolo::INPUT_H;
+static const int INPUT_W = Yolo::INPUT_W;
 static const int OUTPUT_SIZE = 1000 * 7 + 1;  // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1
 const char* INPUT_BLOB_NAME = "data";
 const char* OUTPUT_BLOB_NAME = "prob";
@@ -90,17 +89,17 @@ float iou(float lbox[4], float rbox[4]) {
     return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS);
 }
 
-bool cmp(Detection& a, Detection& b) {
+bool cmp(Yolo::Detection& a, Yolo::Detection& b) {
     return a.det_confidence > b.det_confidence;
 }
 
-void nms(std::vector<Detection>& res, float *output, float nms_thresh = 0.4) {
-    std::map<float, std::vector<Detection>> m;
+void nms(std::vector<Yolo::Detection>& res, float *output, float nms_thresh = 0.4) {
+    std::map<float, std::vector<Yolo::Detection>> m;
     for (int i = 0; i < output[0] && i < 1000; i++) {
         if (output[1 + 7 * i + 4] <= 0.5) continue;
-        Detection det;
+        Yolo::Detection det;
         memcpy(&det, &output[1 + 7 * i], 7 * sizeof(float));
-        if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector<Detection>());
+        if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector<Yolo::Detection>());
         m[det.class_id].push_back(det);
     }
     for (auto it = m.begin(); it != m.end(); it++) {
@@ -535,7 +534,7 @@ int main(int argc, char** argv) {
         // Run inference
         auto start = std::chrono::system_clock::now();
         doInference(*context, data, prob, 1);
-        std::vector<Detection> res;
+        std::vector<Yolo::Detection> res;
         nms(res, prob);
         auto end = std::chrono::system_clock::now();
         std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;