yolov7-onnxruntime

yolov7-onnxruntime调用示例

参考文章

https://cloud.tencent.com/developer/article/2205367

https://github.com/hpc203/yolov7-opencv-onnxrun-cpp-py

官方文档 https://onnxruntime.ai/docs/api/c/index.html

python示例

这里测试可正常推理自定义模型，模型导出参考文章yolov7-opencv

import cv2
import numpy as np
import onnxruntime
import argparse


class YOLOv7:
    def __init__(self, path, conf_thres=0.7, iou_thres=0.5):
        self.conf_threshold = conf_thres
        self.iou_threshold = iou_thres
        #self.class_names = list(map(lambda x: x.strip(), open('coco.names', 'r').readlines()))
        self.class_names =["ceshi1","ceshi2","ceshi3"]
        # Initialize model
        self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        model_inputs = self.session.get_inputs()
        self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
        self.input_shape = model_inputs[0].shape
        self.input_height = self.input_shape[2]
        self.input_width = self.input_shape[3]
        
        model_outputs = self.session.get_outputs()
        self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
        self.has_postprocess = 'score' in self.output_names
    
    def detect(self, image):
        input_tensor = self.prepare_input(image)
        
        # Perform inference on the image
        outputs = self.session.run(self.output_names, {
   
   self.input_names[0]: input_tensor})
        
        if self.has_postprocess:
            boxes, scores, class_ids = self.parse_processed_output(outputs)
        
        else:
            # Process output data
            boxes, scores, class_ids = self.process_output(outputs)
        
        return boxes, scores, class_ids
    
    def prepare_input(self, image):
        self.img_height, self.img_width = image.shape[:2]
        
        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Resize input image
        input_img = cv2.resize(input_img, (self.input_width, self.input_height))
        
        # Scale input pixel values to 0 to 1
        input_img = input_img / 255.0
        input_i