From f15443d1264235bdfd3f26c3a9bb2c7d1caad602 Mon Sep 17 00:00:00 2001 From: itsnine Date: Mon, 27 Sep 2021 23:11:01 +0500 Subject: [PATCH 1/4] added onnxruntime python inference --- deployment/onnxruntime-python/README.md | 25 +++++++ deployment/onnxruntime-python/example.py | 91 ++++++++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 deployment/onnxruntime-python/README.md create mode 100644 deployment/onnxruntime-python/example.py diff --git a/deployment/onnxruntime-python/README.md b/deployment/onnxruntime-python/README.md new file mode 100644 index 00000000..f4d11cfb --- /dev/null +++ b/deployment/onnxruntime-python/README.md @@ -0,0 +1,25 @@ +# ONNXRuntime Inference + +The ONNXRuntime inference for `yolort`, both GPU and CPU are supported. + +## Dependencies + +- Ubuntu 20.04 / Windows 10 +- ONNXRuntime 1.7 + +- OpenCV 4.5 + +- CUDA 11 [Optional] + +*We didn't impose too strong restrictions on the versions of dependencies.* + +## Usage + +1. Update your PyTorch model weights from ultralytics to yolort and export to ONNX following the [notebooks with tutorials](https://github.com/zhiqwang/yolov5-rt-stack/blob/master/notebooks/). + +2. To infer your own images run the script as following. + + ```bash + python example.py [--image ../../../test/assets/zidane.jpg] + [--model_path ../../../test/tracing/yolov5s.onnx] + [--class_names ../../../notebooks/assets/coco.names] + [--gpu] # GPU switch, which is optional, and set False as default + ``` \ No newline at end of file diff --git a/deployment/onnxruntime-python/example.py b/deployment/onnxruntime-python/example.py new file mode 100644 index 00000000..9ed76b16 --- /dev/null +++ b/deployment/onnxruntime-python/example.py @@ -0,0 +1,91 @@ +import argparse +from typing import Tuple, List + +import numpy as np +import cv2 +import onnxruntime as ort + + +class YOLOv5Detector: + def __init__(self, model_path: str, gpu: bool = True) -> None: + ort_device = ort.get_device() + self._providers = None + + if ort_device == 'GPU' and gpu: + self._providers = ['CUDAExecutionProvider'] + print('Inference device: GPU') + elif gpu: + print('GPU is not supported by your ONNXRuntime build. Fallback to CPU.') + else: + self._providers = ['CPUExecutionProvider'] + print('Inference device: CPU') + + self._model = ort.InferenceSession(model_path, providers=self._providers) + self._input_names = self._model.get_inputs()[0].name + print('Model was initialized.') + + def _preprocessing(self, image: np.ndarray) -> np.ndarray: + rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + blob = rgb_image.astype(np.float32) / 255.0 + blob = np.transpose(blob, axes=(2, 0, 1)) + + return blob + + def __call__(self, image: np.ndarray) -> Tuple[List[float], List[int], List[List[int]]]: + blob = self._preprocessing(image) + scores, class_ids, boxes = self._model.run(output_names=None, + input_feed={self._input_names: blob}) + boxes = boxes.astype(np.int32) + return scores.tolist(), class_ids.tolist(), boxes.tolist() + + +def visualize_detection(image: np.ndarray, + class_names: List[str], + scores: List[float], + class_ids: List[int], + boxes: List[List[int]]) -> None: + + for i, class_id in enumerate(class_ids): + x1, y1, x2, y2 = boxes[i] + conf = round(scores[i], 2) + + label = class_names[class_id] + " " + str(conf) + text_size = cv2.getTextSize(label, cv2.FONT_ITALIC, 0.8, 2)[0] + + cv2.rectangle(image, (x1, y1), (x2, y2), (229, 160, 21), 2) + cv2.rectangle(image, (x1, y1 - 25), (x1 + text_size[0], y1), (229, 160, 21), -1) + cv2.putText(image, label, (x1, y1 - 3), cv2.FONT_ITALIC, 0.8, (255, 255, 255), 2) + + +def get_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser('ONNXRuntime inference script.', add_help=True) + parser.add_argument('--model_path', default='yolov5.onnx', type=str, required=True, + help='Path to onnx model.') + parser.add_argument('--image', default='bus.jpg', type=str, required=True, + help='Image source to be detected.') + parser.add_argument('--class_names', default='coco.names', type=str, required=True, + help='Path of dataset labels.') + parser.add_argument('--gpu', action='store_true', + help='To enable inference on GPU.') + + return parser + + +if __name__ == '__main__': + parser = get_parser() + args = parser.parse_args() + print('Command Line Args: ', args) + + with open(args.class_names, 'r', encoding='utf-8') as f: + class_names = f.readlines() + class_names = [class_name.strip() for class_name in class_names] + + detector = YOLOv5Detector(model_path=args.model_path, gpu=args.gpu) + + image = cv2.imread(args.image) + scores, class_ids, boxes = detector(image) + + visualize_detection(image, class_names, scores, class_ids, boxes) + + cv2.imshow('result', image) + cv2.waitKey(0) From 8424a2603fe2f8a5df7eb568e58f8105370b0f5a Mon Sep 17 00:00:00 2001 From: itsnine Date: Mon, 27 Sep 2021 23:35:49 +0500 Subject: [PATCH 2/4] updated boxes format to xywh --- deployment/onnxruntime-python/example.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/deployment/onnxruntime-python/example.py b/deployment/onnxruntime-python/example.py index 9ed76b16..2fd0b707 100644 --- a/deployment/onnxruntime-python/example.py +++ b/deployment/onnxruntime-python/example.py @@ -36,6 +36,8 @@ def __call__(self, image: np.ndarray) -> Tuple[List[float], List[int], List[List scores, class_ids, boxes = self._model.run(output_names=None, input_feed={self._input_names: blob}) boxes = boxes.astype(np.int32) + boxes[:, [2, 3]] -= boxes[:, [0, 1]] # from xyxy to xywh format + return scores.tolist(), class_ids.tolist(), boxes.tolist() @@ -46,15 +48,15 @@ def visualize_detection(image: np.ndarray, boxes: List[List[int]]) -> None: for i, class_id in enumerate(class_ids): - x1, y1, x2, y2 = boxes[i] + x, y, w, h = boxes[i] conf = round(scores[i], 2) label = class_names[class_id] + " " + str(conf) text_size = cv2.getTextSize(label, cv2.FONT_ITALIC, 0.8, 2)[0] - cv2.rectangle(image, (x1, y1), (x2, y2), (229, 160, 21), 2) - cv2.rectangle(image, (x1, y1 - 25), (x1 + text_size[0], y1), (229, 160, 21), -1) - cv2.putText(image, label, (x1, y1 - 3), cv2.FONT_ITALIC, 0.8, (255, 255, 255), 2) + cv2.rectangle(image, (x, y), (x + w, y + h), (229, 160, 21), 2) + cv2.rectangle(image, (x, y - 25), (x + text_size[0], y), (229, 160, 21), -1) + cv2.putText(image, label, (x, y - 3), cv2.FONT_ITALIC, 0.8, (255, 255, 255), 2) def get_parser() -> argparse.ArgumentParser: From 622d7f0de6a58dbb542ff2a8fc7498f717e44025 Mon Sep 17 00:00:00 2001 From: itsnine Date: Mon, 27 Sep 2021 23:49:06 +0500 Subject: [PATCH 3/4] up onnxruntime-python readme --- deployment/onnxruntime-python/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deployment/onnxruntime-python/README.md b/deployment/onnxruntime-python/README.md index f4d11cfb..269798da 100644 --- a/deployment/onnxruntime-python/README.md +++ b/deployment/onnxruntime-python/README.md @@ -19,7 +19,7 @@ The ONNXRuntime inference for `yolort`, both GPU and CPU are supported. ```bash python example.py [--image ../../../test/assets/zidane.jpg] - [--model_path ../../../test/tracing/yolov5s.onnx] - [--class_names ../../../notebooks/assets/coco.names] - [--gpu] # GPU switch, which is optional, and set False as default + [--model_path ../../../notebooks/yolov5s.onnx] + [--class_names ../../../notebooks/assets/coco.names] + [--gpu] # GPU switch, which is optional, and set False as default ``` \ No newline at end of file From c12b3db67982a32371c79977f96349c152da0c2b Mon Sep 17 00:00:00 2001 From: itsnine Date: Mon, 27 Sep 2021 23:50:30 +0500 Subject: [PATCH 4/4] up onnxruntime-python readme --- deployment/onnxruntime-python/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deployment/onnxruntime-python/README.md b/deployment/onnxruntime-python/README.md index 269798da..73af7b1e 100644 --- a/deployment/onnxruntime-python/README.md +++ b/deployment/onnxruntime-python/README.md @@ -19,7 +19,7 @@ The ONNXRuntime inference for `yolort`, both GPU and CPU are supported. ```bash python example.py [--image ../../../test/assets/zidane.jpg] - [--model_path ../../../notebooks/yolov5s.onnx] - [--class_names ../../../notebooks/assets/coco.names] - [--gpu] # GPU switch, which is optional, and set False as default + [--model_path ../../../notebooks/yolov5s.onnx] + [--class_names ../../../notebooks/assets/coco.names] + [--gpu] # GPU switch, which is optional, and set False as default ``` \ No newline at end of file