YOLOV8-ONNX接口预测（第五届计挑赛尝试onnx）

Hoo K

已于 2023-12-26 21:30:29 修改

阅读量342

点赞数 1

文章标签： YOLO opencv 人工智能

于 2023-12-26 20:34:54 首次发布

本文链接：https://blog.youkuaiyun.com/qq_73668945/article/details/135230828

版权

一、文件目录

附加【pt转onnx】

# Pytorch模型转换为Onnx模型
import onnx
from onnxruntime.quantization import quantize_dynamic,QuantType
from onnxconverter_common import float16


#yolov8原生转换
# from ultralytics import YOLO
# model = YOLO(r'C:\YOLOv8-Openvino-master\best.pt')
# result = model.export(format='onnx') 

#fp16
# model = onnx.load(r"C:\Users\Acer\Desktop\new\best.onnx")
# model_fp16 = float16.convert_float_to_float16(model)
# onnx.save(model_fp16,r"C:\Users\Acer\Desktop\new\onnx_fp16\best_fp16.onnx")



# ******************************************
# model_fp32 = r"C:\Users\Acer\Desktop\new\best.onnx"
# model_quant = r"C:\Users\Acer\Desktop\new\onnx_fp16\best_int8.onnx"
# quantized_model = quantize_dynamic(model_fp32,model_quant)
# **********************************************

def quantize_onnx_model(onnx_model_path, quantized_model_path):
    from onnxruntime.quantization import quantize_dynamic, QuantType
    import onnx
    onnx_opt_model = onnx.load(onnx_model_path)
    quantize_dynamic(onnx_model_path,quantized_model_path,weight_type=QuantType.QUInt8)

二、各个目录下代码

其中①运行即可出现[y1,x1,y2,x2,conf,class]的预测结果

-- ①model_predict.py

from yolov8.YOLOv8 import YOLOv8
import numpy as np
import cv2
class predict:
    def __init__(self, **kwargs):
        self.yolov8_detector = YOLOv8(r"./best_new_int8.onnx", conf_thres=0.6, iou_thres=0.4)
    def detect_image(self, image_path):
        img = cv2.imread(image_path)
        boxes, scores, class_ids = self.yolov8_detector(img)
        results = []
        for ii, i in enumerate(boxes):
            ls = i.tolist()
            ls.append(scores[ii])
            ls.append(class_ids[ii])
            results.append(ls)
        return np.array(results)

if __name__ == '__main__':
    image_path = r"C:\Users\Acer\Desktop\data_d\train\images\283.jpg"
    s = predict()
    ls = s.detect_image(image_path)
    print(ls)
# 下面是使用cv检验权重模型的准确性
    # img = cv2.imread(fill)
    # for i in ls:
    #     y1 = int(i[0])
    #     x1 = int(i[1])
    #     y2 = int(i[2])
    #     x2 = int(i[3])
        # cv2.rectangle(img, (x1, y1), (x2, y2), (255, 255, 0), 3)
    # cv2.imshow('img', img)
    # cv2.waitKeyEx(0)
    # cv2.destroyAllWindows()

②-- YOLOV8.py

这个py文件引用了utils.py文件里面的NMS、和坐标转换、draw_detections(画框)

坐标转换：将中心点坐标转换为左上角右下角坐标

NMS：非极大值抑制

—— 1.计算框的面积

——2.计算相交面积（相交、不相交）

—— 3.计算该框与其它框的IOU，去除掉重复的框，即IOU值大的框

—— 4.IOU小于thresh的框保留下来

import time
import cv2
import numpy as np
import onnxruntime

from yolov8.utils import xywh2xyxy,  multiclass_nms  # draw_detections,

