mmdetection之初体验

修改了mmdetection中demo/wencam_demo.py，做个一个检测视频的小实验
代码及分析如下：

import argparse
import os
import cv2
import torch

from mmdet.apis import inference_detector, init_detector

# 获取当前文件的路径
file_path = __file__

# dirname（file_path）表示获取file_path的父路径
dir_path = os.path.dirname(file_path)

# 检测后的MP4路径
default_video_path = os.path.join(dir_path, 'result.mp4')

def parse_args():
    parser = argparse.ArgumentParser(description='MMDetection webcam demo')
    # 选择检测方法，注意这里如果不在命令行键入，则config前面需要加上--。
    parser.add_argument('--config', default='/gpfs/home/xxy_1997032/wjl/mmdetection/configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py', help='test config file path')
    # 选择预训练权重
    parser.add_argument('--checkpoint', default='/gpfs/home/xxy_1997032/wjl/mmdetection/checkpoints/mask_rcnn_r101_fpn_2x_coco_bbox_mAP-0.408__segm_mAP-0.366_20200505_071027-14b391c7.pth', help='checkpoint file')
    parser.add_argument(
        '--device', type=str, default='cuda:0', help='CPU/CUDA device option')
    # parser.add_argument(
    #     '--camera-id', type=int, default=0, help='camera device id')
    # 原视频路径
    parser.add_argument('--file', default='/gpfs/home/xxy_1997032/wjl/mmdetection/utils/a.mp4',type=str, help='test video path')
    parser.add_argument('--out', type=str, default=default_video_path, help='the output video path')
    # 检测的阈值，小于阈值的则不显示检测框
    parser.add_argument(
        '--score-thr', type=float, default=0.5, help='bbox score threshold')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    if not args.file:
        print('No target file!')
        exit(0)
	# 若没有GPU，则使用cpu
    if not torch.cuda.is_available():
        args.device = 'cpu'

    device = torch.device(args.device)

	# init_detector三个输入参数：
  # config：检测模型的配置文件，一般位于 mmdetection/config/ 中
  # checkpoint：即训练好的权重 在mmdetection的model zoo中可以下载
  # device：分配到的设备对象
  # 返回：model（检测模型）
    model = init_detector(args.config, args.checkpoint, device=device)
    
	# 读如MP4文件
    capture = cv2.VideoCapture(args.file)
	
	# 设置检测后得每一帧图像得长和宽
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

	#每一帧一帧得读入视频，视频则被处理成一张张得图片
    fps = capture.get(cv2.CAP_PROP_FPS)
	
	# 处理后得视频的存储地址、存储格式、每一帧、长和宽
    video_writer = cv2.VideoWriter(args.out, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))

   # 这里给每一帧计数，如果中断，可以继续检测
    count = 0
    print('Press "Esc", "q" or "Q" to exit.')
    while True:
        torch.cuda.empty_cache()
        # ret_val是一个布尔值，若这一帧有图片则为true，否则false
        ret_val, img = capture.read()
        if ret_val:
        # 这里如果上一次检测在count=900的时候停止了，则这里应该写为count < 901,表示跳过前900帧，
        # 最后在用其他方式将两个视频拼接起来
            if count < 0:
                count += 1
                print('Write {} in result Successfully!'.format(count))
                continue
            # inference_detector 输入两个参数：
    			 # model: init_detector返回的model
    			 # imgs: 图像的路径
					 # 返回：测试图像的结果
            result = inference_detector(model, img)

            ch = cv2.waitKey(1)
            if ch == 27 or ch == ord('q') or ch == ord('Q'):
                break
            # img :图像的路径（字符串）
    			 # result：inference_detector的返回值
    			 # class_name：训练图像的类别名称
    			 # 阈值
            frame = model.show_result(img, result, score_thr=args.score_thr, thickness=1, font_size=1, show=False)

            if len(frame) >= 1 or frame:
                video_writer.write(frame)
                count += 1
                print('Write {} in result Successfully!'.format(count))
        else:
            print('Try to read video fail at {}'.format(args.file))
            break

    capture.release()
    video_writer.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()