haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

yolov3-spp update

Browse Source
This commit is contained in:
wang-xinyu 2020-04-09 09:55:45 +08:00
parent aeee6ac153
commit 6df8efd0c0
5 changed files with 47 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
yolov3-spp/CMakeLists.txt
View File

@ -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")

55
yolov3-spp/YoloConfigs.h
View File

@ -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

3
yolov3-spp/yololayer.cu
View File

@ -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);

29
yolov3-spp/yololayer.h
View File

@ -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{

17
yolov3-spp/yolov3-spp.cpp
View File

@ -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;