class YOLOv8:

    def __init__(self, path, conf_thres=0.7, iou_thres=0.5):
        self.conf_threshold = conf_thres
        self.iou_threshold = iou_thres

        # Initialize model
        self.initialize_model(path)

    def __call__(self, image):
        return self.detect_objects(image)

    def initialize_model(self, path):
        self.session = onnxruntime.InferenceSession(path,
                                                    providers=onnxruntime.get_available_providers())
        # Get model info
        self.get_input_details()
        self.get_output_details()


    def detect_objects(self, image):
        input_tensor = self.prepare_input(image)
        #
        # input_tensor= input_tensor.astype(np.float16)

        # Perform inference on the image
        outputs = self.inference(input_tensor)

        self.boxes, self.scores, self.class_ids = self.process_output(outputs)

        return self.boxes, self.scores, self.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_img = input_img.transpose(2, 0, 1)
        input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)

        return input_tensor


    def inference(self, input_tensor):
        start = time.perf_counter()
        outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor})

        # print(f"Inference time: {(time.perf_counter() - start)*1000:.2f} ms")
        return outputs

    def process_output(self, output):
        predictions = np.squeeze(output[0]).T

        # Filter out object confidence scores below threshold
        scores = np.max(predictions[:, 4:], axis=1)
        predictions = predictions[scores > self.conf_threshold, :]
        scores = scores[scores > self.conf_threshold]

        if len(scores) == 0:
            return [], [], []

        # Get the class with the highest confidence
        class_ids = np.argmax(predictions[:, 4:], axis=1)

        # Get bounding boxes for each object
        boxes = self.extract_boxes(predictions)

        # Apply non-maxima suppression to suppress weak, overlapping bounding boxes
        # indices = nms(boxes, scores, self.iou_threshold)
        indices = multiclass_nms(boxes, scores, class_ids, self.iou_threshold)

        return boxes[indices], scores[indices], class_ids[indices]

    def extract_boxes(self, predictions):
        # Extract boxes from predictions
        boxes = predictions[:, :4]

        # Scale boxes to original image dimensions
        boxes = self.rescale_boxes(boxes)

        # Convert boxes to xyxy format
        boxes = xywh2xyxy(boxes)

        return boxes

    def rescale_boxes(self, boxes):

        # Rescale boxes to original image dimensions
        input_shape = np.array([self.input_width, self.input_height, self.input_width, self.input_height])
        boxes = np.divide(boxes, input_shape, dtype=np.float32)
        boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
        return boxes

    # def draw_detections(self, image, draw_scores=True, mask_alpha=0.4):
    #
    #     return draw_detections(image, self.boxes, self.scores,
    #                            self.class_ids, mask_alpha)

    def get_input_details(self):
        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]

    def get_output_details(self):
        model_outputs = self.session.get_outputs()
        self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]


if __name__ == '__main__':
    from imread_from_url import imread_from_url

    model_path = "../models/yolov8m.onnx"

    # Initialize YOLOv8 object detector
    yolov8_detector = YOLOv8(model_path, conf_thres=0.3, iou_thres=0.5)

    img_url = "https://live.staticflickr.com/13/19041780_d6fd803de0_3k.jpg"
    img = imread_from_url(img_url)

    # Detect Objects
    yolov8_detector(img)

    # Draw detections
    combined_img = yolov8_detector.draw_detections(img)
    cv2.namedWindow("Output", cv2.WINDOW_NORMAL)
    cv2.imshow("Output", combined_img)
    cv2.waitKey(0)

③-- utils

import numpy as np
import cv2

class_names = ["ship","plane"]

# Create a list of colors for each class where each color is a tuple of 3 integer values
rng = np.random.default_rng(3)
colors = rng.uniform(0, 255, size=(len(class_names), 3))


def nms(boxes, scores, iou_threshold):
    # Sort by score
    sorted_indices = np.argsort(scores)[::-1]

    keep_boxes = []
    while sorted_indices.size > 0:
        # Pick the last box
        box_id = sorted_indices[0]
        keep_boxes.append(box_id)

        # Compute IoU of the picked box with the rest
        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])

        # Remove boxes with IoU over the threshold
        keep_indices = np.where(ious < iou_threshold)[0]

        # print(keep_indices.shape, sorted_indices.shape)
        sorted_indices = sorted_indices[keep_indices + 1]

    return keep_boxes

def multiclass_nms(boxes, scores, class_ids, iou_threshold):

    unique_class_ids = np.unique(class_ids)

    keep_boxes = []
    for class_id in unique_class_ids:
        class_indices = np.where(class_ids == class_id)[0]
        class_boxes = boxes[class_indices,:]
        class_scores = scores[class_indices]

        class_keep_boxes = nms(class_boxes, class_scores, iou_threshold)
        keep_boxes.extend(class_indices[class_keep_boxes])

    return keep_boxes

def compute_iou(box, boxes):
    # Compute xmin, ymin, xmax, ymax for both boxes
    xmin = np.maximum(box[0], boxes[:, 0])
    ymin = np.maximum(box[1], boxes[:, 1])
    xmax = np.minimum(box[2], boxes[:, 2])
    ymax = np.minimum(box[3], boxes[:, 3])

    # Compute intersection area
    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)

    # Compute union area
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union_area = box_area + boxes_area - intersection_area

    # Compute IoU
    iou = intersection_area / union_area

    return iou


def xywh2xyxy(x, decimals=2):
    # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
    y = np.copy(x)
    y[..., 1] = np.round(x[..., 0] - x[..., 2] / 2,decimals)
    y[..., 0] = np.round(x[..., 1] - x[..., 3] / 2,decimals)
    y[..., 3] = np.round(x[..., 0] + x[..., 2] / 2,decimals)
    y[..., 2] = np.round(x[..., 1] + x[..., 3] / 2,decimals)
    np.set_printoptions(suppress=True, precision=decimals)
    return y

#
# def draw_detections(image, boxes, scores, class_ids, mask_alpha=0.3):
#     det_img = image.copy()
#
#     img_height, img_width = image.shape[:2]
#     font_size = min([img_height, img_width]) * 0.0006
#     text_thickness = int(min([img_height, img_width]) * 0.001)
#
#     det_img = draw_masks(det_img, boxes, class_ids, mask_alpha)
#
#     # Draw bounding boxes and labels of detections
#     for class_id, box, score in zip(class_ids, boxes, scores):
#         color = colors[class_id]
#
#         draw_box(det_img, box, color)
#
#         label = class_names[class_id]
#         caption = f'{label} {int(score * 100)}%'
#         draw_text(det_img, caption, box, color, font_size, text_thickness)
#
#     return det_img
#
#
# def draw_box( image: np.ndarray, box: np.ndarray, color: tuple[int, int, int] = (0, 0, 255),
#              thickness: int = 2) -> np.ndarray:
#     x1, y1, x2, y2 = box.astype(int)
#     return cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
#
#
# def draw_text(image: np.ndarray, text: str, box: np.ndarray, color: tuple[int, int, int] = (0, 0, 255),
#               font_size: float = 0.001, text_thickness: int = 2) -> np.ndarray:
#     x1, y1, x2, y2 = box.astype(int)
#     (tw, th), _ = cv2.getTextSize(text=text, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
#                                   fontScale=font_size, thickness=text_thickness)
#     th = int(th * 1.2)
#
#     cv2.rectangle(image, (x1, y1),
#                   (x1 + tw, y1 - th), color, -1)
#
#     return cv2.putText(image, text, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, font_size, (255, 255, 255), text_thickness, cv2.LINE_AA)
#
# def draw_masks(image: np.ndarray, boxes: np.ndarray, classes: np.ndarray, mask_alpha: float = 0.3) -> np.ndarray:
#     mask_img = image.copy()
#
#     # Draw bounding boxes and labels of detections
#     for box, class_id in zip(boxes, classes):
#         color = colors[class_id]
#
#         x1, y1, x2, y2 = box.astype(int)
#
#         # Draw fill rectangle in mask image
#         cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
#
#     return cv2.addWeighted(mask_img, mask_alpha, image, 1 - mask_alpha, 0)

④-- __init__.py