基于YOLOv5的实时目标检测系统【支持图片、视频、RTSP流格式】

1、项目介绍

本项目是一个基于YOLOv5的实时目标检测系统，支持多种输入格式（图片、视频、RTSP流）进行检测。系统采用PyQt5作为GUI框架，用户可以通过图形界面便捷地进行检测任务的设置和执行。项目主要功能包括目标检测、实时视频流处理、结果统计及保存等。

1.1 项目简介

• 项目名称: YOLOv5目标检测系统

• 技术栈: PyQt5, PyTorch, OpenCV, YOLOv5

• 实现功能:

  • 支持图片、视频、RTSP流的输入

  • 实时目标检测和结果显示

  • 结果统计与保存

  • 模型自动加载与切换

  • 灵活的参数设置（置信度阈值、IoU阈值、检测延迟等）

1.2 技术栈

• PyQt5: 用于构建图形用户界面

• PyTorch: 深度学习框架，用于加载和推理YOLOv5模型

• OpenCV: 计算机视觉库，用于图像和视频处理

• YOLOv5: 目标检测模型，用于进行实时目标检测

1.3 实现的功能

1. 输入方式支持:

   1. 支持从文件选择图片或视频进行检测

   2. 支持通过摄像头进行实时视频流检测

   3. 支持通过RTSP流进行实时视频流检测

2. 参数设置:

   1. 置信度阈值（conf_thres）

   2. IoU阈值（iou_thres）

   3. 检测延迟（rate）

   4. 模型选择（支持多种YOLOv5模型）

3. 结果显示:

   1. 实时显示检测结果，包括目标框和类别标签

   2. 实时统计检测结果数量

   3. 支持结果自动保存

4. 模型切换:

   1. 支持在多个YOLOv5模型之间切换

   2. 自动搜索和加载模型文件

2、部分截图

3、部分核心代码

from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMenu, QAction
from main_win.win import Ui_mainWindow
from PyQt5.QtCore import Qt, QPoint, QTimer, QThread, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap, QPainter, QIcon

import sys
import os
import json
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import os
import time
import cv2
import traceback
from models.experimental import attempt_load
from utils.dataloaders import LoadImages, LoadWebcam
from utils.dataloaders import create_dataloader
from utils.CustomMessageBox import MessageBox
from utils.general import check_img_size, check_requirements, check_imshow, colorstr, non_max_suppression, \
    apply_classifier, scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device
from utils.capnums import Camera
from dialog.rtsp_win import Window

class DetThread(QThread):
    send_img = pyqtSignal(np.ndarray)
    send_raw = pyqtSignal(np.ndarray)
    send_statistic = pyqtSignal(dict)
    send_msg = pyqtSignal(str)
    send_percent = pyqtSignal(int)
    send_fps = pyqtSignal(str)

    def __init__(self):
        super(DetThread, self).__init__()
        self.weights = './yolov5s.pt'
        self.current_weight = './yolov5s.pt'
        self.source = '0'
        self.conf_thres = 0.25
        self.iou_thres = 0.45
        self.jump_out = False
        self.is_continue = True
        self.percent_length = 1000
        self.rate_check = True
        self.rate = 100
        self.save_fold = './result'

    @torch.no_grad()
    def run(self, imgsz=640, max_det=1000, device='', view_img=True, save_txt=False, save_conf=False, save_crop=False, 
            nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, 
            project='runs/detect', name='exp', exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, 
            half=False):
        try:
            device = select_device(device)
            half &= device.type != 'cpu'

            model = attempt_load(self.weights, device=device)
            num_params = sum(p.numel() for p in model.parameters())
            stride = int(model.stride.max())
            imgsz = check_img_size(imgsz, s=stride)
            names = model.module.names if hasattr(model, 'module') else model.names
            if half:
                model.half()

            if self.source.isnumeric() or self.source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')):
                view_img = check_imshow()
                cudnn.benchmark = True
                dataset = LoadWebcam(self.source, img_size=imgsz, stride=stride)
            else:
                dataset = LoadImages(self.source, img_size=imgsz, stride=stride)

            if device.type != 'cpu':
                model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))

            count = 0
            start_time = time.time()
            dataset = iter(dataset)

            while True:
                if self.jump_out:
                    self.vid_cap.release()
                    self.send_percent.emit(0)
                    self.send_msg.emit('Stop')
                    if hasattr(self, 'out'):
                        self.out.release()
                    break

                if self.current_weight != self.weights:
                    model = attempt_load(self.weights, device=device)
                    num_params = sum(p.numel() for p in model.parameters())
                    stride = int(model.stride.max())
                    imgsz = check_img_size(imgsz, s=stride)
                    names = model.module.names if hasattr(model, 'module') else model.names
                    if half:
                        model.half()
                    if device.type != 'cpu':
                        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
                    self.current_weight = self.weights

                if self.is_continue:
                    path, img, im0s, self.vid_cap = next(dataset)
                    count += 1
                    if count % 30 == 0 and count >= 30:
                        fps = int(30/(time.time()-start_time))
                        self.send_fps.emit('fps：'+str(fps))
                        start_time = time.time()
                    if self.vid_cap:
                        percent = int(count/self.vid_cap.get(cv2.CAP_PROP_FRAME_COUNT)*self.percent_length)
                        self.send_percent.emit(percent)
                    else:
                        percent = self.percent_length

                    statistic_dic = {names[name]: 0 for name in names}
                    img = torch.from_numpy(img).to(device)
                    img = img.half() if half else img.float()
                    img /= 255.0
                    if img.ndimension() == 3:
                        img = img.unsqueeze(0)

                    pred = model(img, augment=augment)[0]
                    pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes, agnostic_nms, max_det=max_det)
                    for i, det in enumerate(pred):
                        im0 = im0s.copy()
                        annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                        if len(det):
                            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
                            for *xyxy, conf, cls in reversed(det):
                                c = int(cls)
                                statistic_dic[names[c]] += 1
                                label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                                annotator.box_label(xyxy, label, color=colors(c, True))

                    if self.rate_check:
                        time.sleep(1/self.rate)
                    im0 = annotator.result()
                    self.send_img.emit(im0)
                    self.send_raw.emit(im0s if isinstance(im0s, np.ndarray) else im0s[0])
                    self.send_statistic.emit(statistic_dic)
                    if self.save_fold:
                        os.makedirs(self.save_fold, exist_ok=True)
                        if self.vid_cap is None:
                            save_path = os.path.join(self.save_fold, time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime()) + '.jpg')
                            cv2.imwrite(save_path, im0)
                        else:
                            if count == 1:
                                ori_fps = int(self.vid_cap.get(cv2.CAP_PROP_FPS))
                                if ori_fps == 0:
                                    ori_fps = 25
                                width, height = im0.shape[1], im0.shape[0]
                                save_path = os.path.join(self.save_fold, time.strftime('%Y_%m_%d_%H_%M_%S', time.localtime()) + '.mp4')
                                self.out = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), ori_fps, (width, height))
                            self.out.write(im0)

4、运行视频

点击查看运行视频

5、最后

点击了解更多