YOLO检测和分类pt模型的读取和使用

背景：

自从YOLOv5增加分类功能以来，YOLO成为了一个堪称全能的模型。然而，在具体实践中，不可能单单做检测或者单单做分类，往往需要一起使用或者集成到其他的代码中去。那么对于YOLOpt格式的模型如何读取和使用变成了一个极大的问题

以往的方法：

YOLO官方提供了export.py 文件，可以通过该文件将pt文件转换为其他格式的模型文件例如：ONNX。但是既然有YOLO自己读取pt文件的例子在，又何必多此一举转换成其他文件呢？

对于pt模型的使用，检测模型尚可，分类模型的使用极为繁琐，无法直接使用。下面就将介绍YOLOpt文件的读取和使用的具体方式----以YOLOv5为例

YOLO检测模型：

读取文件：

使用torch.hub.load来读取，注意最后一定要加上“local”，不然该方法默认从网上下载模型而不是本地读取。注意，读取模型的路径可以随意（不一定要放在YOLO文件夹内），但是YOLO的地址需要准确。

detecter = torch.hub.load('path of your YOLO root,例如：E:/Yolov5', 'custom', path='path of your pt,例如：E:/detect.pt', source='local')

检测方法：

直接检测，但是对于返回的值需要重点说明

#frame就是一个普通的图片cv2.imread直接读取的那种
results = detecter(frame)
#对于绝大多数信息都保存在results.xyxy[0]中，我也不知道为什么好好的xyxy非要放在另一个列表里
for xyxy in results.xyxy[0]:
    #获取到当前图片真正的xyxy，原来的格式应该是tensor，需要转化为列表
    xyxy = xyxy.tolist()
    #下面四行不用过多说明吧，检测框左上角坐标和右下角坐标
    xStart = xyxy[0]
    yStart = xyxy[1]
    xEnd = xyxy[2]
    yEnd = xyxy[3]
    #检测置信度
    conf = xyxy[4]
    #检测类别    
    cls = xyxy[5]

YOLO分类模型

读取文件：

和检测模型一样的读取方式

classifier_YOLO = torch.hub.load('E:\\yolov5-master', 'custom', path='models\\classify.pt', source='local')

使用方法：

对于YOLO分类模型的使用极为麻烦，如下，关于该函数请无脑使用，想理解可以配合源码服用。不想理解直接调用即可

#这些头文件是YOLO文件夹下 utils文件夹的文件，需要把这个文件复制到你的项目根目录下
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, print_args, strip_optimizer)
from utils.augmentations import classify_transforms

import torch
import torch.nn.functional as F

#classifier为pt模型，frame为待检测图片
def recognition_YOLO(classifier,frame):
    imgsz = (224, 224)
    stride, names, pt = classifier.stride, classifier.names, classifier.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    
    frame = classify_transforms(imgsz[0])(frame)
    
    #转为tensor类型
    frame = torch.Tensor(frame).to(classifier.device)
    if len(frame.shape) == 3:
        frame = frame[None]  # 增加batch维度
    #分类
    finalclasses = F.softmax(classifier(frame),dim=1)
    
    for i, prob in enumerate(finalclasses):  # per image
        #获取置信度最高的两个可能性
        top2i = prob.argsort(0, descending=True)[:2].tolist()  # top 2 indices
        classify_result = f"{', '.join(f'{names[j]} {prob[j]:.2f}' for j in top2i)}, "
    #返回分类结果
    return classify_result