diff --git a/yolov3-spp/CMakeLists.txt b/yolov3-spp/CMakeLists.txt index 18762eb..a466470 100644 --- a/yolov3-spp/CMakeLists.txt +++ b/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") diff --git a/yolov3-spp/YoloConfigs.h b/yolov3-spp/YoloConfigs.h deleted file mode 100644 index 2d2547b..0000000 --- a/yolov3-spp/YoloConfigs.h +++ /dev/null @@ -1,55 +0,0 @@ -#ifndef _YOLO_CONFIGS_H_ -#define _YOLO_CONFIGS_H_ - - -namespace Yolo -{ - static constexpr int CHECK_COUNT = 3; - static constexpr float IGNORE_THRESH = 0.5f; - static constexpr int CLASS_NUM = 80; - static constexpr int INPUT_H = 256; - static constexpr int INPUT_W = 416; - - struct YoloKernel - { - int width; - int height; - float anchors[CHECK_COUNT*2]; - }; - - //YOLO 608 - YoloKernel yolo1 = { - 13, - 8, - {116,90, 156,198, 373,326} - }; - YoloKernel yolo2 = { - 26, - 16, - {30,61, 62,45, 59,119} - }; - YoloKernel yolo3 = { - 52, - 32, - {10,13, 16,30, 33,23} - }; - - //YOLO 416 - // YoloKernel yolo1 = { - // 13, - // 13, - // {116,90, 156,198, 373,326} - // }; - // YoloKernel yolo2 = { - // 26, - // 26, - // {30,61, 62,45, 59,119} - // }; - // YoloKernel yolo3 = { - // 52, - // 52, - // {10,13, 16,30, 33,23} - // }; -} - -#endif diff --git a/yolov3-spp/yololayer.cu b/yolov3-spp/yololayer.cu index 91ad035..7db4107 100644 --- a/yolov3-spp/yololayer.cu +++ b/yolov3-spp/yololayer.cu @@ -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); diff --git a/yolov3-spp/yololayer.h b/yolov3-spp/yololayer.h index 2d250c6..406c954 100644 --- a/yolov3-spp/yololayer.h +++ b/yolov3-spp/yololayer.h @@ -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{ diff --git a/yolov3-spp/yolov3-spp.cpp b/yolov3-spp/yolov3-spp.cpp index 7eb73ef..e007cfb 100644 --- a/yolov3-spp/yolov3-spp.cpp +++ b/yolov3-spp/yolov3-spp.cpp @@ -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& res, float *output, float nms_thresh = 0.4) { - std::map> m; +void nms(std::vector& res, float *output, float nms_thresh = 0.4) { + std::map> 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()); + if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); 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 res; + std::vector res; nms(res, prob); auto end = std::chrono::system_clock::now(); std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl;