From bc9a325c87f9ffc9d38b55b8f8df6c383690e22a Mon Sep 17 00:00:00 2001 From: Pratik Bhave <31940311+pratikbhave2@users.noreply.github.com> Date: Thu, 3 Oct 2019 14:04:37 -0500 Subject: [PATCH 1/2] Added Video Support Using this file, we can perform inference on a video file and save the output as a new video file. --- tools/infer_simple_video.py | 170 ++++++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 tools/infer_simple_video.py diff --git a/tools/infer_simple_video.py b/tools/infer_simple_video.py new file mode 100644 index 000000000..aa0be0b52 --- /dev/null +++ b/tools/infer_simple_video.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python + + +"""Perform inference on a single video file and save +output as a new video file. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from __future__ import unicode_literals + +from collections import defaultdict +import argparse +import cv2 # NOQA (Must import before importing caffe2 due to bug in cv2) +import glob +import logging +import os +import sys +import time + +from caffe2.python import workspace + +from detectron.core.config import assert_and_infer_cfg +from detectron.core.config import cfg +from detectron.core.config import merge_cfg_from_file +from detectron.utils.io import cache_url +from detectron.utils.logging import setup_logging +from detectron.utils.timer import Timer +import detectron.core.test_engine as infer_engine +import detectron.datasets.dummy_datasets as dummy_datasets +import detectron.utils.c2 as c2_utils +import detectron.utils.vis as vis_utils + +c2_utils.import_detectron_ops() + +# OpenCL may be enabled by default in OpenCV3; disable it because it's not +# thread safe and causes unwanted GPU memory allocations. +cv2.ocl.setUseOpenCL(False) + + +def parse_args(): + parser = argparse.ArgumentParser(description='End-to-end inference') + parser.add_argument( + '--cfg', + dest='cfg', + help='cfg model file (/path/to/model_config.yaml)', + default=None, + type=str + ) + parser.add_argument( + '--wts', + dest='weights', + help='weights model file (/path/to/model_weights.pkl)', + default=None, + type=str + ) + parser.add_argument( + '--output-name', + dest='output_name', + help='output name of video along with entire path', + required = True, + type=str + ) + parser.add_argument( + '--video-dir', + dest='video_dir', + help='path to video file', + required= True, + type=str + ) + parser.add_argument( + '--always-out', + dest='out_when_no_box', + help='output image even when no object is found', + action='store_true' + ) + parser.add_argument( + '--thresh', + dest='thresh', + help='Threshold for visualizing detections', + default=0.5, + type=float + ) + parser.add_argument( + '--kp-thresh', + dest='kp_thresh', + help='Threshold for visualizing keypoints', + default=2.0, + type=float + ) + if len(sys.argv) == 1: + parser.print_help() + sys.exit(1) + return parser.parse_args() + + +def main(args): + logger = logging.getLogger(__name__) + + merge_cfg_from_file(args.cfg) + cfg.NUM_GPUS = 1 + args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE) + assert_and_infer_cfg(cache_urls=False) + + assert not cfg.MODEL.RPN_ONLY, \ + 'RPN models are not supported' + assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \ + 'Models that require precomputed proposals are not supported' + + model = infer_engine.initialize_model_from_cfg(args.weights) + dummy_coco_dataset = dummy_datasets.get_coco_dataset() + + vs = cv2.VideoCapture(args.video_dir) + writer = None + count = 0 + + while True: + grab,im = vs.read() + if not grab: + break + count += 1 + timers = defaultdict(Timer) + t = time.time() + with c2_utils.NamedCudaScope(0): + cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all( + model, im, None, timers=timers + ) + logger.info('Inference time: {:.3f}s'.format(time.time() - t)) + for k, v in timers.items(): + logger.info(' | {}: {:.3f}s'.format(k, v.average_time)) + if count == 1: + logger.info( + ' \ Note: inference on the first image will be slower than the ' + 'rest (caches and auto-tuning need to warm up)' + ) + + image_ret = vis_utils.vis_one_image_opencv_copy( + im, + cls_boxes, + cls_segms, + cls_keyps, + dataset=dummy_coco_dataset, + show_box = True, + show_class=True, + thresh=args.thresh, + kp_thresh=args.kp_thresh, + ) + + cv2.imshow("img", image_ret) + if cv2.waitKey(1) & 0xFF == ord('q'): + cv2.destroyAllWindows() + break + + if writer is None: + # initialize our video writer + fourcc = cv2.VideoWriter_fourcc(*"MJPG") + writer = cv2.VideoWriter(args.output_name, fourcc, 30, + (image_ret.shape[1], image_ret.shape[0]), True) + + writer.write(image_ret) + + writer.release() + vs.release() + +if __name__ == '__main__': + workspace.GlobalInit(['caffe2', '--caffe2_log_level=0']) + setup_logging(__name__) + args = parse_args() + main(args) From a33f350c8f6cfa53fa062fbd3352ddb1b77d2f13 Mon Sep 17 00:00:00 2001 From: Pratik Bhave <31940311+pratikbhave2@users.noreply.github.com> Date: Thu, 3 Oct 2019 14:07:43 -0500 Subject: [PATCH 2/2] updated infer_simple_video.py --- tools/infer_simple_video.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/infer_simple_video.py b/tools/infer_simple_video.py index aa0be0b52..2aa3d3995 100644 --- a/tools/infer_simple_video.py +++ b/tools/infer_simple_video.py @@ -135,7 +135,7 @@ def main(args): 'rest (caches and auto-tuning need to warm up)' ) - image_ret = vis_utils.vis_one_image_opencv_copy( + image_ret = vis_utils.vis_one_image_opencv( im, cls_boxes, cls_segms,