aboutsummaryrefslogtreecommitdiff
path: root/python/pyarmnn/examples/object_detection/run_video_file.py
diff options
context:
space:
mode:
authorJakub Sujak <jakub.sujak@arm.com>2020-11-24 16:39:21 +0000
committerFrancis Murtagh <francis.murtagh@arm.com>2020-11-24 18:33:16 +0000
commit3346375a629ff30db25b115e3f4a1c3fdb16adf0 (patch)
tree15627567fcff445196db96a9dcef1db9268e6fee /python/pyarmnn/examples/object_detection/run_video_file.py
parentafc18650dcbe18b76d2ff44bcec31061da9abb60 (diff)
downloadarmnn-3346375a629ff30db25b115e3f4a1c3fdb16adf0.tar.gz
Removing labels and replacing with links to model+labels
Change-Id: Iec6c0b7cf55e0aa3ec1f0013f2da40e93b9bbcfc Signed-off-by: Jakub Sujak <jakub.sujak@arm.com>
Diffstat (limited to 'python/pyarmnn/examples/object_detection/run_video_file.py')
-rw-r--r--python/pyarmnn/examples/object_detection/run_video_file.py12
1 files changed, 5 insertions, 7 deletions
diff --git a/python/pyarmnn/examples/object_detection/run_video_file.py b/python/pyarmnn/examples/object_detection/run_video_file.py
index fc3e214721..e31b779458 100644
--- a/python/pyarmnn/examples/object_detection/run_video_file.py
+++ b/python/pyarmnn/examples/object_detection/run_video_file.py
@@ -36,11 +36,9 @@ 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(script_dir, 'ssd_labels.txt')
- return labels, ssd_processing, ssd_resize_factor(video)
+ return ssd_processing, ssd_resize_factor(video)
elif model_name == 'yolo_v3_tiny':
- labels = os.path.join(script_dir, 'yolo_labels.txt')
- return labels, yolo_processing, yolo_resize_factor(video, input_binding_info)
+ return yolo_processing, yolo_resize_factor(video, input_binding_info)
else:
raise ValueError(f'{model_name} is not a valid model name')
@@ -49,8 +47,8 @@ def main(args):
video, video_writer, frame_count = init_video_file_capture(args.video_file_path, args.output_video_file_path)
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)
+ process_output, resize_factor = get_model_processing(args.model_name, video, executor.input_binding_info)
+ labels = dict_labels(args.label_path, include_rgb=True)
for _ in tqdm(frame_count, desc='Processing frames'):
frame_present, frame = video.read()
@@ -73,7 +71,7 @@ if __name__ == '__main__':
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,
+ parser.add_argument('--label_path', required=True, type=str,
help='Path to the labelset for the provided model file')
parser.add_argument('--output_video_file_path', type=str,
help='Path to the output video file with detections added in')