blob: dba615b97ef30cbba862c3e0b5b41ca383179ba4 [file] [log] [blame]
alexanderf42f5682021-07-16 11:30:56 +01001# Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
2# SPDX-License-Identifier: MIT
3
4"""
5Object detection demo that takes a video stream from a device, runs inference
6on each frame producing bounding boxes and labels around detected objects,
7and displays a window with the latest processed frame.
8"""
9
10import os
11import sys
12script_dir = os.path.dirname(__file__)
13sys.path.insert(1, os.path.join(script_dir, '..', 'common'))
14
15import cv2
16from argparse import ArgumentParser
17
18from ssd import ssd_processing, ssd_resize_factor
19from yolo import yolo_processing, yolo_resize_factor
20from utils import dict_labels
21from cv_utils import init_video_stream_capture, preprocess, draw_bounding_boxes
22from network_executor import ArmnnNetworkExecutor
23
24
25def get_model_processing(model_name: str, video: cv2.VideoCapture, input_binding_info: tuple):
26 """
27 Gets model-specific information such as model labels and decoding and processing functions.
28 The user can include their own network and functions by adding another statement.
29
30 Args:
31 model_name: Name of type of supported model.
32 video: Video capture object, contains information about data source.
33 input_binding_info: Contains shape of model input layer, used for scaling bounding boxes.
34
35 Returns:
36 Model labels, decoding and processing functions.
37 """
38 if model_name == 'ssd_mobilenet_v1':
39 return ssd_processing, ssd_resize_factor(video)
40 elif model_name == 'yolo_v3_tiny':
41 return yolo_processing, yolo_resize_factor(video, input_binding_info)
42 else:
43 raise ValueError(f'{model_name} is not a valid model name')
44
45
46def main(args):
47 video = init_video_stream_capture(args.video_source)
48 executor = ArmnnNetworkExecutor(args.model_file_path, args.preferred_backends)
49
50 model_name = args.model_name
51 process_output, resize_factor = get_model_processing(args.model_name, video, executor.input_binding_info)
52 labels = dict_labels(args.label_path, include_rgb=True)
53
54 while True:
55 frame_present, frame = video.read()
56 frame = cv2.flip(frame, 1) # Horizontally flip the frame
57 if not frame_present:
58 raise RuntimeError('Error reading frame from video stream')
59
60 if model_name == "ssd_mobilenet_v1":
61 input_tensors = preprocess(frame, executor.input_binding_info, True)
62 else:
63 input_tensors = preprocess(frame, executor.input_binding_info, False)
64 print("Running inference...")
65 output_result = executor.run(input_tensors)
66 detections = process_output(output_result)
67 draw_bounding_boxes(frame, detections, resize_factor, labels)
68 cv2.imshow('PyArmNN Object Detection Demo', frame)
69 if cv2.waitKey(1) == 27:
70 print('\nExit key activated. Closing video...')
71 break
72 video.release(), cv2.destroyAllWindows()
73
74
75if __name__ == '__main__':
76 parser = ArgumentParser()
77 parser.add_argument('--video_source', type=int, default=0,
78 help='Device index to access video stream. Defaults to primary device camera at index 0')
79 parser.add_argument('--model_file_path', required=True, type=str,
80 help='Path to the Object Detection model to use')
81 parser.add_argument('--model_name', required=True, type=str,
82 help='The name of the model being used. Accepted options: ssd_mobilenet_v1, yolo_v3_tiny')
83 parser.add_argument('--label_path', required=True, type=str,
84 help='Path to the labelset for the provided model file')
85 parser.add_argument('--preferred_backends', type=str, nargs='+', default=['CpuAcc', 'CpuRef'],
86 help='Takes the preferred backends in preference order, separated by whitespace, '
87 'for example: CpuAcc GpuAcc CpuRef. Accepted options: [CpuAcc, CpuRef, GpuAcc]. '
88 'Defaults to [CpuAcc, CpuRef]')
89 args = parser.parse_args()
90 main(args)