Added Video Support

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(
+            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)