首页 > 编程语言 >YOLOV8 det 多batch TensorRT 推理（python ）

YOLOV8 det 多batch TensorRT 推理（python ）

时间：2024-09-18 17:51:48浏览次数：3  
标签：args python det bboxes list TensorRT device scores 0.0000

由于我当前的项目需求是推理四张，所以，demo部分也是基于4张进行演示的，不过基于此套路，可以实现 N C H W的任意尺度推理，推理代码如下：

import numpy as np
from numpy import ndarray
from typing import List, Tuple, Union

from models import TRTModule  # isort:skip
import argparse
from pathlib import Path

import cv2
import torch
from torch import Tensor

from config import CLASSES, COLORS
from models.torch_utils import det_postprocess
from models.utils import blob, letterbox, path_to_list


def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
    """
    Performs non-maximum suppression (NMS) on the boxes according
    to their intersection-over-union (IoU).

    NMS iteratively removes lower scoring boxes which have an
    IoU greater than iou_threshold with another (higher scoring)
    box.

    If multiple boxes have the exact same score and satisfy the IoU
    criterion with respect to a reference box, the selected box is
    not guaranteed to be the same between CPU and GPU. This is similar
    to the behavior of argsort in PyTorch when repeated values are present.

    Args:
        boxes (Tensor[N, 4])): boxes to perform NMS on. They
            are expected to be in ``(x1, y1, x2, y2)`` format with ``0 <= x1 < x2`` and
            ``0 <= y1 < y2``.
        scores (Tensor[N]): scores for each one of the boxes
        iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold

    Returns:
        Tensor: int64 tensor with the indices of the elements that have been kept
        by NMS, sorted in decreasing order of scores
    """
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)

def blob1(im: ndarray, return_seg: bool = False) -> Union[ndarray, Tuple]:
    seg = None
    if return_seg:
        seg = im.astype(np.float32) / 255
    im = im.transpose([2, 0, 1])
    # im = im[np.newaxis, ...]
    im = np.ascontiguousarray(im).astype(np.float32) / 255
    if return_seg:
        return im, seg
    else:
        return im


def det_postprocess1(data: Tuple[Tensor, Tensor, Tensor, Tensor], i:int):
    assert len(data) == 4
    iou_thres: float = 0.65
    num_dets, bboxes, scores, labels = data[0][i], data[1][i], data[2][
        i], data[3][i]
    nums = num_dets.item()
    if nums == 0:
        return bboxes.new_zeros((0, 4)), scores.new_zeros(
            (0, )), labels.new_zeros((0, ))
    # check score negative
    scores[scores < 0] = 1 + scores[scores < 0]
    # add nms
    idx = nms(bboxes, scores, iou_thres)
    bboxes, scores, labels = bboxes[idx], scores[idx], labels[idx]
    bboxes = bboxes[:nums]
    scores = scores[:nums]
    labels = labels[:nums]

    return bboxes, scores, labels


