onnx快速部署YOLO模型

news/2024/7/10 3:08:34 标签: YOLO, python, 目标检测, 深度学习

1、准备和环境

 首先需要将yolov5模型训练好的最佳权重文件转化为.onnx格式以备使用。不会的小伙伴可以参考yolov5的官方文档,使用yolov5官方的 export.py 脚本进行转换,或者参考一些博客链接,这里不做详细解析。
 基本环境配置,相比于yolov5模型源码部署,使用onnx的方式部署会省下不少配置环境的问题,只需要几个关键的第三方库即可完成。

numpy>=1.22.3
onnxruntime>=1.13.1
Pillow>=9.3.0
python-multipart>=0.0.5
fastapi>=0.88.0
python-multipart>=0.0.5
uvicorn[standard]

2、部署代码

2.1、main.py

 主文件主要调用

python">from PIL import Image,ImageDraw,ImageFont
from utils.operation import YOLO
from utils.Colors import colors
import numpy as np
from matplotlib import pyplot as plt
def draw_anchor(img,det_obj):
    img = Image.open(img)
    draw = ImageDraw.Draw(img)
    font = ImageFont.truetype('arial.ttf', 30)
    imgw,imgh = img.size
    colors_dt = dict() 
    for i in range(len(det_obj)):
        if colors_dt.get(det_obj[i]['classes']) is None:
            colors_dt[det_obj[i]['classes']] = colors(i,True)
        draw.rectangle(det_obj[i]['crop'],width=3,outline=colors_dt[det_obj[i]['classes']])
        x1, y1, x2, y2 = tuple(det_obj[i]['crop'])
        draw.text((x1,y1-35),det_obj[i]['classes'],fill=colors_dt[det_obj[i]['classes']],font=font)
    imgarr = np.array(img)
    plt.imshow(imgarr)
    plt.show()
    img.show() 
def detect(onnx_path='ReqFile/yolov5s.onnx',img=r'ReqFile/bus.jpg',show=True):
    '''
    检测目标,返回目标所在坐标如:
    {'crop': [57, 390, 207, 882], 'classes': 'person'},...]
    :param onnx_path:onnx模型路径
    :param img:检测用的图片
    :param show:是否展示
    '''
    #加载yolo
    yolo = YOLO(onnx_path=onnx_path)  # 加载yolo类
    #检测
    det_obj = yolo.decect(img,conf_thres=0.5, iou_thres=0.25)  # 检测
    #画锚框
    draw_anchor(img,det_obj)
if __name__ == '__main__':
    detect()
    pass

2.2、operation.py

python">from io import BytesIO

import onnxruntime
import numpy as np
from PIL import Image

from utils.orientation import non_max_suppression, tag_images


class ONNXModel(object):
    def __init__(self, onnx_path):
        """
        :param onnx_path:
        """
        self.onnx_session = onnxruntime.InferenceSession(onnx_path)
        self.input_name = self.get_input_name(self.onnx_session)
        self.output_name = self.get_output_name(self.onnx_session)
    def get_output_name(self, onnx_session):
        """
        output_name = onnx_session.get_outputs()[0].name
        :param onnx_session:
        :return:
        """
        output_name = []
        for node in onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def get_input_name(self, onnx_session):
        """
        input_name = onnx_session.get_inputs()[0].name
        :param onnx_session:
        :return:
        """
        input_name = []
        for node in onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_input_feed(self, input_name, image_numpy):
        """
        input_feed={self.input_name: image_numpy}
        :param input_name:
        :param image_numpy:
        :return:
        """
        input_feed = {}
        for name in input_name:
            input_feed[name] = image_numpy
        return input_feed

    def to_numpy(self, file, shape, gray=False):
        if isinstance(file, np.ndarray):
            img = Image.fromarray(file)
        elif isinstance(file, bytes):
            img = Image.open(BytesIO(file))
            pass
        else:
            img = Image.open(file)
        widht, hight = shape
         # 改变大小 并保证其不失真
        img = img.convert('RGB')
        if gray:
            img = img.convert('L')
        img = img.resize((widht, hight), Image.ANTIALIAS)
        # 转换成矩阵
        image_numpy = np.array(img) # (widht, hight, 3)
        if gray:
            image_numpy = np.expand_dims(image_numpy,0)
            image_numpy = image_numpy.transpose(0, 1, 2)
        else:
            image_numpy = image_numpy.transpose(2,0,1) # 转置 (3, widht, hight)
        image_numpy = np.expand_dims(image_numpy,0)
        # 数据归一化
        image_numpy = image_numpy.astype(np.float32) / 255.0
        return image_numpy
