Inferring YOLOv5-CLS converted ONNX classification models using OpenCV DNN

YOLOv5 is an advanced target detection algorithm, and YOLOv5-CLS is a variant of YOLOv5, which is specially used for image classification tasks. In order to use the YOLOv5-CLS model in practical applications, we need to convert it to the Open Neural Network Exchange (ONNX) format and use the OpenCV DNN library for inference.

Step 1: Install OpenCV and ONNX First, you need to make sure you have installed OpenCV and ONNX. It can be installed with the following command:

pip install opencv-python
pip install onnx

Step 2: Convert YOLOv5-CLS to ONNX format In this step, we will use YOLOv5's official codebase to convert the YOLOv5-CLS model to ONNX format. Follow the steps below:

1. Clone the official code base of YOLOv5:

   git clone https://github.com/ultralytics/yolov5.git
   

2. Enter the yolov5 directory and download the pre-trained YOLOv5-CLS model weights:

   cd yolov5
   wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s6.pt
   

3. Run the export.py script to convert the model to ONNX format:

   python export.py --weights yolov5s6.pt --include onnx --img 640
   

   This step will generate a file named 'yolov5s6.onnx', this is the ONNX version of the YOLOv5-CLS model we are going to use.

   Step 3: Inference with OpenCV DNN Now, we are ready for inference. Below is a simple sample code showing how to load and run the YOLOv5-CLS model using the OpenCV DNN library:

   import cv2
   
   # 加载YOLOv5-CLS模型
   net = cv2.dnn.readNetFromONNX("yolov5s6.onnx")
   
   # 加载图像
   image = cv2.imread("test.jpg")
   
   # 创建blob并设置输入
   blob = cv2.dnn.blobFromImage(image, 1/255., (640, 640), swapRB=True)
   net.setInput(blob)
   
   # 进行推理
   output = net.forward()
   
   # 解析推理结果
   classes = open("coco.names").read().strip().split("\n")
   for detection in output[0, 0]:
       scores = detection[5:]
       class_id = np.argmax(scores)
       confidence = scores[class_id]
       if confidence > 0.5:
           label = f"{classes[class_id]}: {confidence:.2f}"
           print(label)
   
   # 显示图像
   cv2.imshow("Image", image)
   cv2.waitKey(0)
   cv2.destroyAllWindows()
   

   Note that in the above code we assume there is already a test image called 'test.jpg' and a 'coco.names' file containing the class names.

   Conclusion: This article describes how to use the OpenCV DNN library for inference of the YOLOv5-CLS model. We first convert the YOLOv5-CLS model to ONNX format, then use the OpenCV DNN library to load and run the model for image classification. By following the steps and sample code provided in this article, you can easily use the YOLOv5-CLS model for image classification tasks in practical applications.