def main():
parser = ArgumentParser()
parser.add_argument('config', help='Config file')
parser.add_argument('--input', help='input video file')
parser.add_argument('--output', help='output video file (mp4 format)')
parser.add_argument('--checkpoint', help='Checkpoint file')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--color',
default=(0, 255, 0),
help='Color of tracked bbox lines.')
parser.add_argument('--thickness',
default=3,
type=int,
help='Thickness of bbox lines.')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
cap = cv2.VideoCapture(args.input)
save_out_video = False
if args.output is not None:
save_out_video = True
fps = cap.get(cv2.CAP_PROP_FPS)
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(args.output, fourcc, fps, size)
frame_id = 0
while (cap.isOpened()):
flag, frame = cap.read()
if not flag:
break
if frame_id == 0:
init_bbox = list(cv2.selectROI(args.input, frame, False, False))
# convert (x1, y1, w, h) to (x1, y1, x2, y2)
init_bbox[2] += init_bbox[0]
init_bbox[3] += init_bbox[1]
# test a single image
result = inference_sot(model, frame, init_bbox, frame_id)
vis_frame = model.show_result(frame,
result,
color=args.color,
thickness=args.thickness,
show=False)
if save_out_video:
videoWriter.write(vis_frame)
if args.show:
cv2.imshow(args.input, vis_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
frame_id += 1
cap.release()
if save_out_video:
videoWriter.release()
cv2.destroyAllWindows()
def main():
parser = ArgumentParser()
parser.add_argument('config', help='Config file')
parser.add_argument('--input', help='input video file')
parser.add_argument('--output', help='output video file (mp4 format)')
parser.add_argument('--checkpoint', help='Checkpoint file')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--color',
default=(0, 255, 0),
help='Color of tracked bbox lines.')
parser.add_argument('--thickness',
default=3,
type=int,
help='Thickness of bbox lines.')
parser.add_argument('--fps', help='FPS of the output video')
parser.add_argument('--gt_bbox_file', help='The path of gt_bbox file')
args = parser.parse_args()
# load images
if osp.isdir(args.input):
imgs = sorted(filter(lambda x: x.endswith(('.jpg', '.png', '.jpeg')),
os.listdir(args.input)),
key=lambda x: int(x.split('.')[0]))
IN_VIDEO = False
else:
imgs = mmcv.VideoReader(args.input)
IN_VIDEO = True
# define output
if args.output is not None:
if args.output.endswith('.mp4'):
OUT_VIDEO = True
out_dir = tempfile.TemporaryDirectory()
out_path = out_dir.name
_out = args.output.rsplit(os.sep, 1)
if len(_out) > 1:
os.makedirs(_out[0], exist_ok=True)
else:
OUT_VIDEO = False
out_path = args.output
os.makedirs(out_path, exist_ok=True)
fps = args.fps
if args.show or OUT_VIDEO:
if fps is None and IN_VIDEO:
fps = imgs.fps
if not fps:
raise ValueError('Please set the FPS for the output video.')
fps = int(fps)
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
prog_bar = mmcv.ProgressBar(len(imgs))
# test and show/save the images
for i, img in enumerate(imgs):
if isinstance(img, str):
img = osp.join(args.input, img)
img = mmcv.imread(img)
if i == 0:
if args.gt_bbox_file is not None:
bboxes = mmcv.list_from_file(args.gt_bbox_file)
init_bbox = list(map(float, bboxes[0].split(',')))
else:
init_bbox = list(cv2.selectROI(args.input, img, False, False))
# convert (x1, y1, w, h) to (x1, y1, x2, y2)
init_bbox[2] += init_bbox[0]
init_bbox[3] += init_bbox[1]
result = inference_sot(model, img, init_bbox, frame_id=i)
if args.output is not None:
if IN_VIDEO or OUT_VIDEO:
out_file = osp.join(out_path, f'{i:06d}.jpg')
else:
out_file = osp.join(out_path, img.rsplit('os.sep', 1)[-1])
else:
out_file = None
model.show_result(img,
result,
show=args.show,
wait_time=int(1000. / fps) if fps else 0,
out_file=out_file,
thickness=args.thickness)
prog_bar.update()
if args.output and OUT_VIDEO:
print(
f'\nmaking the output video at {args.output} with a FPS of {fps}')
mmcv.frames2video(out_path, args.output, fps=fps, fourcc='mp4v')
out_dir.cleanup()