class YOLO(ONNXModel):
    def __init__(self, onnx_path="ReqFile/yolov5n-7-k5.onnx"):
        super(YOLO, self).__init__(onnx_path)
        # 训练所采用的输入图片大小
        self.img_size = 640
        self.img_size_h = self.img_size_w = self.img_size
        self.batch_size = 1
        #数量
        self.num_classes = 2
        #标签
        self.classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']
    def to_numpy(self, file, shape, gray=False):
        def letterbox_image(image, size):
            iw, ih = image.size
            w, h = size
            scale = min(w / iw, h / ih)
            nw = int(iw * scale)
            nh = int(ih * scale)
            image = image.resize((nw, nh), Image.BICUBIC)
            new_image = Image.new('RGB', size, (128, 128, 128))
            new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
            return new_image
        if isinstance(file, np.ndarray):
            img = Image.fromarray(file)
        elif isinstance(file, bytes):
            img = Image.open(BytesIO(file))
        else:
            img = Image.open(file)
        resized = letterbox_image(img, (self.img_size_w, self.img_size_h))
        img_in = np.transpose(resized, (2, 0, 1)).astype(np.float32)  # HWC -> CHW
        img_in = np.expand_dims(img_in, axis=0)
        img_in /= 255.0
        return img_in
    def decect(self, file,conf_thres=0.25, iou_thres=0.45):
        # 图片转换为矩阵
        image_numpy = self.to_numpy(file, shape=(self.img_size, self.img_size))
        input_feed = self.get_input_feed(self.input_name, image_numpy)
        outputs = self.onnx_session.run(self.output_name, input_feed=input_feed)
        pred = non_max_suppression(outputs[0],conf_thres, iou_thres)
        if pred:
            res = tag_images(np.array(Image.open(file)), pred, self.img_size, self.classes)
        else:
            res = []
        return res

2.3、orientation.py

