Python Yolo8:YOLOV6

1. 准备

# 创建虚拟环境
conda create -n yolov6env python=3.8

# 这种方式创建出来没有name，该文件夹
PS D:\Cumtb_Code> conda create --prefix=D:\Cumtb_Code\mincondatest python=3.8

# 进入虚拟环境，我这里将文件路径移动到了目录下
D:\Cumtb_Code\yolov6env>conda activate d:\Cumtb_Code\yolov6env

2. YOLOV6

2.1 拉取Yolov6

git clone https://github.com/meituan/YOLOv6
cd YOLOv6

2.2 安装相应的库

pip install -r requirements.txt

# 验证Torch
import torch
print(torch.__version__)
print(torch.cuda.is_available())

PyTorch

Reduce inference costs by 71% and drive scale out using PyTorch, TorchServe, and AWS Inferentia. Pushing the state of the art in NLP and Multi-task learning. Using PyTorch's flexibility to efficiently research new algorithmic approaches.

https://pytorch.org/

conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge

2.2 下载测试模型pt

Release YOLOv6 0.1.0 · meituan/YOLOv6

YOLOv6 is a single-stage object detection framework dedicated to industrial application, with hardware-friendly efficient design and high performance, outperforming YOLOv5, YOLOX and PP-YOLOE. YOLOv6-nano achieves 35.0 mAP on COCO val2017 dataset with 1242 FPS on T4 using TensorRT FP16 for bs32 inference, and YOLOv6-s achieves 43.1 mAP on COCO val2017 dataset with 520 FPS on T4 using TensorRT FP16 for bs32 inference.

https://github.com/meituan/YOLOv6/releases/tag/0.1.0

2.3 运行测试

# 单个文件
python tools/infer.py --weights weight/yolov6s.pt --source data/images/image1.jpg

# 文件夹下的多个文件
python tools/infer.py --weights weight/yolov6s.pt --source data/images

3. YOLOV6导出模型

3.1 导出FP32 ONNX模型

python ./deploy/ONNX/export_onnx.py --weights weight/yolov6s.pt --img 640 --batch 1 --end2end

3.2 导出FP16 ONNX模型

python ./deploy/ONNX/export_onnx.py --weights weight/yolov6s.pt --img 640 --batch 1 --end2end --half

3.3 导出engine模型

# fp32

trtexec --onnx=yolov6s.onnx --saveEngine=yolov6s.engine --workspace=30

# fp16
trtexec --onnx=yolov6s.onnx --saveEngine=yolov6s.engine --fp16 --workspace=30

安装tensorrt

安装

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5>cd python

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5\python>pip install tensorrt-8.4.1.5-cp38-none-win_amd64.whl

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5>cd uff 

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5\uff>pip install uff-0.6.9-py2.py3-none-any.whl

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5>cd graphsurgeon

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5\graphsurgeon>pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5>cd onnx_graphsurgeon

(d:\Cumtb_Code\yolov6env) D:\TensorRT-8.4.1.5\onnx_graphsurgeon>pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

(d:\Cumtb_Code\yolov6env) D:\Cumtb_Code\yolov6env>python
Python 3.8.13 (default, Mar 28 2022, 06:59:08) [MSC v.1916 64 bit (AMD64)] :: Anaconda, Inc. on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorrt
>>> print(tensorrt.__version__)
8.4.1.5

运行

• inferDemo.py

import cv2
import torch
import numpy as np
import tensorrt as trt
from collections import OrderedDict,namedtuple

