yolov5的detect.py代码修改＋调用摄像头+ROS话题发布

 一、detect重要参数详解

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='C:/Users/29294/Desktop/Celery_cabbage/runs/train/exp31/weights/best.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    opt = parser.parse_args()
    print(opt)
    check_requirements(exclude=('pycocotools', 'thop'))

weights:参数为模型文件所在目录;
source:参数为要识别的图片的路径(多个图片可以是其文件夹路径);

Imgse:表示图像输入大小，yolo会根据这个参数剪裁图片;

--conf-thres为识别的置信度;

iou-thres:“预测的边框”和“真实的边框”的交集和并集的比值，这个IoU值设置小一些可以防止出现重复框;

max-det:参数为最大侦测的目标数;
device参数为用来识别的设备，数字表示gpu号，若需要cpu运行则将参数改为cpu;

view-img参数表示是否展示识别图像;
save-txt、--save-txt--save-conf、--save-crop、--nosave都是保存相关结果的;

classes参数表示只检测某一目标;

agnostic-nms表示是否进行nms去除不同类别之间的框

augment表示推理的时候进行多尺度，翻转等操作推理;
visualize表示保存可视化特征图片;update参数表示更新模型;
project、name、exist-ok为识别结果图保存地址设置;

line-thickness为画框的粗细;
hide-labels、hide-conf为是否在识别结果里添加标签名字置信度
half参数表示是否使用FP16半精度推理:
dnn参数表示是否使用OpenCV DNN进行ONNX推断

 二、化简detect.py

        保留的图像识别必须的模块，将其他冗余的函数去除掉

"""
import rospy
from std_msgs.msg import Int8

import argparse
import os
import sys
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.backends.cudnn as cudnn

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.experimental import attempt_load
from utils.datasets import LoadImages, LoadStreams
from utils.general import apply_classifier, check_img_size, check_imshow, check_requirements, check_suffix, colorstr, \
    increment_path, non_max_suppression, print_args, save_one_box, scale_coords, set_logging, \
    strip_optimizer, xyxy2xywh
from utils.plots import Annotator, colors
from utils.torch_utils import load_classifier, select_device, time_sync


#被torch.no_grad()包住的代码,不需要计算梯度,也不会进行反向传播 
@torch.no_grad()


#检测函数
def run(weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
        source='0',  # file/dir/URL/glob, 0 for webcam
        project=ROOT / 'runs/detect/exp',  # save results to project/name
        imgsz=640,  # inference size (pixels)
        conf_thres=0.6,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS
        max_det=1000,  # maximum detections per image
        device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        ):

 # Directories
    save_dir = increment_path(Path(project), exist_ok=False, mkdir=True)  # increment and make dir
    print(save_dir)

    # Initialize
    set_logging()
    device = select_device(device)


    ##################################首先加载模型#########################################
    model = attempt_load(weights, map_location=device)
    stride = int(model.stride.max())  # model stride
    names = model.module.names if hasattr(model, 'module') else model.names  # get class names


    ################################然后加载识别对象#######################################
    dataset = LoadStreams(str(source), img_size=imgsz, stride=stride, auto=True)
    a = ""

    ####################################开始识别##########################################
    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters())))  # run once


    dt, seen = [0.0, 0.0, 0.0], 0



    #注意这个循环是，每有一张图片循环一次
    for path, img, im0s, vid_cap in dataset:
        #首先将图片转为张量
        t1 = time_sync()
        img = torch.from_numpy(img).to(device)
        img = img / 255.0  # 0 - 255 to 0.0 - 1.0
        if len(img.shape) == 3:
            img = img[None]  # expand for batch dim
        t2 = time_sync()
        dt[0] += t2 - t1

        #开始推理，也就是检测
        pred = model(img)[0]
        t3 = time_sync()
        dt[1] += t3 - t2

        #NMS非极大值抑制，过滤、去除框选
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        dt[2] += time_sync() - t3


        #注意这个循环是，每检测到一种特征循环一次
        for i, det in enumerate(pred):  # per image

            im0 = im0s[i].copy()
            person = 0
            

            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

               

                #注意这个循环是，某种特征每检测到一个循环一次
                #其中xyxy[0]-xyxy[3]对应目标特征的x0\y0\x1\y1
                #int(cls)对应获得识别到的特征种类的编号，names[int(cls)]是识别到的种类
                #config对应此次特征的置信度
                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)
                    if c == 0 :
                        person += 1
                    
                        #画框、保存识别效果图片
                        label = f'{names[int(cls)]} {conf:.2f}'
                        annotator = Annotator(im0, line_width=3 , example=str(names))
                        annotator.box_label(xyxy, label, color=colors(c, True))

三、ROS话题发布

        1.导入rospy

    import rospy

        2.创建ros节点和发布者

    #初始化ROS节点
    rospy.init_node('yolo', anonymous=True)
    # 创建话题发布者
    pub = rospy.Publisher('/person', Int8, queue_size=10)

        3.发布话题消息

    # 发布识别信息话题
    pub.publish(person)

        4.查看消息

    rostopic echo /person

四、与摄像头通讯

        1.将摄像头数据保存为视频

        # 保存识别结果
        if a != path:  # new video
            a = path
            vid_writer = None
            if isinstance(vid_writer, cv2.VideoWriter):
                vid_writer.release()  # release previous video writer
            fps, w, h = 30, im0.shape[1], im0.shape[0]
            vid_writer = cv2.VideoWriter(str(save_dir / source) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        vid_writer.write(im0)

        2.显示识别结果

        # 显示识别结果
        cv2.imshow('', im0)
        cv2.waitKey(1)  # 1 millisecond

五、其他

        ROS默认支持的是PYTHON2，而YOLOv5需要PYTHON3环境，所以这里我们要保证PYTHON3可以使用rospy安装:

    pip3 install rospkg
    pip3 install catkin-tools

        然后在识别代码里 import rospy就可以正常发布话题

        最后加一个主函数调用即可：

    if __name__ == "__main__":
        run()