python">import time
import numpy as np
#用于控制Python中小数的显示精度
np.set_printoptions(precision=4)
def rescale_boxes(boxes, current_dim, original_shape):
    """ Rescales bounding boxes to the original shape """
    orig_h, orig_w = original_shape
    # The amount of padding that was added
    pad_x = max(orig_h - orig_w, 0) * (current_dim / max(original_shape))
    pad_y = max(orig_w - orig_h, 0) * (current_dim / max(original_shape))
    # Image height and width after padding is removed
    unpad_h = current_dim - pad_y
    unpad_w = current_dim - pad_x
    # Rescale bounding boxes to dimension of original image
    boxes[:, 0] = ((boxes[:, 0] - pad_x // 2) / unpad_w) * orig_w
    boxes[:, 1] = ((boxes[:, 1] - pad_y // 2) / unpad_h) * orig_h
    boxes[:, 2] = ((boxes[:, 2] - pad_x // 2) / unpad_w) * orig_w
    boxes[:, 3] = ((boxes[:, 3] - pad_y // 2) / unpad_h) * orig_h
    return boxes
def tag_images(imgs, img_detections, img_size, classes):
    imgs = [imgs]
    results = []
    for img_i, (img, detections) in enumerate(zip(imgs, img_detections)):
        # Create plot
        if detections is not None:
            # Rescale boxes to original image
            detections = rescale_boxes(detections, img_size, img.shape[:2])
            for x1, y1, x2, y2, conf, cls_pred in detections:
                results.append(
                    {
                        "crop": [int(i) for i in (x1, y1, x2, y2)],
                        "classes": classes[int(cls_pred)]
                    }
                )
        else:
            print("识别失败")
    return results
# 识别结果解析
def xywh2xyxy(x):
    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] 
    # where xy1=top-left, xy2=bottom-right
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y
def nms(dets, scores, thresh):
    """Pure Python NMS baseline."""
    # x1、y1、x2、y2、以及score赋值
    x1 = dets[:, 0]  # xmin
    y1 = dets[:, 1]  # ymin
    x2 = dets[:, 2]  # xmax
    y2 = dets[:, 3]  # ymax
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    # argsort()返回数组值从小到大的索引值
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:  # 还有数据
        i = order[0]
        keep.append(i)
        if order.size == 1: break
        # 计算当前概率最大矩形框与其他矩形框的相交框的坐标
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        # 计算相交框的面积
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        # 计算重叠度IOU:重叠面积/(面积1+面积2-重叠面积)
        IOU = inter / (areas[i] + areas[order[1:]] - inter)
        left_index = (np.where(IOU <= thresh))[0]
        # 将order序列更新,由于前面得到的矩形框索引要比矩形框在原order序列中的索引小1,所以要把这个1加回来
        order = order[left_index + 1]
    return np.array(keep)

def non_max_suppression(prediction, conf_thres, iou_thres, classes=None, agnostic=False, multi_label=False,
                        labels=()):
    """Runs Non-Maximum Suppression (NMS) on inference results

        Returns:
             list of detections, on (n,6) tensor per image [xyxy, conf, cls]
        """
    nc = prediction.shape[2] - 5
    xc = prediction[..., 4] > conf_thres  # candidates
    # Settings
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
    max_det = 300  # maximum number of detections per image
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    t = time.time()
    output = [np.zeros((0, 6))] * prediction.shape[0]
    for xi, x in enumerate(prediction):
        x = x[xc[xi]]  # confidence
        # Cat apriori labels if autolabelling
        if labels and len(labels[xi]):
            l = labels[xi]
            v = np.zeros((len(l), nc + 5))
            v[:, :4] = l[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
            x = np.concatenate((x, v), 0)
        # If none remain process next image
        if not x.shape[0]:
            continue
        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero()
            x = np.concatenate((box[i], x[i, j + 5, None], j[:, None]), 1)
        else:  # best class only
            conf = x[:, 5:].max(1, keepdims=True)
            j = x[:, 5:].argmax(1)
            j = np.expand_dims(j, 0).T
            x = np.concatenate((box, conf, j), 1)[conf.reshape(1, -1)[0] > conf_thres]
        # Filter by class
        if classes is not None:
            x = x[(x[:, 5:6] == np.array(classes)).any(1)]
        # Check shape
        n = x.shape[0]  # number of boxes
        if not n:  # no boxes
            continue
        elif n > max_nms:  # excess boxes
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence
        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = nms(boxes, scores, iou_thres)  # NMS
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
            break  # time limit exceeded
        return output

3、测试

在这里插入图片描述

完整文件请参考:
https://pan.baidu.com/s/1X1DiywM8yJtBzysKfzqbcA
提取码:6666


http://www.niftyadmin.cn/n/5226285.html

相关文章

JS快速入门

一、JS介绍 JavaScript &#xff08;简称JS&#xff09;&#xff0c;是一门跨平台、面向对象的脚本语言&#xff08;弱类型语言&#xff09;&#xff0c;而Java语言也是跨平台的、面向对象的语言&#xff0c;只不过Java是编译语言&#xff0c;是需要编译成字节码文件才能运行的…

CRM系统的数据分析和报表功能对企业重要吗?

竞争日益激烈&#xff0c;企业需要更加高效地管理客户关系&#xff0c;以获取更多的商机。为此&#xff0c;许多企业选择使用CRM系统。在CRM中&#xff0c;数据分析功能扮演着重要的角色。下面就来详细说说&#xff0c;CRM系统数据分析与报表功能对企业来说重要吗&#xff1f; …

福州大学《嵌入式系统综合设计》 实验十一:OpenCV视频解码

一、实验目的 掌握OpenCV视频解码的流程&#xff0c;对比和FFMPEG的区别。 二、实验内容 搭建实验开发环境&#xff0c;编译并运行解码程序&#xff0c;对编码后的视频码流通过OpenCV进行解码。 三、开发环境 开发主机&#xff1a;Ubuntu 20.04.6 LTS 硬件&#xff1a;算…

安网AC智能路由系统actpt_5g.data敏感信息泄露漏洞复现 [附POC]

文章目录 安网AC智能路由系统actpt_5g.data敏感信息泄露漏洞复现 [附POC]0x01 前言0x02 漏洞描述0x03 影响版本0x04 漏洞环境0x05 漏洞复现1.访问漏洞环境2.构造POC3.复现 安网AC智能路由系统actpt_5g.data敏感信息泄露漏洞复现 [附POC] 0x01 前言 免责声明&#xff1a;请勿利…

CAN总线学习(STM32的CAN寄存器使用)(笔记二)

CAN总线基础基础知识的文章&#xff1a;CAN总线学习&#xff08;CAN总线基础知识&#xff09;&#xff08;笔记一&#xff09;-CSDN博客 在前面CAN总线基础知识和报文中介绍的是报文内容是比较全面的&#xff0c;STM32在CAN协议的基础上做了一些简单的简化&#xff0c;例如下图…

BatchOutput PDF for Mac(PDF 批量处理软件)

BatchOutput PDF是一款适用于 Mac 的 PDF 批量处理软件。它可以帮助用户将多个 PDF 文件进行异步处理&#xff0c;提高工作效率。 BatchOutput PDF 可以自动化执行许多任务&#xff0c;包括 PDF 文件的打印、转换、分割、压缩、加密、重命名等&#xff0c;而且它还可以将自定义…

layui提示框没有渲染bug解决

bug&#xff1a;使用layui时或许是依赖导入又或是ideal和浏览器缓存问题导致前面明明正常的页面显示&#xff0c;后面出现提示框没有css样式&#xff0c;弹出框没有背景css 效果如下 解决后 解决方法 在你的代码中引入layer.js 我这是jsp页面 <script type"text/jav…

Redis 实战缓存

本篇概要&#xff1a; 1. 设置、查询、获取过期时间&#xff1b;2. 缓存穿透&#xff1a;设置空键&#xff1b;3. 封杀单ip&#xff1b;4. 封杀ip段&#xff1b;5. 缓存预热&#xff1b;6. 使用 hash 数据类型保存新闻的缓存&#xff0c;增加点击量&#xff1b;7. Sorted set&a…