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 <sys/time.h> #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_d", 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