diff options
Diffstat (limited to 'python/pyarmnn/examples/object_detection/run_video_stream.py')
| -rw-r--r-- | python/pyarmnn/examples/object_detection/run_video_stream.py | 79 |
1 files changed, 32 insertions, 47 deletions
diff --git a/python/pyarmnn/examples/object_detection/run_video_stream.py b/python/pyarmnn/examples/object_detection/run_video_stream.py index 94dc6c8b13..9a303e8129 100644 --- a/python/pyarmnn/examples/object_detection/run_video_stream.py +++ b/python/pyarmnn/examples/object_detection/run_video_stream.py @@ -8,47 +8,18 @@ and displays a window with the latest processed frame. """
import os
+import sys
+script_dir = os.path.dirname(__file__)
+sys.path.insert(1, os.path.join(script_dir, '..', 'common'))
+
import cv2
-import pyarmnn as ann
-from tqdm import tqdm
from argparse import ArgumentParser
from ssd import ssd_processing, ssd_resize_factor
from yolo import yolo_processing, yolo_resize_factor
-from utils import create_network, dict_labels, preprocess, execute_network, draw_bounding_boxes
-
-
-parser = ArgumentParser()
-parser.add_argument('--video_source', type=int, default=0,
- help='Device index to access video stream. Defaults to primary device camera at index 0')
-parser.add_argument('--model_file_path', required=True, type=str,
- help='Path to the Object Detection model to use')
-parser.add_argument('--model_name', required=True, type=str,
- help='The name of the model being used. Accepted options: ssd_mobilenet_v1, yolo_v3_tiny')
-parser.add_argument('--label_path', type=str,
- help='Path to the labelset for the provided model file')
-parser.add_argument('--preferred_backends', type=str, nargs='+', default=['CpuAcc', 'CpuRef'],
- help='Takes the preferred backends in preference order, separated by whitespace, '
- 'for example: CpuAcc GpuAcc CpuRef. Accepted options: [CpuAcc, CpuRef, GpuAcc]. '
- 'Defaults to [CpuAcc, CpuRef]')
-args = parser.parse_args()
-
-
-def init_video(video_source: int):
- """
- Creates a video capture object from a device.
-
- Args:
- video_source: Device index used to read video stream.
-
- Returns:
- Video capture object used to capture frames from a video stream.
- """
- video = cv2.VideoCapture(video_source)
- if not video.isOpened:
- raise RuntimeError(f'Failed to open video capture for device with index: {video_source}')
- print('Processing video stream. Press \'Esc\' key to exit the demo.')
- return video
+from utils import dict_labels
+from cv_utils import init_video_stream_capture, preprocess, draw_bounding_boxes
+from network_executor import ArmnnNetworkExecutor
def get_model_processing(model_name: str, video: cv2.VideoCapture, input_binding_info: tuple):
@@ -65,31 +36,31 @@ def get_model_processing(model_name: str, video: cv2.VideoCapture, input_binding Model labels, decoding and processing functions.
"""
if model_name == 'ssd_mobilenet_v1':
- labels = os.path.join('ssd_labels.txt')
+ labels = os.path.join(script_dir, 'ssd_labels.txt')
return labels, ssd_processing, ssd_resize_factor(video)
elif model_name == 'yolo_v3_tiny':
- labels = os.path.join('yolo_labels.txt')
+ labels = os.path.join(script_dir, 'yolo_labels.txt')
return labels, yolo_processing, yolo_resize_factor(video, input_binding_info)
else:
raise ValueError(f'{model_name} is not a valid model name')
def main(args):
- video = init_video(args.video_source)
- net_id, runtime, input_binding_info, output_binding_info = create_network(args.model_file_path,
- args.preferred_backends)
- output_tensors = ann.make_output_tensors(output_binding_info)
- labels, process_output, resize_factor = get_model_processing(args.model_name, video, input_binding_info)
- labels = dict_labels(labels if args.label_path is None else args.label_path)
+ video = init_video_stream_capture(args.video_source)
+ executor = ArmnnNetworkExecutor(args.model_file_path, args.preferred_backends)
+
+ labels, process_output, resize_factor = get_model_processing(args.model_name, video, executor.input_binding_info)
+ labels = dict_labels(labels if args.label_path is None else args.label_path, include_rgb=True)
while True:
frame_present, frame = video.read()
frame = cv2.flip(frame, 1) # Horizontally flip the frame
if not frame_present:
raise RuntimeError('Error reading frame from video stream')
- input_tensors = preprocess(frame, input_binding_info)
- inference_output = execute_network(input_tensors, output_tensors, runtime, net_id)
- detections = process_output(inference_output)
+ input_tensors = preprocess(frame, executor.input_binding_info)
+ print("Running inference...")
+ output_result = executor.run(input_tensors)
+ detections = process_output(output_result)
draw_bounding_boxes(frame, detections, resize_factor, labels)
cv2.imshow('PyArmNN Object Detection Demo', frame)
if cv2.waitKey(1) == 27:
@@ -99,4 +70,18 @@ def main(args): if __name__ == '__main__':
+ parser = ArgumentParser()
+ parser.add_argument('--video_source', type=int, default=0,
+ help='Device index to access video stream. Defaults to primary device camera at index 0')
+ parser.add_argument('--model_file_path', required=True, type=str,
+ help='Path to the Object Detection model to use')
+ parser.add_argument('--model_name', required=True, type=str,
+ help='The name of the model being used. Accepted options: ssd_mobilenet_v1, yolo_v3_tiny')
+ parser.add_argument('--label_path', type=str,
+ help='Path to the labelset for the provided model file')
+ parser.add_argument('--preferred_backends', type=str, nargs='+', default=['CpuAcc', 'CpuRef'],
+ help='Takes the preferred backends in preference order, separated by whitespace, '
+ 'for example: CpuAcc GpuAcc CpuRef. Accepted options: [CpuAcc, CpuRef, GpuAcc]. '
+ 'Defaults to [CpuAcc, CpuRef]')
+ args = parser.parse_args()
main(args)
|