当前位置:

新鲜出炉的 yoloV5可视化实战项目(1)

访客 2023-07-03 1271 0

闲谈

    了解到目标检测算法,越来越觉得有意思,希望能做一些更有趣的东西,发现yolov5是个非常牛逼的框架,做点东西实践下,顺便记录下学习的历程,方便复习。 可以到, 目前已更新到v7版本,我们还是用v5版本 本项目的代码已上传至github 如遇到问题或疑问,可以直接留言或私信 祝大家学习愉快

展示效果

开始

    一般官网的代码会比较完整,需要做一些修改才能符合我们的需求 下载完官网的项目后,直接运行安装依赖即可
pip install -r requirements.txt
    安装完运行
import torch# Modelmodel = torch.hub.load(ultralytics/yolov5, yolov5s)  # or yolov5n - yolov5x6, custom# Imagesimg = https://ultralytics.com/images/zidane.jpg  # or file, Path, PIL, OpenCV, numpy, list# Inferenceresults = model(img)# Resultsresults.print()  # or .show(), .save(), .crop(), .pandas(), etc.
    方案二
python detect.py --source data/images/bus.jpg

改装

    制作一个可视化的界面 编写detect.py文件 逻辑关联处理

可视化工具

    我们这里用的是pyqt5的工具 关于pyqt5 可以参考这个教程,写的很好了

界面制作

做完大概长这样

逻辑交互制作

def init_slots(self):                self.ui.imgScan.clicked.connect(self.button_image_open)                self.ui.videoScan.clicked.connect(self.button_video_open)                self.ui.capScan.clicked.connect(self.button_camera_open)                self.ui.loadWeight.clicked.connect(self.open_model)                self.ui.initModel.clicked.connect(self.model_init)                self.ui.start.clicked.connect(self.toggleState)                self.ui.end.clicked.connect(self.endVideo)                        # self.ui.pushButton_stop.clicked.connect(self.button_video_stop)                        # self.ui.pushButton_finish.clicked.connect(self.finish_detect)                self.timer_video.timeout.connect(self.show_video_frame) # 定时器超时,将槽绑定至show_video_frame
    这部分很简单就是绑定下函数, 这里制作了一个读取视频的定时器,方便控制视频的读取

关键函数detect

    这个是本项目的关键代码,核心代码,其他的都只能是辅助,同时也是最难的部分。
def detect(self, name_list, img):                                :param name_list: 文件名列表                :param img: 待检测图片                :return: info_show:检测输出的文字信息                                showimg = img                # 不使用图计算                with torch.no_grad():                        # 填充为标准的大小,letterbox 不同于resize, 使用的是边缘填充                        img = letterbox(img, new_shape=self.opt.img_size)[0]                        # Convert                        # 转换成Rgb模式                        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416                        # 转换成连续的内存格式,加快计算速度                        img = np.ascontiguousarray(img)                        # 图和ndarray的转换,并制定cpu或gpu                        img = torch.from_numpy(img).to(self.device)                        # 设置精度                        img = img.half() if self.half else img.float()  # uint8 to fp16/32                        # 归一化                        img /= 255.0  # 0 - 255 to 0.0 - 1.0                        # 统一格式                        if img.ndimension() == 3:                                # 增加一个维度                                img = img.unsqueeze(0)                        # Inference                        # 使用模型预测                        pred = self.model(img, augment=self.opt.augment)[0]                        # Apply NMS                        # 非极大值抑制                        pred = non_max_suppression(pred, self.opt.conf_thres, self.opt.iou_thres, classes=self.opt.classes,                        agnostic=self.opt.agnostic_nms)                        info_show = ""                        # Process detections                                                for i, det in enumerate(pred):                            if det is not None and len(det):                                # Rescale boxes from img_size to im0 size                                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()                                for *xyxy, conf, cls in reversed(det):                                        label = %s %.2f % (self.names[int(cls)], conf)                                        name_list.append(self.names[int(cls)])                                        single_info = plot_one_box2(xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)                                        # print(single_info)                                        info_show = info_show + single_info + ""                return  info_show
    具体的都有注释