def main(args: argparse.Namespace) -> None:
    device = torch.device(args.device)
    Engine = TRTModule(args.engine, device)
    H, W = Engine.inp_info[0].shape[-2:]

    # set desired output names order
    Engine.set_desired(['num_dets', 'bboxes', 'scores', 'labels'])

    images = path_to_list(args.imgs)
    save_path = Path(args.out_dir)

    if not args.show and not save_path.exists():
        save_path.mkdir(parents=True, exist_ok=True)


    draw_list = []
    save_path_list = []
    img_list = []
    ratio_list = []
    dwdh_list = []
    for image in images:
        save_image = save_path / image.name
        bgr = cv2.imread(str(image))
        draw_list.append(bgr)
        save_path_list.append(save_image)

        bgr, ratio, dwdh = letterbox(bgr, (W, H))

        ratio_list.append(ratio)

        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        tensor = blob1(rgb, return_seg=False)
        dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device)
        dwdh_list.append(dwdh)
        tensor = torch.asarray(tensor, device=device)
        img_list.append(tensor)
    batch_tensor = torch.stack([img_list[0], img_list[1], img_list[2], img_list[3]], dim=0)
    print(batch_tensor.shape)
    data = Engine(batch_tensor)
    print(data)
    print(len(data))
    print(len(data[0]))

    for i in range(len(data[0])):
        bboxes, scores, labels = det_postprocess1(data, i)
        bboxes -= dwdh_list[i]
        bboxes /= ratio_list[i]

        for (bbox, score, label) in zip(bboxes, scores, labels):
            bbox = bbox.round().int().tolist()
            cls_id = int(label)
            cls = CLASSES[cls_id]
            color = COLORS[cls]
            cv2.rectangle(draw_list[i], bbox[:2], bbox[2:], color, 2)
            cv2.putText(draw_list[i],
                        f'{cls}:{score:.3f}', (bbox[0], bbox[1] - 2),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.75, [225, 255, 255],
                        thickness=2)
        if args.show:
            cv2.imshow('result', draw_list[i])
            cv2.waitKey(0)
        else:
            cv2.imwrite(str(save_path_list), draw_list[i])


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--engine', default="yolov8n_bach4.engine",type=str, help='Engine file')
    parser.add_argument('--imgs',default="data", type=str, help='Images file')
    parser.add_argument('--show', default=True,
                        action='store_true',
                        help='Show the detection results')
    parser.add_argument('--out-dir',
                        type=str,
                        default='./output',
                        help='Path to output file')
    parser.add_argument('--device',
                        type=str,
                        default='cuda:0',
                        help='TensorRT infer device')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    main(args)

模型输入尺寸：

torch.Size([4, 3, 640, 640])


模型输出：

