使用mmdetection一帧一帧检测图片并拼接成视频

在demo文件夹下创建demo.py，使用mmdetection训练好的模型，对视频每一帧进行检测并画上bounding boxes，然后把这些帧拼接回一个视频。

import argparse
import numpy as np
import torch
from mmdet.apis import inference_detector, init_detector, show_result
import cv2
from moviepy.editor import VideoFileClip

def parse_args():
	#config和checpoint必填
    parser = argparse.ArgumentParser(description='MMDetection webcam demo')
    parser.add_argument('--config', help='test config file path',default="/home/liucs/mmdetection/config_clc/faster_rcnn_r101_fpn_1x.py")
    parser.add_argument('--checkpoint', help='checkpoint file',default="/home/liucs/mmdetection/work_dirs/clc/epoch_12.pth")
    parser.add_argument('--device', type=int, default=0, help='CUDA device id')
    parser.add_argument(
        '--camera-id', type=int, default=0, help='camera device id')
    parser.add_argument(
        '--score-thr', type=float, default=0.5, help='bbox score threshold')
    args = parser.parse_args()
    return args
    
    
def draw_bbox(result,img,score_thr,class_names):
    if isinstance(result, tuple):
        bbox_result, segm_result = result
    else:
        bbox_result, segm_result = result, None#n,5 for eachclass
 
    bboxes = np.vstack(bbox_result)#
   
    labels = [
        np.full(bbox.shape[0], i, dtype=np.int32)
        for i, bbox in enumerate(bbox_result)#bbox_result=[[first_class,first_class,first_class],[secon_class],,,,]  
    ]
    labels = np.concatenate(labels)
    
    
    assert bboxes.ndim == 2
    assert labels.ndim == 1
    assert bboxes.shape[0] == labels.shape[0]
    assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
    
    img = np.ascontiguousarray(img)

    if score_thr > 0:
        assert bboxes.shape[1] == 5
        scores = bboxes[:, -1]
        inds = scores > score_thr
        bboxes = bboxes[inds, :]
        labels = labels[inds]
    #[mmcv中imshow_det_bboxes函数](https://github.com/open-mmlab/mmcv/blob/master/mmcv/visualization/image.py)
    bbox_color = (0,255,0)#这里只有一类，所以框的颜色我都用绿色来标注了。多种类参考上面链接
    text_color = (0,255,0)
    thickness=1
    font_scale=0.5
    for bbox, label in zip(bboxes, labels):
        bbox_int = bbox.astype(np.int32)
        left_top = (bbox_int[0], bbox_int[1])
        right_bottom = (bbox_int[2], bbox_int[3])
        
        cv2.rectangle(
            img, left_top, right_bottom, bbox_color, thickness=thickness)
        label_text = class_names[
            label] if class_names is not None else f'cls {label}'
        if len(bbox) > 4:
            label_text += f'|{bbox[-1]:.02f}'
        cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2),
                    cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
    return img
    #cv2.imwrite('/home/liucs/mmdetection/demo_clc/demo_008.jpg',img)


def process_image_1( frame): #返回每一帧的检测结果
    args = parse_args()
    model = init_detector(
        args.config, args.checkpoint, device=torch.device('cuda', args.device))
    score_thr=args.score_thr
    class_names=model.CLASSES
    result = inference_detector(model, frame)
    img = draw_bbox(result,frame,score_thr,class_names)
    return img
    
def main():
    white_output = "/home/chenlc/am/demo/DJI_0006_crop_demo.MP4"#输出视频的路径
    clip1 = VideoFileClip("/home/chenlc/am/demo/DJI_0006_resize.avi").subclip(5,7)#输入视频。（5,7）表示输入视频的5到7秒。
    white_clip = clip1.fl_image(process_image_1) #NOTE: this function expects color images!!
    white_clip.write_videofile(white_output, audio=False)
if __name__ == '__main__':
    main()