From c8dea4668d6223e96e2cf3df5eb92e6df2cf75e4 Mon Sep 17 00:00:00 2001 From: Francesco Gatti Date: Mon, 4 Feb 2019 20:34:15 +0000 Subject: [PATCH] compute detections --- include/Layer.h | 24 +++++- src/Yolo.cpp | 106 +++++++++++++++++++++++- tests/yolo3_berkeley/yolo3_berkeley.cpp | 16 +++- 3 files changed, 140 insertions(+), 6 deletions(-) diff --git a/include/Layer.h b/include/Layer.h index 12ffd31..4e53ffb 100644 --- a/include/Layer.h +++ b/include/Layer.h @@ -339,13 +339,33 @@ struct sortable_bbox { class Yolo : public Layer { public: - Yolo(Network *net, int classes, int num); + struct box { + float x, y, w, h; + }; + + typedef struct detection{ + Yolo::box bbox; + int classes; + float *prob; + float *mask; + float objectness; + int sort_class; + }; + + Yolo(Network *net, int classes, int num, const char* fname_weights); virtual ~Yolo(); virtual layerType_t getLayerType() { return LAYER_YOLO; }; int classes, num; - + dnnType *mask_h, *mask_d; //anchors + dnnType *bias_h, *bias_d; //anchors + virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); + int computeDetections(int w, int h, float thresh); + + const int MAX_DETECTIONS = 256; + Yolo::detection *dets; + int detected; }; /** diff --git a/src/Yolo.cpp b/src/Yolo.cpp index 6ebfb7c..bef4ccc 100644 --- a/src/Yolo.cpp +++ b/src/Yolo.cpp @@ -11,12 +11,28 @@ namespace tk { namespace dnn { -Yolo::Yolo(Network *net, int classes, int num) : +Yolo::detection *make_network_boxes(int nboxes, int classes) { + + int i; + Yolo::detection *dets = (Yolo::detection*) calloc(nboxes, sizeof(Yolo::detection)); + for(i = 0; i < nboxes; ++i){ + dets[i].prob = (float*) calloc(classes, sizeof(float)); + } + return dets; +} + +Yolo::Yolo(Network *net, int classes, int num, const char* fname_weights) : Layer(net) { this->classes = classes; this->num = num; + // load anchors + int seek = 0; + readBinaryFile(fname_weights, num, &mask_h, &mask_d); + seek += num; + readBinaryFile(fname_weights, 3*num, &bias_h, &bias_d); + // same output_dim.n = input_dim.n; output_dim.c = input_dim.c; @@ -25,6 +41,9 @@ Yolo::Yolo(Network *net, int classes, int num) : output_dim.l = input_dim.l; checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) ); + + dets = make_network_boxes(MAX_DETECTIONS, classes); + detected = 0; } Yolo::~Yolo() { @@ -39,6 +58,43 @@ int entry_index(int batch, int location, int entry, entry*input_dim.w*input_dim.h + loc; } +Yolo::box get_yolo_box(float *x, float *biases, int n, int index, int i, int j, int lw, int lh, int w, int h, int stride) { + Yolo::box b; + b.x = (i + x[index + 0*stride]) / lw; + b.y = (j + x[index + 1*stride]) / lh; + b.w = exp(x[index + 2*stride]) * biases[2*n] / w; + b.h = exp(x[index + 3*stride]) * biases[2*n+1] / h; + return b; +} + +void correct_yolo_boxes(Yolo::detection *dets, int n, int w, int h, int netw, int neth, int relative) +{ + int i; + int new_w=0; + int new_h=0; + if (((float)netw/w) < ((float)neth/h)) { + new_w = netw; + new_h = (h * netw)/w; + } else { + new_h = neth; + new_w = (w * neth)/h; + } + for (i = 0; i < n; ++i){ + Yolo::box b = dets[i].bbox; + b.x = (b.x - (netw - new_w)/2./netw) / ((float)new_w/netw); + b.y = (b.y - (neth - new_h)/2./neth) / ((float)new_h/neth); + b.w *= (float)netw/new_w; + b.h *= (float)neth/new_h; + if(!relative){ + b.x *= w; + b.w *= w; + b.y *= h; + b.h *= h; + } + dets[i].bbox = b; + } +} + dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) { @@ -58,4 +114,52 @@ dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) { return dstData; } +int Yolo::computeDetections(int w, int h, float thresh) { + + dnnType *predictions = new dnnType[output_dim.tot()]; + checkCuda( cudaMemcpy(predictions, dstData, output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToHost)); + + int relative = 1; + + int lw = output_dim.w; + int lh = output_dim.h; + int netw = net->input_dim.w; + int neth = net->input_dim.h; + + if (output_dim.n == 2) { + FatalError("BATCH of 2 not supported"); + //avg_flipped_yolo(l); + } + int i,j,n; + int count = 0; + for (i = 0; i < lw*lh; ++i){ + int row = i / lw; + int col = i % lw; + for(n = 0; n < num; ++n){ + int obj_index = entry_index(0, n*lw*lh + i, 4, classes, input_dim, output_dim); + float objectness = predictions[obj_index]; + if(objectness <= thresh) continue; + int box_index = entry_index(0, n*lw*lh + i, 0, classes, input_dim, output_dim); + + dets[count].bbox = get_yolo_box(predictions, bias_h, mask_h[n], box_index, col, row, lw, lh, netw, neth, lw*lh); + dets[count].objectness = objectness; + dets[count].classes = classes; + for(j = 0; j < classes; ++j){ + int class_index = entry_index(0, n*lw*lh + i, 4 + 1 + j, classes, input_dim, output_dim); + float prob = objectness*predictions[class_index]; + dets[count].prob[j] = (prob > thresh) ? prob : 0; + } + + ++count; + if(count >= MAX_DETECTIONS) + FatalError("reach max boxes"); + } + } + correct_yolo_boxes(dets, count, w, h, netw, neth, relative); + + std::cout<<"DETECTED: "<