mirror of https://github.com/AlexeyAB/darknet.git
AlexeyAB
8 years ago
parent
a71bdd7a83
commit
a6cbaeecde
9 changed files with 828 additions and 2 deletions
@ -0,0 +1,57 @@ |
||||
#include <iostream> |
||||
#include <string> |
||||
#include <vector> |
||||
|
||||
//#define OPENCV
|
||||
|
||||
#include "yolo_v2_class.hpp" // imported functions from DLL |
||||
|
||||
|
||||
#ifdef OPENCV |
||||
#include <opencv2/opencv.hpp> // C++ |
||||
#pragma comment(lib, "opencv_core249.lib") |
||||
#pragma comment(lib, "opencv_imgproc249.lib") |
||||
#pragma comment(lib, "opencv_highgui249.lib") |
||||
void draw_boxes(cv::Mat mat_img, std::vector<bbox_t> result_vec) { |
||||
for (auto &i : result_vec) { |
||||
cv::rectangle(mat_img, cv::Rect(i.x, i.y, i.w, i.h), cv::Scalar(50, 200, 50), 3); |
||||
} |
||||
cv::imshow("window name", mat_img); |
||||
cv::waitKey(0); |
||||
} |
||||
#endif // OPENCV
|
||||
|
||||
|
||||
void show_result(std::vector<bbox_t> result_vec) { |
||||
for (auto &i : result_vec) { |
||||
std::cout << "obj_id = " << i.obj_id << " - x = " << i.x << ", y = " << i.y
|
||||
<< ", w = " << i.w << ", h = " << i.h |
||||
<< ", prob = " << i.prob << std::endl; |
||||
} |
||||
} |
||||
|
||||
|
||||
|
||||
int main()
|
||||
{ |
||||
Detector detector("yolo-voc.cfg", "yolo-voc.weights"); |
||||
|
||||
while (true)
|
||||
{ |
||||
std::string filename; |
||||
std::cout << "input image filename: "; |
||||
std::cin >> filename; |
||||
if (filename.size() == 0) break; |
||||
|
||||
#ifdef OPENCV |
||||
cv::Mat mat_img = cv::imread(filename); |
||||
std::vector<bbox_t> result_vec = detector.detect(mat_img); |
||||
draw_boxes(mat_img, result_vec); |
||||
#else |
||||
std::vector<bbox_t> result_vec = detector.detect(filename); |
||||
#endif |
||||
show_result(result_vec); |
||||
} |
||||
|
||||
return 0; |
||||
} |
||||
@ -0,0 +1,165 @@ |
||||
#include "yolo_v2_class.hpp" |
||||
|
||||
|
||||
#include "network.h" |
||||
|
||||
extern "C" { |
||||
#include "detection_layer.h" |
||||
#include "region_layer.h" |
||||
#include "cost_layer.h" |
||||
#include "utils.h" |
||||
#include "parser.h" |
||||
#include "box.h" |
||||
#include "image.h" |
||||
#include "demo.h" |
||||
|
||||
#include "option_list.h" |
||||
|
||||
} |
||||
//#include <sys/time.h>
|
||||
|
||||
#include <vector> |
||||
#include <iostream> |
||||
|
||||
|
||||
#define FRAMES 3 |
||||
#define ROI_PER_DETECTOR 100 |
||||
|
||||
|
||||
struct detector_gpu_t{ |
||||
float **probs; |
||||
box *boxes; |
||||
network net; |
||||
//image det;
|
||||
//image det_s;
|
||||
image images[FRAMES]; |
||||
float *avg; |
||||
float *predictions[FRAMES]; |
||||
}; |
||||
|
||||
|
||||
|
||||
YOLODLL_API Detector::Detector(std::string cfg_filename, std::string weight_filename, int gpu_id) |
||||
{ |
||||
int old_gpu_index; |
||||
cudaGetDevice(&old_gpu_index); |
||||
|
||||
detector_gpu_ptr = std::make_shared<detector_gpu_t>(); |
||||
|
||||
detector_gpu_t &detector_gpu = *reinterpret_cast<detector_gpu_t *>(detector_gpu_ptr.get()); |
||||
|
||||
cudaSetDevice(gpu_id); |
||||
network &net = detector_gpu.net; |
||||
net.gpu_index = gpu_id; |
||||
//gpu_index = i;
|
||||
|
||||
char *cfgfile = const_cast<char *>(cfg_filename.data()); |
||||
char *weightfile = const_cast<char *>(weight_filename.data()); |
||||
|
||||
net = parse_network_cfg(cfgfile); |
||||
if (weightfile) { |
||||
load_weights(&net, weightfile); |
||||
} |
||||
set_batch_network(&net, 1); |
||||
net.gpu_index = gpu_id; |
||||
|
||||
layer l = net.layers[net.n - 1]; |
||||
int j; |
||||
|
||||
detector_gpu.avg = (float *)calloc(l.outputs, sizeof(float)); |
||||
for (j = 0; j < FRAMES; ++j) detector_gpu.predictions[j] = (float *)calloc(l.outputs, sizeof(float)); |
||||
for (j = 0; j < FRAMES; ++j) detector_gpu.images[j] = make_image(1, 1, 3); |
||||
|
||||
detector_gpu.boxes = (box *)calloc(l.w*l.h*l.n, sizeof(box)); |
||||
detector_gpu.probs = (float **)calloc(l.w*l.h*l.n, sizeof(float *)); |
||||
for (j = 0; j < l.w*l.h*l.n; ++j) detector_gpu.probs[j] = (float *)calloc(l.classes, sizeof(float)); |
||||
|
||||
cudaSetDevice(old_gpu_index); |
||||
} |
||||
|
||||
YOLODLL_API Detector::~Detector()
|
||||
{ |
||||
detector_gpu_t &detector_gpu = *reinterpret_cast<detector_gpu_t *>(detector_gpu_ptr.get()); |
||||
layer l = detector_gpu.net.layers[detector_gpu.net.n - 1]; |
||||
|
||||
free(detector_gpu.boxes); |
||||
free(detector_gpu.avg); |
||||
free(detector_gpu.predictions); |
||||
for (int j = 0; j < l.w*l.h*l.n; ++j) free(detector_gpu.probs[j]); |
||||
free(detector_gpu.probs); |
||||
} |
||||
|
||||
|
||||
YOLODLL_API std::vector<bbox_t> Detector::detect(std::string image_filename, float thresh) |
||||
{ |
||||
char *input = const_cast<char *>(image_filename.data()); |
||||
image im = load_image_color(input, 0, 0); |
||||
|
||||
image_t img; |
||||
img.c = im.c; |
||||
img.data = im.data; |
||||
img.h = im.h; |
||||
img.w = im.w; |
||||
|
||||
return detect(img, thresh); |
||||
} |
||||
|
||||
|
||||
YOLODLL_API std::vector<bbox_t> Detector::detect(image_t img, float thresh) |
||||
{ |
||||
|
||||
detector_gpu_t &detector_gpu = *reinterpret_cast<detector_gpu_t *>(detector_gpu_ptr.get()); |
||||
network &net = detector_gpu.net; |
||||
int old_gpu_index; |
||||
cudaGetDevice(&old_gpu_index); |
||||
cudaSetDevice(net.gpu_index); |
||||
//std::cout << "net.gpu_index = " << net.gpu_index << std::endl;
|
||||
|
||||
float nms = .4; |
||||
|
||||
image im; |
||||
im.c = img.c; |
||||
im.data = img.data; |
||||
im.h = img.h; |
||||
im.w = img.w; |
||||
|
||||
image sized = resize_image(im, net.w, net.h); |
||||
layer l = net.layers[net.n - 1]; |
||||
|
||||
//box *boxes = (box *)calloc(l.w*l.h*l.n, sizeof(box));
|
||||
//float **probs = (float **)calloc(l.w*l.h*l.n, sizeof(float *));
|
||||
// (int j = 0; j < l.w*l.h*l.n; ++j) probs[j] = (float *)calloc(l.classes, sizeof(float *));
|
||||
|
||||
float *X = sized.data; |
||||
|
||||
network_predict(net, X); |
||||
|
||||
get_region_boxes(l, 1, 1, thresh, detector_gpu.probs, detector_gpu.boxes, 0, 0); |
||||
if (nms) do_nms_sort(detector_gpu.boxes, detector_gpu.probs, l.w*l.h*l.n, l.classes, nms); |
||||
//draw_detections(im, l.w*l.h*l.n, thresh, boxes, probs, names, alphabet, l.classes);
|
||||
|
||||
std::vector<bbox_t> bbox_vec; |
||||
|
||||
for (size_t i = 0; i < (l.w*l.h*l.n); ++i) { |
||||
box b = detector_gpu.boxes[i]; |
||||
int const obj_id = max_index(detector_gpu.probs[i], l.classes); |
||||
float const prob = detector_gpu.probs[i][obj_id]; |
||||
|
||||
if (prob > thresh)
|
||||
{ |
||||
bbox_t bbox; |
||||
bbox.x = (b.x - b.w / 2.)*im.w; |
||||
bbox.y = (b.y - b.h / 2.)*im.h; |
||||
bbox.w = b.w*im.w; |
||||
bbox.h = b.h*im.h; |
||||
bbox.obj_id = obj_id; |
||||
bbox.prob = prob; |
||||
|
||||
bbox_vec.push_back(bbox); |
||||
} |
||||
} |
||||
|
||||
cudaSetDevice(old_gpu_index); |
||||
|
||||
return bbox_vec; |
||||
} |
||||
@ -0,0 +1,120 @@ |
||||
#pragma once |
||||
#include <memory> |
||||
#include <vector> |
||||
|
||||
#ifdef OPENCV |
||||
#include <opencv2/opencv.hpp> // C++ |
||||
#include "opencv2/highgui/highgui_c.h" // C |
||||
#include "opencv2/imgproc/imgproc_c.h" // C |
||||
#endif // OPENCV
|
||||
|
||||
//extern "C" {
|
||||
//#include "image.h"
|
||||
//}
|
||||
|
||||
#ifdef YOLODLL_EXPORTS |
||||
#define YOLODLL_API __declspec(dllexport) |
||||
#else |
||||
#define YOLODLL_API __declspec(dllimport) |
||||
#endif |
||||
|
||||
struct bbox_t { |
||||
float x, y, w, h; |
||||
float prob; |
||||
unsigned int obj_id; |
||||
}; |
||||
|
||||
typedef struct { |
||||
int h; |
||||
int w; |
||||
int c; |
||||
float *data; |
||||
} image_t; |
||||
|
||||
|
||||
class Detector { |
||||
std::shared_ptr<void> detector_gpu_ptr; |
||||
public: |
||||
|
||||
YOLODLL_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id = 0); |
||||
YOLODLL_API ~Detector(); |
||||
|
||||
YOLODLL_API std::vector<bbox_t> Detector::detect(std::string image_filename, float thresh = 0.2); |
||||
|
||||
YOLODLL_API std::vector<bbox_t> detect(image_t img, float thresh = 0.2); |
||||
|
||||
|
||||
#ifdef OPENCV |
||||
std::vector<bbox_t> detect(cv::Mat mat, float thresh = 0.2) { |
||||
std::shared_ptr<image_t> image_ptr(new image_t, [](image_t *img) { free_image(*img); } ); |
||||
*image_ptr = mat_to_image(mat); |
||||
return detect(*image_ptr, thresh); |
||||
} |
||||
|
||||
private: |
||||
static image_t mat_to_image(cv::Mat img) |
||||
{ |
||||
std::shared_ptr<IplImage> ipl_small = std::make_shared<IplImage>(img); |
||||
image_t im_small = ipl_to_image(ipl_small.get()); |
||||
rgbgr_image(im_small); |
||||
return im_small; |
||||
} |
||||
|
||||
static image_t ipl_to_image(IplImage* src) |
||||
{ |
||||
unsigned char *data = (unsigned char *)src->imageData; |
||||
int h = src->height; |
||||
int w = src->width; |
||||
int c = src->nChannels; |
||||
int step = src->widthStep; |
||||
image_t out = make_image_custom(w, h, c); |
||||
int i, j, k, count = 0;; |
||||
|
||||
for (k = 0; k < c; ++k) { |
||||
for (i = 0; i < h; ++i) { |
||||
for (j = 0; j < w; ++j) { |
||||
out.data[count++] = data[i*step + j*c + k] / 255.; |
||||
} |
||||
} |
||||
} |
||||
return out; |
||||
} |
||||
|
||||
static image_t make_empty_image(int w, int h, int c) |
||||
{ |
||||
image_t out; |
||||
out.data = 0; |
||||
out.h = h; |
||||
out.w = w; |
||||
out.c = c; |
||||
return out; |
||||
} |
||||
|
||||
static image_t make_image_custom(int w, int h, int c) |
||||
{ |
||||
image_t out = make_empty_image(w, h, c); |
||||
out.data = (float *)calloc(h*w*c, sizeof(float)); |
||||
return out; |
||||
} |
||||
|
||||
static void rgbgr_image(image_t im) |
||||
{ |
||||
int i; |
||||
for (i = 0; i < im.w*im.h; ++i) { |
||||
float swap = im.data[i]; |
||||
im.data[i] = im.data[i + im.w*im.h * 2]; |
||||
im.data[i + im.w*im.h * 2] = swap; |
||||
} |
||||
} |
||||
|
||||
static void free_image(image_t m) |
||||
{ |
||||
if (m.data) { |
||||
free(m.data); |
||||
} |
||||
} |
||||
#endif // OPENCV
|
||||
}; |
||||
|
||||
|
||||
|
||||
Loading…
Reference in new issue