将yolov5中的detect.py封装成一个方法（可结合大恒相机或者海康威视相机使用）

放心！能跑通！不会来问我！

import numpy as np
from torch import load,from_numpy,tensor
#import torch
from utils.general import check_img_size,scale_coords,xyxy2xywh,non_max_suppression
from utils.augmentations import letterbox
import cv2
from utils.plots import plot_one_box

# 初始化参数
# 初始化目录
# FILE = Path(__file__).resolve()
# ROOT = FILE.parents[0]  # 定义YOLOv5的根目录
# if str(ROOT) not in sys.path:
#     sys.path.append(str(ROOT))  # 将YOLOv5的根目录添加到环境变量中（程序结束后删除）
# ROOT = Path(os.path.relpath(ROOT, Path.cwd()))
#
# weights = ROOT / 'yolov5s.pt'  # 权重文件地址   .pt文件
# source = ROOT / 'data/images'  # 测试数据文件(图片或视频)的保存路径
# data = ROOT / 'data/coco128.yaml'  # 标签文件地址   .yaml文件


imgsz = (640, 640)  # 输入图片的大小 默认640(pixels)
conf_thres = 0.25  # object置信度阈值 默认0.25  用在nms中
iou_thres = 0.45  # 做nms的iou阈值 默认0.45   用在nms中
max_det = 1000  # 每张图片最多的目标数量  用在nms中
device = 'cpu'  # 设置代码执行的设备 cuda device, i.e. 0 or 0,1,2,3 or cpu
classes = None  # 在nms中是否是只保留某些特定的类 默认是None 就是所有类只要满足条件都可以保留 --class 0, or --class 0 2 3
agnostic_nms = False  # 进行nms是否也除去不同类别之间的框 默认False
augment = False  # 预测是否也要采用数据增强 TTA 默认False
visualize = False  # 特征图可视化 默认FALSE
half = False  # 是否使用半精度 Float16 推理 可以缩短推理时间 但是默认是False
dnn = False  # 使用OpenCV DNN进行ONNX推理

# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = select_device(device)
# print(device)


# 导入模型
model = load('./yolov5s.pt')["model"].float().eval()  # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
if half:
    model.half()  # to FP16

# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names  # 获取标签
# colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]


def detect_(img):
    im0 = img
    im = letterbox(im0)[0]
    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    im = np.ascontiguousarray(im)  # www 函数将一个内存不连续存储的数组转换为内存连续存储的数组，使得运行速度更快。
    im = from_numpy(im).to(device)
    im = im.half() if half else im.float()  # uint8 to fp16/32
    im /= 255  # 0 - 255 to 0.0 - 1.0
    if im.ndimension() == 3:
        im = im[None]  # 增加维度

    # Inference
    pred = model(im)[0]
    # NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

    # 用于存放结果
    # detections = []
    for i, det in enumerate(pred):  # per image 每张图片
        if len(det):
            det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
            # result
            for *xyxy, conf, cls in reversed(det):
                
                # xywh = (xyxy2xywh(tensor(xyxy).view(1, 4))).view(-1).tolist()
                # dic = {
                #    'class': f'{names[int(cls)]}',  # 检测目标对应的类别名
                #    'location': xywh,  # 坐标信息
                #    'score': round(float(conf) * 100, 2)  # 目标检测分数
                #}
                # detections.append(dic)
                # if dict['class'] == 'person':
                if names[int(cls)] == 'person':
                    label = f'{names[int(cls)]} {conf:.2f}'
                    plot_one_box(xyxy, im0, label=label, color=[255, 0, 255], line_thickness=2)                   label = f'{names[int(cls)]} {conf:.2f}'
                  
if __name__ == '__main__':
    pic=cv2.imread('./data/images/001.jpg')
    detect_(pic)