(tensor([[5],
        [5],
        [3],
        [3]], device='cuda:0', dtype=torch.int32), tensor([[[477.7500, 224.0000, 560.0000, 521.5000],
         [211.2500, 241.5000, 283.5000, 507.0000],
         [109.8750, 235.8750, 224.6250, 536.0000],
         ...,
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000]],

        [[477.7500, 224.0000, 560.0000, 521.5000],
         [211.2500, 241.5000, 283.5000, 507.0000],
         [109.8750, 235.8750, 224.6250, 536.0000],
         ...,
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000]],

        [[373.5000, 160.6250, 572.0000, 494.7500],
         [ 59.7500, 238.1250, 555.0000, 495.2500],
         [218.5000, 358.5000, 261.7500, 497.0000],
         ...,
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000]],

        [[373.5000, 160.6250, 572.0000, 494.7500],
         [ 59.7500, 238.1250, 555.0000, 495.2500],
         [218.5000, 358.5000, 261.7500, 497.0000],
         ...,
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000],
         [  0.0000,   0.0000,   0.0000,   0.0000]]], device='cuda:0'), tensor([[0.9014, 0.8833, 0.8755, 0.8418, 0.4358, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000],
        [0.9014, 0.8833, 0.8755, 0.8418, 0.4358, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000],
        [0.8052, 0.7959, 0.3684, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000],
        [0.8052, 0.7959, 0.3684, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000]], device='cuda:0'), tensor([[ 0,  0,  0,  5,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
        [ 0,  0,  0,  5,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
        [ 0,  0, 27,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
        [ 0,  0, 27,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
          0,  0,  0,  0,  0,  0,  0,  0,  0,  0]], device='cuda:0',
       dtype=torch.int32))

batch推理后的结果中，对应的num_dets、bboxes、scores、labels输出维度变化如下：
num_dets：[1]   变为  [4, 1]

bboxes： [100,  4]      变为 [ 4 , 100, 4]

scores：  [100]      变为  [4 ,100]

labels：  [80]       变为  [4, 80]

推理结果图片如下：

后续抽空再补上TensorRT  C++的推理部分。


本完整项目链接如下：
https://download.csdn.net/download/weixin_59701401/89767509

标签：args,python,det,bboxes,list,TensorRT,device,scores,0.0000
From： https://blog.csdn.net/weixin_59701401/article/details/142337617

相关文章

  • centos7下安装Python3.7
    centos7默认安装的是python2.7，然而python2基本上要淘汰了，所以有必要安装最新的python3python,g++这些工具一般安装在/usr/bin目录里通过指令llpython*可以看到python指向的是python2.7我们要安装python3，使python指向python3下面开始具体步骤（参考其他大佬的方法，也是学了......
  • 使用Python计算多个集合的交集详解
    集合（Set）是Python中的一种常用数据结构，专门用于存储不重复的元素。在数据处理中，集合操作常被用来处理去重、并集、交集等问题。尤其在处理多个数据集时，交集操作尤为重要，因为它可以帮助我们找到多个集合中都存在的共同元素。本文将详细探讨如何在Python中高效地进行多个集合的交集操......
  • ffplay python 播放rtsp ffmpeg播放rtsp流
    ffmpeg播放RTSP的一点优化： AVDictionary参数配置。 https://www.ffmpeg.org/doxygen/trunk/libavformat_2options__table_8h-source.html 基于ffmpeg的播放器起播延迟优化 :probesize ONVIF、RTSP/RTP、FFMPEG的开发实录： 这里注意多线程问题，如果你用FFMPEG解多......
  • Python 中常见的数据结构（一）
    Python中常见的数据结构（一）Python是一种功能强大且灵活的编程语言，它提供了多种内置的数据结构，可以帮助我们更好地组织和处理数据。在这个文章中，我们将探讨Python中最常见的一些数据结构，并结合实例来演示它们的使用。1.字典（Dictionary）字典是一种键值对的数据结构，它允许我们根据......
  • Python 中常见的数据结构（二）
    Python中常见的数据结构（二）6.栈（Stack）栈是一种后进先出数据结构，Python中，可以使用list类型创建一个栈，例如：stack=[]stack.append('apple')stack.append('banana')print(stack.pop())#Output:banana在上面的示例中，我们创建了一个名为stack的栈，然后使用append方法添加......
  • Python 实现自动配置华为交换机
    Python实现自动配置华为交换机在网络运维中，配置交换机是非常重要的一步。如果我们可以使用Python来实现配置交换机，那么我们的工作效率将会大大提高。在本文中，我们将学习如何使用Python配置华为交换机。背景知识华为交换机是一种常用的网络设备，用于连接和转发数据包。为了配置......
  • Python 语法糖：让编程更简单（续三）
    Python语法糖：让编程更简单（续三）15.DictionarycomprehensionsDictionarycomprehensions是Python中的一种语法糖，用于简化字典的创建。例如：numbers=[1,2,3,4,5]squared_numbers_dict={x:x**2forxinnumbers}print(squared_numbers_dict)#prints{1:1,2:4,......
  • 让 Python 和 Web 世界合二为一
    Python操作浏览器：让Python和Web世界合二为一在日常开发中，我们经常需要与浏览器进行交互，以实现特定的任务，如爬取网页信息、自动化测试或执行某些操作。这时，Python的一些库和框架可以帮助我们轻松地操作浏览器。下面，我们将探讨Python操作浏览器的方法和示例。SeleniumSeleniu......
  • Python 语法糖：让编程更简单（续二）
    Python语法糖：让编程更简单（续）10.TypehintsTypehints是Python中的一种语法糖，用于指定函数或变量的类型。例如：defgreet(name:str)->None:print(f"Hello,{name}!")这段代码将定义一个名为greet的函数，它接受一个字符串参数name，并打印出Hello消息。11.ContextlibC......
  • Python 语法糖：让编程更简单（续）
    Python语法糖：让编程更简单（续）6.SlicenotationSlicenotation是Python中的一种语法糖，用于从列表或字符串中获取子串或子列表。例如：numbers=[1,2,3,4,5]print(numbers[1:3])#Output:[2,3]这段代码将从numbers列表中获取索引为1到3的子列表。7.f-stringsf-str......

赞助商

阅读排行

网站主页关于我们联系我们网站地图

本网站内容转载自其他媒体，侵权联系[admin##ips99.com]。

Copyright © 2020-2023 IPS99 版权所有 IPS99