diff --git a/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md b/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md index 7851df315f..3747ec0b39 100644 --- a/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md +++ b/examples/YOLOv8-Classification-ONNXRuntime-Python/README.md @@ -40,4 +40,4 @@ After installing the required packages, you can run YOLOv8 classification with t 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 +Replace yolov8n.onnx with the path to your YOLOv8 ONNX model file and image.jpg with the path to your input image. diff --git a/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py b/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py index 532b659756..1b0df02110 100644 --- a/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py +++ b/examples/YOLOv8-Classification-ONNXRuntime-Python/yolov8_classifier.py @@ -1,10 +1,11 @@ import argparse + import cv2 import numpy as np import onnxruntime as ort import torch -from ultralytics.utils import ASSETS, yaml_load +from ultralytics.utils import yaml_load from ultralytics.utils.checks import check_requirements, check_yaml @@ -35,7 +36,6 @@ class YOLOv8Classifier: 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) @@ -109,6 +109,7 @@ class YOLOv8Classifier: # 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.")