修改了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()