def demo(opt):
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str # 指定gpu的id号
opt.debug = max(opt.debug, 1) # 选择可视化的等级1,2,3,4
detector = Detector(opt) # 检测器对象
if opt.demo == 'webcam' or \ # 判断输入的是video还是webcam
opt.demo[opt.demo.rfind('.') + 1:].lower() in video_ext:
is_video = True # vedio
# demo on video stream
cam = cv2.VideoCapture(0 if opt.demo == 'webcam' else opt.demo)
else:
is_video = False # webcam
# Demo on images sequences
if os.path.isdir(opt.demo): # 判断输入的是文件夹还是图片
image_names = []
ls = os.listdir(opt.demo)
for file_name in sorted(ls):
ext = file_name[file_name.rfind('.') + 1:].lower()
if ext in image_ext:
image_names.append(os.path.join(opt.demo, file_name)) # 输入为图片序列
else:
image_names = [opt.demo]
# Initialize output video
out = None
out_name = opt.demo[opt.demo.rfind('/') + 1:]
if opt.save_video: # 输出结果是否保存为视频
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('../results/{}.mp4'.format( # 初始化输出的视频结果
opt.exp_id + '_' + out_name),fourcc, opt.save_framerate, (
opt.video_w, opt.video_h))
if opt.debug < 5:
detector.pause = False
cnt = 0 # cnt记录读取第几帧,每次循环加一
results = {} # 结果保存到results
while True:
if is_video: # 输入为视频序列,每次读取一帧到img
_, img = cam.read()
if img is None:
save_and_exit(opt, out, results, out_name)
else: # 读取输入的图片序列,每次读取一帧到img
if cnt < len(image_names):
img = cv2.imread(image_names[cnt])
else:
save_and_exit(opt, out, results, out_name)
cnt += 1
# resize the original video for saving video results
if opt.resize_video: # resize输入的图片
img = cv2.resize(img, (opt.video_w, opt.video_h))
# skip the first X frames of the video
if cnt < opt.skip_first:
continue
cv2.imshow('input', img) # 显示本帧图片
# track or detect the image.
ret = detector.run(img) # 使用检测器对本帧图片检测,返回到ret中
# log run time
time_str = 'frame {} |'.format(cnt)
for stat in time_stats:
time_str = time_str + '{} {:.3f}s |'.format(stat, ret[stat])
print(time_str)
# results[cnt] is a list of dicts:
# [{'bbox': [x1, y1, x2, y2], 'tracking_id': id, 'category_id': c, ...}]
results[cnt] = ret['results']
# save debug image to video
if opt.save_video:
out.write(ret['generic'])
if not is_video:
cv2.imwrite('../results/demo{}.jpg'.format(cnt), ret['generic'])
# esc to quit and finish saving video
if cv2.waitKey(1) == 27:
save_and_exit(opt, out, results, out_name)
return
save_and_exit(opt, out, results)
demo函数
最新推荐文章于 2022-03-12 14:43:54 发布