YOLO应用化之保存视频

YOLO的原生代码没有保存视频的功能，但我们往往需要一个视频文件作为检测结果。

使用命令

./darknet detector demo cfg/coco.data cfg/yolo.cfg yolo.weights

可以打开camera进行实时检测MSCOCO的８０种物体。（附加参数-c　camera_index 可以选择camera）

如果没有camera也可以使用命令

./darknet detector demo cfg/coco.data cfg/yolo.cfg yolo.weights input.avi

检测视频文件中的物体。

通过附加参数-prefix 可以将每一帧输出保存为jpg图片。

我们需要修改下代码将检测结果以视频文件形式保存。
已将代码提交到Github，先放链接
https://github.com/PaulChongPeng/darknet/commit/0366ba1fc0f8a640a11d300a3ef16382fbc530ef

共需要修改两个文件

demo.c

@@ -10,13 +10,18 @@
 #include 

 #define FRAMES 3
+#define SAVEVIDEO

 #ifdef OPENCV
 #include "opencv2/highgui/highgui_c.h"
 #include "opencv2/imgproc/imgproc_c.h"
 #include "opencv2/videoio/videoio_c.h"
 image get_image_from_stream(CvCapture *cap);

+#ifdef SAVEVIDEO
+static CvVideoWriter *mVideoWriter;
+#endif
+
 static char **demo_names;
 static image **demo_alphabet;
 static int demo_classes;

@@ -115,8 +121,20 @@ void demo(char *cfgfile, char *weightfile, float thresh, int cam_index, const ch
     if(filename){
         printf("video file: %s\n", filename);
         cap = cvCaptureFromFile(filename);
+#ifdef SAVEVIDEO
+        if(cap){
+            int mfps = cvGetCaptureProperty(cap,CV_CAP_PROP_FPS);
+            mVideoWriter=cvCreateVideoWriter("Output.avi",CV_FOURCC('M','J','P','G'),mfps,cvSize(cvGetCaptureProperty(cap,CV_CAP_PROP_FRAME_WIDTH),cvGetCaptureProperty(cap,CV_CAP_PROP_FRAME_HEIGHT)),1);
+        }
+#endif
     }else{
         cap = cvCaptureFromCAM(cam_index);
+#ifdef SAVEVIDEO
+        if(cap){
+            int mfps = cvGetCaptureProperty(cap,CV_CAP_PROP_FPS);
+            mVideoWriter=cvCreateVideoWriter("Output.avi",CV_FOURCC('M','J','P','G'),mfps,cvSize(cvGetCaptureProperty(cap,CV_CAP_PROP_FRAME_WIDTH),cvGetCaptureProperty(cap,CV_CAP_PROP_FRAME_HEIGHT)),1);
+        }
+#endif
     }

     if(!cap) error("Couldn't connect to webcam.\n");
@@ -169,6 +187,9 @@ void demo(char *cfgfile, char *weightfile, float thresh, int cam_index, const ch
             if(pthread_create(&detect_thread, 0, detect_in_thread, 0)) error("Thread creation failed");

             if(!prefix){
+#ifdef SAVEVIDEO
+                save_video(disp,mVideoWriter);    
+#endif
                 show_image(disp, "Demo");
                 int c = cvWaitKey(1);
                 if (c == 10){
@@ -180,7 +201,11 @@ void demo(char *cfgfile, char *weightfile, float thresh, int cam_index, const ch
             }else{
                 char buff[256];
                 sprintf(buff, "%s_%08d", prefix, count);
+#ifdef SAVEVIDEO
+                save_video(disp,mVideoWriter);
+#else
                 save_image(disp, buff);
+#endif
             }

             pthread_join(fetch_thread, 0);

image.c

     cvReleaseImage(&disp);
     free_image(copy);
 }
+
+
+void save_video(image p, CvVideoWriter *mVideoWriter)
+{
+    image copy = copy_image(p);
+    if(p.c == 3) rgbgr_image(copy);
+    int x,y,k;
+
+    IplImage *disp = cvCreateImage(cvSize(p.w,p.h), IPL_DEPTH_8U, p.c);
+    int step = disp->widthStep;
+    for(y = 0; y < p.h; ++y){
+        for(x = 0; x < p.w; ++x){
+            for(k= 0; k < p.c; ++k){
+                disp->imageData[y*step + x*p.c + k] = (unsigned char)(get_pixel(copy,x,y,k)*255);
+            }
+        }
+    }
+    cvWriteFrame(mVideoWriter,disp);
+    cvReleaseImage(&disp);
+    free_image(copy);
+}
 #endif

 void save_image_png(image im, const char *name)

测试

make clean
make -j8
./darknet detector demo cfg/coco.data cfg/yolo.cfg yolo.weights -c 0 -thresh 0.3