模型的初始化和权重参数的加载

    权重加载
def open_model(self):                self.openfile_name_model, _ = QFileDialog.getOpenFileName(self, 选择权重文件, directory=./yolov5yoloYoloV5_PyQt5-mainweights)                print(self.openfile_name_model)                if not self.openfile_name_model:                #    QtWidgets.QMessageBox.warning(self, u"Warning" u未选择权重文件,请重试, buttons=QtWidgets.QMessageBox.Ok)                   self.log("warining 未选择权重文件,请重试")                else :                   print(self.openfile_name_model)                   self.log("权重文件路径为:%s"%self.openfile_name_model)                pass
    模型加载
def model_init(self):                # 模型相关参数配置                parser = argparse.ArgumentParser()                parser.add_argument(--weights, nargs=+, type=str, default=weights/yolov5s.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.25, help=IOU threshold for NMS)                parser.add_argument(--device, default=cpu, 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(--save, default=False, help=save video to local)                parser.add_argument(--exist-ok, action=store_true, help=existing project/name ok, do not increment)                self.opt = parser.parse_args()                print(self.opt)                # 使用模型的默认参数                source, self.weights, imgsz = self.opt.source, self.opt.weights, self.opt.img_size                                # 若有配置自定义权重,                if hasattr(self, "openfile_name_model") and self.openfile_name_model.endswith(.pt):                     self.weights = self.openfile_name_model                     self.log(记载自定义权重成功)                else :                     self.log(warning:权重文件有误,请重新加载)                                # 选择cpu 和gpu                self.device = select_device(self.opt.device)                self.half = self.device.type != cpu                cudnn.benchmark = True                # Load model                self.model = attempt_load(self.weights, map_location=self.device)                stride = int(self.model.stride.max())  # model stride                self.imgsz = check_img_size(imgsz, s=stride)  # check img_size                if self.half:                    self.model.half()  # to FP16                # Get names and colors                self.names = self.model.module.names if hasattr(self.model, module) else self.model.names                self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]                QtWidgets.QMessageBox.information(self, u"Notice", u"模型加载完成", buttons=QtWidgets.QMessageBox.Ok,                                      defaultButton=QtWidgets.QMessageBox.Ok)                self.log(模型加载完成)

设置图片识别

def button_image_open(self):            print(button_image_open)            name_list = []            try:                img_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "选择文件")            except OSError as reason:                print(文件出错啦)                QtWidgets.QMessageBox.warning(self, Warning, 文件出错, buttons=QtWidgets.QMessageBox.Ok)            else:                if not img_name:                   QtWidgets.QMessageBox.warning(self,"Warning", 文件出错, buttons=QtWidgets.QMessageBox.Ok)                   self.log(文件出错)                else:                    img = cv.imread(img_name)                    info_show = self.detect(name_list, img)                    date = time.strftime(%Y-%m-%d-%H-%M-%S, time.localtime(time.time())) # 当前时间                    file_extaction = img_name.split(.)[-1]                    new_fileName = date + . + file_extaction                    file_path = self.output_folder + img_output/ + new_fileName                    cv.imwrite(file_path, img)                    self.show_img(info_show, img)                                    #     self.log(info_show) #检测信息                                    #     self.result = cv.cvtColor(img, cv.COLOR_BGR2BGRA)                                    #     self.result =  letterbox(self.result, new_shape=self.opt.img_size)[0] #cv.resize(self.result, (640, 480), interpolation=cv.INTER_AREA)                #     self.QtImg = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0], QtGui.QImage.Format_RGB32)                #     print(type(self.ui.show))                #     self.ui.show.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))                #     self.ui.show.setScaledContents(True) # 设置图像自适应界面大小

视频和摄像头

