From 225ecd6e741bf75e3d1b06e114de3d5c1290b591 Mon Sep 17 00:00:00 2001 From: semihhdemirel Date: Fri, 8 Nov 2024 15:02:32 +0300 Subject: [PATCH] [Example] YOLOv8-Classification-ONNXRuntime-Python --- examples/README.md | 1 + .../README.md | 43 ++++++ .../yolov8_classifier.py | 130 ++++++++++++++++++ 3 files changed, 174 insertions(+) create mode 100644 examples/YOLOv8-Classification-ONNXRuntime-Python/README.md create mode 100644 examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py diff --git a/examples/README.md b/examples/README.md index 76f078bde2..b4fb63a3c3 100644 --- a/examples/README.md +++ b/examples/README.md @@ -12,6 +12,7 @@ This directory features a collection of real-world applications and walkthroughs | [YOLO .Net ONNX Detection C#](https://www.nuget.org/packages/Yolov8.Net) | C# .Net | [Samuel Stainback](https://github.com/sstainba) | | [YOLOv8 on NVIDIA Jetson(TensorRT and DeepStream)](https://wiki.seeedstudio.com/YOLOv8-DeepStream-TRT-Jetson/) | Python | [Lakshantha](https://github.com/lakshanthad) | | [YOLOv8 ONNXRuntime Python](./YOLOv8-ONNXRuntime) | Python/ONNXRuntime | [Semih Demirel](https://github.com/semihhdemirel) | +| [YOLOv8 Classification ONNXRuntime Python](./YOLOv8-Classification-ONNXRuntime-Python) | Python/ONNXRuntime | [Semih Demirel](https://github.com/semihhdemirel) | | [YOLOv8 ONNXRuntime CPP](./YOLOv8-ONNXRuntime-CPP) | C++/ONNXRuntime | [DennisJcy](https://github.com/DennisJcy), [Onuralp Sezer](https://github.com/onuralpszr) | | [RTDETR ONNXRuntime C#](https://github.com/Kayzwer/yolo-cs/blob/master/RTDETR.cs) | C#/ONNX | [Kayzwer](https://github.com/Kayzwer) | | [YOLOv8 SAHI Video Inference](https://github.com/RizwanMunawar/ultralytics/blob/main/examples/YOLOv8-SAHI-Inference-Video/yolov8_sahi.py) | Python | [Muhammad Rizwan Munawar](https://github.com/RizwanMunawar) | diff --git a/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md b/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md new file mode 100644 index 0000000000..7851df315f --- /dev/null +++ b/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md @@ -0,0 +1,43 @@ +# YOLOv8 - Classification with ONNX Runtime + +This project implements classification using YOLOv8 and ONNX Runtime. + +## Installation + +To run this project, you'll need to install the necessary dependencies. Follow the instructions below to set up your environment. + +### Installing Required Dependencies + +Install the required dependencies with this command: + +```bash +pip install -r requirements.txt +``` + +### Installing `onnxruntime-gpu` + +If you have an NVIDIA GPU and want to leverage GPU acceleration, you can install the onnxruntime-gpu package using the following command: + +```bash +pip install onnxruntime-gpu +``` + +Note: Make sure you have the appropriate GPU drivers installed on your system. + +### Installing `onnxruntime` (CPU version) + +If you don't have an NVIDIA GPU or prefer to use the CPU version of onnxruntime, you can install the onnxruntime package using the following command: + +```bash +pip install onnxruntime +``` + +### Usage + +After installing the required packages, you can run YOLOv8 classification with the following command: + +```bash +python yolov8_classifier.py --model yolov8n.onnx --img image.jpg +``` + +Replace yolov8n.onnx with the path to your YOLOv8 ONNX model file and image.jpg with the path to your input image. \ No newline at end of file diff --git a/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py b/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py new file mode 100644 index 0000000000..532b659756 --- /dev/null +++ b/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py @@ -0,0 +1,130 @@ +import argparse +import cv2 +import numpy as np +import onnxruntime as ort +import torch + +from ultralytics.utils import ASSETS, yaml_load +from ultralytics.utils.checks import check_requirements, check_yaml + + +class YOLOv8Classifier: + """YOLOv8 model for image classification and visualization.""" + + def __init__(self, model_path, image_path): + """ + Initializes the YOLOv8Classifier instance for classification and visualization. + + Args: + model_path (str): Path to the ONNX model file. + image_path (str): Path to the input image file. + """ + self.image_path = image_path + self.session = ort.InferenceSession(model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + self.model_input = self.session.get_inputs() + self.input_width = self.model_input[0].shape[2] + self.input_height = self.model_input[0].shape[3] + + # Load class names from the ImageNet dataset + self.class_names = yaml_load(check_yaml("ImageNet.yaml"))["names"] + + def annotate_class_label(self, class_index): + """ + Annotates the class label on the image. + + Args: + class_index (int): Index of the detected class to annotate. + """ + + # Draw the class label text on the image + label = f"{self.class_names[class_index]}" + cv2.putText(self.image, label, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) + + def preprocess_image(self): + """ + Preprocesses the input image for model inference. + + Returns: + np.array: Preprocessed image data ready for inference. + """ + # Read the image from the given path + self.image = cv2.imread(self.image_path) + + # Get the height and width of the image + self.image_height, self.image_width = self.image.shape[:2] + + # Convert the image from BGR to RGB color space + image_rgb = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB) + + # Resize the image to the input dimensions + image_resized = cv2.resize(image_rgb, (self.input_width, self.input_height)) + + # Normalize the image to the range [0, 1] + image_normalized = np.array(image_resized) / 255.0 + + # Transpose the image to channel-first format + image_transposed = np.transpose(image_normalized, (2, 0, 1)) + + # Expand dimensions to add batch size and convert to float32 + image_input = np.expand_dims(image_transposed, axis=0).astype(np.float32) + + # Return the preprocessed image input + return image_input + + def process_output(self, model_output): + """ + Processes the model's output to annotate the class label on the image. + + Args: + model_output (np.array): Raw output from the model. + + Returns: + np.array: Image with annotated class label. + """ + # Generate binary output based on the model's score threshold + out = [1 if score > 0.5 else 0 for score in model_output[0][0]] + + # Find the index of the detected class + detected_class_index = out.index(1) + + # Annotate the image with the detected class label + self.annotate_class_label(detected_class_index) + + # Return the annotated image + return self.image + + def classify_image(self): + """ + Executes the classification on the input image using the ONNX model. + + Returns: + np.array: Output image with annotations. + """ + # Preprocess the image for model input + image_data = self.preprocess_image() + + # Run the model inference + model_output = self.session.run(None, {self.model_input[0].name: image_data}) + + # Process and return the model output + return self.process_output(model_output) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, default="yolo8n-cls.onnx", help="Path to the ONNX model file.") + parser.add_argument("--img", type=str, default="image.jpg", help="Path to the input image.") + args = parser.parse_args() + + # Check for dependencies and set up ONNX runtime + check_requirements("onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime") + + # Create the classifier instance with specified parameters + classifier = YOLOv8Classifier(args.model, args.img) + + # Perform classification and get the output image + annotated_image = classifier.classify_image() + + # Display the annotated output image + cv2.namedWindow("Annotated Output", cv2.WINDOW_NORMAL) + cv2.imshow("Annotated Output", annotated_image) + cv2.waitKey(0)