AlexeyAB
2017-09-11 c6ccb80f17fc3fbdf51fc75650c52686ea97e618
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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#include <iostream>
#include <iomanip> 
#include <string>
#include <vector>
#include <fstream>
#include <thread>
#include <atomic>
 
#define OPENCV
 
#include "yolo_v2_class.hpp"    // imported functions from DLL
 
#ifdef OPENCV
#include <opencv2/opencv.hpp>           // C++
#include "opencv2/core/version.hpp"
#ifndef CV_VERSION_EPOCH
#include "opencv2/videoio/videoio.hpp"
#pragma comment(lib, "opencv_world320.lib")  
#else
#pragma comment(lib, "opencv_core2413.lib")  
#pragma comment(lib, "opencv_imgproc2413.lib")  
#pragma comment(lib, "opencv_highgui2413.lib") 
#endif
 
 
void draw_boxes(cv::Mat mat_img, std::vector<bbox_t> result_vec, std::vector<std::string> obj_names, 
    unsigned int wait_msec = 0, int current_fps = -1)
{
    for (auto &i : result_vec) {
        cv::Scalar color(60, 160, 260);
        cv::rectangle(mat_img, cv::Rect(i.x, i.y, i.w, i.h), color, 5);
        if (obj_names.size() > i.obj_id) {
            std::string obj_name = obj_names[i.obj_id];
            if (i.track_id > 0) obj_name += " - " + std::to_string(i.track_id);
            cv::Size const text_size = getTextSize(obj_name, cv::FONT_HERSHEY_COMPLEX_SMALL, 1.2, 2, 0);
            int const max_width = (text_size.width > i.w + 2) ? text_size.width : (i.w + 2);
            cv::rectangle(mat_img, cv::Point2f(std::max((int)i.x - 3, 0), std::max((int)i.y - 30, 0)), 
                cv::Point2f(std::min((int)i.x + max_width, mat_img.cols-1), std::min((int)i.y, mat_img.rows-1)), 
                color, CV_FILLED, 8, 0);
            putText(mat_img, obj_name, cv::Point2f(i.x, i.y - 10), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.2, cv::Scalar(0, 0, 0), 2);
        }
    }
    if(current_fps >= 0)
        putText(mat_img, "FPS: " + std::to_string(current_fps), cv::Point2f(10, 20), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.2, cv::Scalar(50, 255, 0), 2);
    cv::imshow("window name", mat_img);
    cv::waitKey(wait_msec);
}
#endif  // OPENCV
 
 
void show_result(std::vector<bbox_t> const result_vec, std::vector<std::string> const obj_names) {
    for (auto &i : result_vec) {
        if (obj_names.size() > i.obj_id) std::cout << obj_names[i.obj_id] << " - ";
        std::cout << "obj_id = " << i.obj_id << ",  x = " << i.x << ", y = " << i.y 
            << ", w = " << i.w << ", h = " << i.h
            << std::setprecision(3) << ", prob = " << i.prob << std::endl;
    }
}
 
std::vector<std::string> objects_names_from_file(std::string const filename) {
    std::ifstream file(filename);
    std::vector<std::string> file_lines;
    if (!file.is_open()) return file_lines;
    for(std::string line; file >> line;) file_lines.push_back(line);
    std::cout << "object names loaded \n";
    return file_lines;
}
 
 
int main(int argc, char *argv[])
{
    std::string filename;
    if (argc > 1) filename = argv[1];
 
    Detector detector("yolo-voc.cfg", "yolo-voc.weights");
 
    auto obj_names = objects_names_from_file("data/voc.names");
    std::string out_videofile = "result.avi";
    bool const save_output_videofile = false;
 
    while (true) 
    {       
        std::cout << "input image or video filename: ";
        if(filename.size() == 0) std::cin >> filename;
        if (filename.size() == 0) break;
        
        try {
#ifdef OPENCV
            std::string const file_ext = filename.substr(filename.find_last_of(".") + 1);
            std::string const protocol = filename.substr(0, 7);
            if (file_ext == "avi" || file_ext == "mp4" || file_ext == "mjpg" || file_ext == "mov" ||    // video file
                protocol == "rtsp://" || protocol == "http://" || protocol == "https:/")    // video network stream
            {
                cv::Mat frame, prev_frame, det_frame;
                std::vector<bbox_t> result_vec, thread_result_vec;
                detector.nms = 0.02;    // comment it - if track_id is not required
                std::thread td([]() {});
                std::atomic<int> ready_flag;
                ready_flag = true;
                std::chrono::steady_clock::time_point steady_start, steady_end;
                int fps_counter = 0, current_fps = 0;
                cv::VideoWriter output_video;
                for (cv::VideoCapture cap(filename); cap >> frame, cap.isOpened();) {
                    if(!output_video.isOpened() && save_output_videofile)
                        output_video.open(out_videofile, CV_FOURCC('D','I','V','X'), cap.get(CV_CAP_PROP_FPS), frame.size(), true);
                    if (ready_flag || (protocol != "rtsp://" && protocol != "http://" && protocol != "https:/")) {
                        td.join();
                        ++fps_counter;
                        ready_flag = false;
                        result_vec = thread_result_vec;
                        result_vec = detector.tracking(result_vec); // comment it - if track_id is not required
                        det_frame = frame.clone();                      
                        td = std::thread([&]() { thread_result_vec = detector.detect(det_frame, 0.24, true); ready_flag = true; });
                    }
                    if (!prev_frame.empty()) {  
                        steady_end = std::chrono::steady_clock::now();
                        if (std::chrono::duration<double>(steady_end - steady_start).count() >= 1) {
                            current_fps = fps_counter;
                            steady_start = steady_end;
                            fps_counter = 0;
                        }
                        draw_boxes(prev_frame, result_vec, obj_names, 3, current_fps);
                        show_result(result_vec, obj_names);
                        if (output_video.isOpened())
                            output_video << prev_frame;
                    }
                    prev_frame = frame.clone();                 
                }
            }
            else if (file_ext == "txt") {   // list of image files
                std::ifstream file(filename);
                if (!file.is_open()) std::cout << "File not found! \n";
                else 
                    for (std::string line; file >> line;) {
                        std::cout << line << std::endl;
                        cv::Mat mat_img = cv::imread(line);
                        std::vector<bbox_t> result_vec = detector.detect(mat_img);
                        show_result(result_vec, obj_names);
                        //draw_boxes(mat_img, result_vec, obj_names);
                        //cv::imwrite("res_" + line, mat_img);
                    }
                
            }
            else {  // image file
                cv::Mat mat_img = cv::imread(filename);
                std::vector<bbox_t> result_vec = detector.detect(mat_img);
                result_vec = detector.tracking(result_vec); // comment it - if track_id is not required
                draw_boxes(mat_img, result_vec, obj_names);
                show_result(result_vec, obj_names);
            }
#else
            //std::vector<bbox_t> result_vec = detector.detect(filename);
 
            auto img = detector.load_image(filename);
            std::vector<bbox_t> result_vec = detector.detect(img);
            detector.free_image(img);
            show_result(result_vec, obj_names);
#endif          
        }
        catch (std::exception &e) { std::cerr << "exception: " << e.what() << "\n"; getchar(); }
        catch (...) { std::cerr << "unknown exception \n"; getchar(); }
        filename.clear();
    }
 
    return 0;
}