def button_video_open(self):                video_path, _ = QtWidgets.QFileDialog.getOpenFileName(self, 选择检测视频, filter="*.mp4;;*.avi;;All Files(*)")                flag = self.cap.open(video_path)                if not flag:                        QtWidgets.QMessageBox.warning(self,"Warning", 打开视频失败, buttons=QtWidgets.QMessageBox.Ok)                else:                         self.timer_video.start(1000/self.cap.get(cv.CAP_PROP_FPS)) # 以30ms为间隔,启动或重启定时器                        if self.opt.save:                                fps, w, h, path = self.set_video_name_and_path()                                self.vid_writer = cv.VideoWriter(path, cv.VideoWriter_fourcc(*mp4v), fps, (w, h))def button_camera_open(self):                camera_num = 0                self.cap = cv.VideoCapture(camera_num)                if not self.cap.isOpened():                        QtWidgets.QMessageBox.warning(self, u"Warning", u摄像头打开失败, buttons=QtWidgets.QMessageBox.Ok)                else:                        self.timer_video.start(1000/60)                        if self.opt.save:                                fps, w, h, path = self.set_video_name_and_path()                                self.vid_writer = cv.VideoWriter(path, cv.VideoWriter_fourcc(*mp4v), fps, (w, h))

知识点

    图片的显示用qlabel就可以了
def show_img(self, info_show, img):                self.log(info_show)                show = cv.resize(img, (640, 480)) # 直接将原始img上的检测结果进行显示                self.result = cv.cvtColor(show, cv.COLOR_BGR2RGB)                self.result =  letterbox(self.result, new_shape=self.opt.img_size)[0]                showImage = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0],                                     QtGui.QImage.Format_RGB888)                self.ui.show.setPixmap(QtGui.QPixmap.fromImage(showImage))                self.ui.show.setScaledContents(True)  # 设置图像自适应界面大小

完整的代码

