AlexeyAB
2017-08-09 ee0dec90de9802912a2c5ccd4e9854c3e72bf701
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
#pragma once
#include <memory>
#include <vector>
#include <deque>
#include <algorithm>
 
#ifdef OPENCV
#include <opencv2/opencv.hpp>           // C++
#include "opencv2/highgui/highgui_c.h"  // C
#include "opencv2/imgproc/imgproc_c.h"  // C
#endif  // OPENCV
 
#ifdef YOLODLL_EXPORTS
#if defined(_MSC_VER)
#define YOLODLL_API __declspec(dllexport) 
#else
#define YOLODLL_API __attribute__((visibility("default")))
#endif
#else
#if defined(_MSC_VER)
#define YOLODLL_API __declspec(dllimport) 
#else
#define YOLODLL_API
#endif
#endif
 
struct bbox_t {
    unsigned int x, y, w, h;    // (x,y) - top-left corner, (w, h) - width & height of bounded box
    float prob;                 // confidence - probability that the object was found correctly
    unsigned int obj_id;        // class of object - from range [0, classes-1]
    unsigned int track_id;      // tracking id for video (0 - untracked, 1 - inf - tracked object)
};
 
struct image_t {
    int h;                      // height
    int w;                      // width
    int c;                      // number of chanels (3 - for RGB)
    float *data;                // pointer to the image data
};
 
 
class Detector {
    std::shared_ptr<void> detector_gpu_ptr;
public:
    float nms = .4;
 
    YOLODLL_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id = 0);
    YOLODLL_API ~Detector();
 
    YOLODLL_API std::vector<bbox_t> detect(std::string image_filename, float thresh = 0.2, bool use_mean = false);
    YOLODLL_API std::vector<bbox_t> detect(image_t img, float thresh = 0.2, bool use_mean = false);
    static YOLODLL_API image_t load_image(std::string image_filename);
    static YOLODLL_API void free_image(image_t m);
    YOLODLL_API int get_net_width();
    YOLODLL_API int get_net_height();
 
    YOLODLL_API std::vector<bbox_t> tracking(std::vector<bbox_t> cur_bbox_vec, int const frames_story = 6);
 
#ifdef OPENCV
    std::vector<bbox_t> detect(cv::Mat mat, float thresh = 0.2, bool use_mean = false)
    {
        if(mat.data == NULL)
            throw std::runtime_error("file not found");
        cv::Mat det_mat;
        cv::resize(mat, det_mat, cv::Size(get_net_width(), get_net_height()));
        auto image_ptr = mat_to_image(det_mat);
        auto detection_boxes = detect(*image_ptr, thresh, use_mean);
        float wk = (float)mat.cols / det_mat.cols, hk = (float)mat.rows / det_mat.rows;
        for (auto &i : detection_boxes) i.x*=wk, i.w*= wk, i.y*=hk, i.h*=hk;
        return detection_boxes;
    }
 
    static std::shared_ptr<image_t> mat_to_image(cv::Mat img)
    {
        std::shared_ptr<image_t> image_ptr(new image_t, [](image_t *img) { free_image(*img); delete img; });
        std::shared_ptr<IplImage> ipl_small = std::make_shared<IplImage>(img);
        *image_ptr = ipl_to_image(ipl_small.get());
        rgbgr_image(*image_ptr);
        return image_ptr;
    }
 
private:
    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;
        }
    }
 
#endif  // OPENCV
 
    std::deque<std::vector<bbox_t>> prev_bbox_vec_deque;
};