Go
Joseph Redmon
2016-03-14 d1965bdb969920c85f72785ec6e1f3d7bda957de
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
#include "network.h"
#include "detection_layer.h"
#include "cost_layer.h"
#include "utils.h"
#include "parser.h"
#include "box.h"
#include "image.h"
#include <sys/time.h>
 
#define FRAMES 1
 
#ifdef OPENCV
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
void convert_coco_detections(float *predictions, int classes, int num, int square, int side, int w, int h, float thresh, float **probs, box *boxes, int only_objectness);
 
extern char *coco_classes[];
extern image coco_labels[];
 
static float **probs;
static box *boxes;
static network net;
static image in   ;
static image in_s ;
static image det  ;
static image det_s;
static image disp ;
static CvCapture * cap;
static float fps = 0;
static float demo_thresh = 0;
 
static float *predictions[FRAMES];
static int demo_index = 0;
static image images[FRAMES];
static float *avg;
 
void *fetch_in_thread_coco(void *ptr)
{
    in = get_image_from_stream(cap);
    in_s = resize_image(in, net.w, net.h);
    return 0;
}
 
void *detect_in_thread_coco(void *ptr)
{
    float nms = .4;
 
    detection_layer l = net.layers[net.n-1];
    float *X = det_s.data;
    float *prediction = network_predict(net, X);
 
    memcpy(predictions[demo_index], prediction, l.outputs*sizeof(float));
    mean_arrays(predictions, FRAMES, l.outputs, avg);
 
    free_image(det_s);
    convert_coco_detections(avg, l.classes, l.n, l.sqrt, l.side, 1, 1, demo_thresh, probs, boxes, 0);
    if (nms > 0) do_nms(boxes, probs, l.side*l.side*l.n, l.classes, nms);
    printf("\033[2J");
    printf("\033[1;1H");
    printf("\nFPS:%.0f\n",fps);
    printf("Objects:\n\n");
 
    images[demo_index] = det;
    det = images[(demo_index + FRAMES/2 + 1)%FRAMES];
    demo_index = (demo_index + 1)%FRAMES;
 
    draw_detections(det, l.side*l.side*l.n, demo_thresh, boxes, probs, coco_classes, coco_labels, 80);
    return 0;
}
 
void demo_coco(char *cfgfile, char *weightfile, float thresh, int cam_index, const char *filename)
{
    demo_thresh = thresh;
    printf("COCO demo\n");
    net = parse_network_cfg(cfgfile);
    if(weightfile){
        load_weights(&net, weightfile);
    }
    set_batch_network(&net, 1);
 
    srand(2222222);
 
    if(filename){
        cap = cvCaptureFromFile(filename);
    }else{
        cap = cvCaptureFromCAM(cam_index);
    }
 
    if(!cap) error("Couldn't connect to webcam.\n");
    cvNamedWindow("COCO", CV_WINDOW_NORMAL); 
    cvResizeWindow("COCO", 512, 512);
 
    detection_layer l = net.layers[net.n-1];
    int j;
 
    avg = (float *) calloc(l.outputs, sizeof(float));
    for(j = 0; j < FRAMES; ++j) predictions[j] = (float *) calloc(l.outputs, sizeof(float));
    for(j = 0; j < FRAMES; ++j) images[j] = make_image(1,1,3);
 
    boxes = (box *)calloc(l.side*l.side*l.n, sizeof(box));
    probs = (float **)calloc(l.side*l.side*l.n, sizeof(float *));
    for(j = 0; j < l.side*l.side*l.n; ++j) probs[j] = (float *)calloc(l.classes, sizeof(float *));
 
    pthread_t fetch_thread;
    pthread_t detect_thread;
 
    fetch_in_thread_coco(0);
    det = in;
    det_s = in_s;
 
    fetch_in_thread_coco(0);
    detect_in_thread_coco(0);
    disp = det;
    det = in;
    det_s = in_s;
 
    for(j = 0; j < FRAMES/2; ++j){
        fetch_in_thread_coco(0);
        detect_in_thread_coco(0);
        disp = det;
        det = in;
        det_s = in_s;
    }
 
    while(1){
        struct timeval tval_before, tval_after, tval_result;
        gettimeofday(&tval_before, NULL);
        if(pthread_create(&fetch_thread, 0, fetch_in_thread_coco, 0)) error("Thread creation failed");
        if(pthread_create(&detect_thread, 0, detect_in_thread_coco, 0)) error("Thread creation failed");
        show_image(disp, "COCO");
        //save_image(disp, "COCO");
        free_image(disp);
        cvWaitKey(10);
        pthread_join(fetch_thread, 0);
        pthread_join(detect_thread, 0);
 
        disp  = det;
        det   = in;
        det_s = in_s;
 
        gettimeofday(&tval_after, NULL);
        timersub(&tval_after, &tval_before, &tval_result);
        float curr = 1000000.f/((long int)tval_result.tv_usec);
        fps = .9*fps + .1*curr;
    }
}
#else
void demo_coco(char *cfgfile, char *weightfile, float thresh, int cam_index){
    fprintf(stderr, "YOLO-COCO demo needs OpenCV for webcam images.\n");
}
#endif