import argparseimport randomimport sysimport timeimport torchfrom PyQt5.QtWidgets import *from PyQt5 import QtCore, QtGui, QtWidgetsfrom PyQt5.QtWidgets import QApplication,QMainWindowimport namefrom functools import partialfrom utils.torch_utils import select_deviceimport torch.backends.cudnn as cudnnimport cv2 as cvfrom utils.datasets import letterboxfrom utils.plots import plot_one_box2import numpy as npfrom models.experimental import attempt_loadfrom utils.general import check_img_size, non_max_suppression, scale_coords# pyuic5 -o name.py test.uiclass UI_Logic_Window(QtWidgets.QMainWindow):        def __init__(self, parent = None):                super(UI_Logic_Window, self).__init__(parent)                self.timer_video = QtCore.QTimer() # 创建定时器                #创建一个窗口                self.w = QMainWindow()                self.ui = name.Ui_MainWindow()                self.ui.setupUi(self)                self.init_slots()                self.output_folder = output/                self.cap = cv.VideoCapture()                # 日志                self.logging =        # 控件绑定相关操作        def init_slots(self):                self.ui.imgScan.clicked.connect(self.button_image_open)                self.ui.videoScan.clicked.connect(self.button_video_open)                self.ui.capScan.clicked.connect(self.button_camera_open)                self.ui.loadWeight.clicked.connect(self.open_model)                self.ui.initModel.clicked.connect(self.model_init)                self.ui.start.clicked.connect(self.toggleState)                self.ui.end.clicked.connect(self.endVideo)                        # self.ui.pushButton_stop.clicked.connect(self.button_video_stop)                        # self.ui.pushButton_finish.clicked.connect(self.finish_detect)                self.timer_video.timeout.connect(self.show_video_frame) # 定时器超时,将槽绑定至show_video_frame        def button_image_open(self):            print(button_image_open)            name_list = []            try:                img_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "选择文件")            except OSError as reason:                print(文件出错啦)                QtWidgets.QMessageBox.warning(self, Warning, 文件出错, buttons=QtWidgets.QMessageBox.Ok)            else:                if not img_name:                   QtWidgets.QMessageBox.warning(self,"Warning", 文件出错, buttons=QtWidgets.QMessageBox.Ok)                   self.log(文件出错)                else:                    img = cv.imread(img_name)                    info_show = self.detect(name_list, img)                    date = time.strftime(%Y-%m-%d-%H-%M-%S, time.localtime(time.time())) # 当前时间                    file_extaction = img_name.split(.)[-1]                    new_fileName = date + . + file_extaction                    file_path = self.output_folder + img_output/ + new_fileName                    cv.imwrite(file_path, img)                    self.show_img(info_show, img)                                    #     self.log(info_show) #检测信息                                    #     self.result = cv.cvtColor(img, cv.COLOR_BGR2BGRA)                                    #     self.result =  letterbox(self.result, new_shape=self.opt.img_size)[0] #cv.resize(self.result, (640, 480), interpolation=cv.INTER_AREA)                #     self.QtImg = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0], QtGui.QImage.Format_RGB32)                #     print(type(self.ui.show))                #     self.ui.show.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))                #     self.ui.show.setScaledContents(True) # 设置图像自适应界面大小        def show_img(self, info_show, img):                self.log(info_show)                show = cv.resize(img, (640, 480)) # 直接将原始img上的检测结果进行显示                self.result = cv.cvtColor(show, cv.COLOR_BGR2RGB)                self.result =  letterbox(self.result, new_shape=self.opt.img_size)[0]                showImage = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0],                                     QtGui.QImage.Format_RGB888)                self.ui.show.setPixmap(QtGui.QPixmap.fromImage(showImage))                self.ui.show.setScaledContents(True)  # 设置图像自适应界面大小        def toggleState(self):                print(toggle)                state = self.timer_video.signalsBlocked()                self.timer_video.blockSignals(not state)                text = 继续 if not state else 暂停                self.ui.start.setText(text)        def endVideo(self):                print(end)                self.timer_video.blockSignals(True)                self.releaseRes()        def detect(self, name_list, img):                                :param name_list: 文件名列表                :param img: 待检测图片                :return: info_show:检测输出的文字信息                                showimg = img                # 不使用图计算                with torch.no_grad():                        # 填充为标准的大小,letterbox 不同于resize, 使用的是边缘填充                        img = letterbox(img, new_shape=self.opt.img_size)[0]                        # Convert                        # 转换成Rgb模式                        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416                        # 转换成连续的内存格式,加快计算速度                        img = np.ascontiguousarray(img)                        # 图和ndarray的转换,并制定cpu或gpu                        img = torch.from_numpy(img).to(self.device)                        # 设置精度                        img = img.half() if self.half else img.float()  # uint8 to fp16/32                        # 归一化                        img /= 255.0  # 0 - 255 to 0.0 - 1.0                        # 统一格式                        if img.ndimension() == 3:                                # 增加一个维度                                img = img.unsqueeze(0)                        # Inference                        # 使用模型预测                        pred = self.model(img, augment=self.opt.augment)[0]                        # Apply NMS                        # 非极大值抑制                        pred = non_max_suppression(pred, self.opt.conf_thres, self.opt.iou_thres, classes=self.opt.classes,                        agnostic=self.opt.agnostic_nms)                        info_show = ""                        # Process detections                                                for i, det in enumerate(pred):                            if det is not None and len(det):                                # Rescale boxes from img_size to im0 size                                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()                                for *xyxy, conf, cls in reversed(det):                                        label = %s %.2f % (self.names[int(cls)], conf)                                        name_list.append(self.names[int(cls)])                                        single_info = plot_one_box2(xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)                                        # print(single_info)                                        info_show = info_show + single_info + ""                return  info_show                   def button_video_open(self):                video_path, _ = QtWidgets.QFileDialog.getOpenFileName(self, 选择检测视频, filter="*.mp4;;*.avi;;All Files(*)")                flag = self.cap.open(video_path)                if not flag:                        QtWidgets.QMessageBox.warning(self,"Warning", 打开视频失败, buttons=QtWidgets.QMessageBox.Ok)                else:                         self.timer_video.start(1000/self.cap.get(cv.CAP_PROP_FPS)) # 以30ms为间隔,启动或重启定时器                        if self.opt.save:                                fps, w, h, path = self.set_video_name_and_path()                                self.vid_writer = cv.VideoWriter(path, cv.VideoWriter_fourcc(*mp4v), fps, (w, h))        def set_video_name_and_path(self):                # 获取当前系统时间,作为img和video的文件名                now = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(time.time()))                # if vid_cap:  # video                fps = self.cap.get(cv.CAP_PROP_FPS)                w = int(self.cap.get(cv.CAP_PROP_FRAME_WIDTH))                h = int(self.cap.get(cv.CAP_PROP_FRAME_HEIGHT))                # 视频检测结果存储位置                save_path = self.output_folder + video/ + now + .mp4                return fps, w, h, save_path        def button_camera_open(self):                camera_num = 0                self.cap = cv.VideoCapture(camera_num)                if not self.cap.isOpened():                        QtWidgets.QMessageBox.warning(self, u"Warning", u摄像头打开失败, buttons=QtWidgets.QMessageBox.Ok)                else:                        self.timer_video.start(1000/60)                        if self.opt.save:                                fps, w, h, path = self.set_video_name_and_path()                                self.vid_writer = cv.VideoWriter(path, cv.VideoWriter_fourcc(*mp4v), fps, (w, h))                                def open_model(self):                self.openfile_name_model, _ = QFileDialog.getOpenFileName(self, 选择权重文件, directory=./yolov5yoloYoloV5_PyQt5-mainweights)                print(self.openfile_name_model)                if not self.openfile_name_model:                #    QtWidgets.QMessageBox.warning(self, u"Warning" u未选择权重文件,请重试, buttons=QtWidgets.QMessageBox.Ok)                   self.log("warining 未选择权重文件,请重试")                else :                   print(self.openfile_name_model)                   self.log("权重文件路径为:%s"%self.openfile_name_model)                pass        def model_init(self):                # 模型相关参数配置                parser = argparse.ArgumentParser()                parser.add_argument(--weights, nargs=+, type=str, default=weights/yolov5s.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.25, help=IOU threshold for NMS)                parser.add_argument(--device, default=cpu, 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(--save, default=False, help=save video to local)                parser.add_argument(--exist-ok, action=store_true, help=existing project/name ok, do not increment)                self.opt = parser.parse_args()                print(self.opt)                # 使用模型的默认参数                source, self.weights, imgsz = self.opt.source, self.opt.weights, self.opt.img_size                                # 若有配置自定义权重,                if hasattr(self, "openfile_name_model") and self.openfile_name_model.endswith(.pt):                     self.weights = self.openfile_name_model                     self.log(记载自定义权重成功)                else :                     self.log(warning:权重文件有误,请重新加载)                                # 选择cpu 和gpu                self.device = select_device(self.opt.device)                self.half = self.device.type != cpu                cudnn.benchmark = True                # Load model                self.model = attempt_load(self.weights, map_location=self.device)                stride = int(self.model.stride.max())  # model stride                self.imgsz = check_img_size(imgsz, s=stride)  # check img_size                if self.half:                    self.model.half()  # to FP16                # Get names and colors                self.names = self.model.module.names if hasattr(self.model, module) else self.model.names                self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]                QtWidgets.QMessageBox.information(self, u"Notice", u"模型加载完成", buttons=QtWidgets.QMessageBox.Ok,                                      defaultButton=QtWidgets.QMessageBox.Ok)                self.log(模型加载完成)        def show_video_frame(self):                name_list = []                flag, img = self.cap.read()                if img is None:                                              self.releaseRes()                 else:                        info_show = self.detect(name_list, img)                        if self.opt.save:                                self.vid_writer.write(img) # 检测结果写入视频                        self.show_img(info_show, img)        def releaseRes(self):                        print(读取结束)                        self.log(检测结束)                        self.timer_video.stop()                        self.cap.release() # 释放video_capture资源                        self.ui.show.clear()                        if self.opt.save:                                self.vid_writer.release()                       def log(self, msg):                self.logging += %s%msg                self.ui.log.setText(self.logging)if __name__==__main__:    # 创建QApplication实例    app=QApplication(sys.argv)#获取命令行参数    current_ui = UI_Logic_Window()    current_ui.show()    sys.exit(app.exec_())

演示效果

新鲜出炉的 yoloV5可视化实战项目(1)

发表评论

  • 评论列表
还没有人评论,快来抢沙发吧~