class TRT_engine():
    def __init__(self, weight) -> None:
        self.imgsz = [640,640]
        self.weight = weight
        self.device = torch.device('cuda:0')
        self.init_engine()

    def init_engine(self):
        # Infer TensorRT Engine
        self.Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
        self.logger = trt.Logger(trt.Logger.INFO)
        trt.init_libnvinfer_plugins(self.logger, namespace="")
        with open(self.weight, 'rb') as self.f, trt.Runtime(self.logger) as self.runtime:
            self.model = self.runtime.deserialize_cuda_engine(self.f.read())
        self.bindings = OrderedDict()
        self.fp16 = False
        for index in range(self.model.num_bindings):
            self.name = self.model.get_binding_name(index)
            self.dtype = trt.nptype(self.model.get_binding_dtype(index))
            self.shape = tuple(self.model.get_binding_shape(index))
            self.data = torch.from_numpy(np.empty(self.shape, dtype=np.dtype(self.dtype))).to(self.device)
            self.bindings[self.name] = self.Binding(self.name, self.dtype, self.shape, self.data, int(self.data.data_ptr()))
            if self.model.binding_is_input(index) and self.dtype == np.float16:
                self.fp16 = True
        self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())
        self.context = self.model.create_execution_context()

    def letterbox(self,im,color=(114, 114, 114), auto=False, scaleup=True, stride=32):
        # Resize and pad image while meeting stride-multiple constraints
        shape = im.shape[:2]  # current shape [height, width]
        new_shape = self.imgsz
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)
        # Scale ratio (new / old)
        self.r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        if not scaleup:  # only scale down, do not scale up (for better val mAP)
            self.r = min(self.r, 1.0)
        # Compute padding
        new_unpad = int(round(shape[1] * self.r)), int(round(shape[0] * self.r))
        self.dw, self.dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
        if auto:  # minimum rectangle
            self.dw, self.dh = np.mod(self.dw, stride), np.mod(self.dh, stride)  # wh padding
        self.dw /= 2  # divide padding into 2 sides
        self.dh /= 2
        if shape[::-1] != new_unpad:  # resize
            im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(self.dh - 0.1)), int(round(self.dh + 0.1))
        left, right = int(round(self.dw - 0.1)), int(round(self.dw + 0.1))
        self.img = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
        return self.img,self.r,self.dw,self.dh

    def preprocess(self,image):
        self.img,self.r,self.dw,self.dh = self.letterbox(image)
        self.img = cv2.cvtColor(self.img,cv2.COLOR_BGR2RGB)
        self.img = self.img.transpose((2, 0, 1))
        self.img = np.expand_dims(self.img,0)
        self.img = np.ascontiguousarray(self.img)
        self.img = torch.from_numpy(self.img).to(self.device)
        self.img = self.img.half() if self.fp16 else self.img.float()
        self.img /= 255.
        return self.img

    def predict(self,img,threshold):
        img = self.preprocess(img)
        self.binding_addrs['image_arrays'] = int(img.data_ptr())
        self.context.execute_v2(list(self.binding_addrs.values()))
        nums = self.bindings['num_dets'].data[0].tolist()
        boxes = self.bindings['det_boxes'].data[0].tolist()
        scores =self.bindings['det_scores'].data[0].tolist()
        classes = self.bindings['det_classes'].data[0].tolist()
        num = int(nums[0])
        new_bboxes = []
        for i in range(num):
            if(scores[i] < threshold):
                continue
            xmin = (boxes[i][0] - self.dw)/self.r
            ymin = (boxes[i][1] - self.dh)/self.r
            xmax = (boxes[i][2] - self.dw)/self.r
            ymax = (boxes[i][3] - self.dh)/self.r
            new_bboxes.append([classes[i],scores[i],xmin,ymin,xmax,ymax])
        return new_bboxes

def visualize(img,bbox_array):
    for temp in bbox_array:
        xmin = int(temp[2])
        ymin = int(temp[3])
        xmax = int(temp[4])
        ymax = int(temp[5])
        clas = int(temp[0])
        score = temp[1]
        cv2.rectangle(img,(xmin,ymin),(xmax,ymax), (105, 237, 249), 2)
        img = cv2.putText(img, "class:"+str(clas)+" "+str(round(score,2)), (xmin,int(ymin)-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (105, 237, 249), 1)
    return img

trt_engine = TRT_engine("./weight/yolov6s.engine")
img1 = cv2.imread("./data/images/image1.jpg")
results = trt_engine.predict(img1,threshold=0.5)
img = visualize(img1,results)
cv2.imshow("img",img)
cv2.waitKey(0)

• 这个地方运行报错了Torch版本问题

python inferDemo.py

报错

SyntaxError: invalid syntax

File "D:\Cumtb_Code\yolov6env\YOLOv6\yolov6\utils\torch_utils.py", line 105
    flops, params = profile(deepcopy(model), inputs=(img,), verbose=False)
    ^
SyntaxError: invalid syntax

ModuleNotFoundError: No module named 'tensorrt’

安装

raise AssertionError("Torch not compiled with CUDA enabled")

import torch
print(torch.__version__)
print(torch.cuda.is_available())

PyTorch

Reduce inference costs by 71% and drive scale out using PyTorch, TorchServe, and AWS Inferentia. Pushing the state of the art in NLP and Multi-task learning. Using PyTorch's flexibility to efficiently research new algorithmic approaches.

https://pytorch.org/

conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge