基于YOLOv5n6和Tkinter实现的检测模型可视化界面

墨色流年 2024-06-28 ⋅ 27 阅读

YOLOv5 Logo

在计算机视觉领域中,目标检测是重要的任务之一,其可以通过算法自动识别并定位图像中的物体。YOLOv5 是最新推出的目标检测算法,具有高度精确度和快速的实时推理速度。而 Tkinter 是 Python 的一个标准库,用于创建图形用户界面(GUI)。将两者结合起来,我们可以实现一个基于 YOLOv5n6 和 Tkinter 的检测模型的可视化界面,方便用户直观地观察图像中的目标检测结果。

安装步骤

  1. 首先,确保你已经安装了 Python,推荐使用 Python 3版本。
  2. 接下来,我们需要安装 YOLOv5n6 和 Tkinter,可以使用以下命令安装依赖:
pip install yolov5
pip install tkinter
  1. 然后,我们需要下载 YOLOv5 项目的代码并导入到我们的项目中。

构建界面

使用 Tkinter 从零开始构建图形用户界面可能比较繁琐,但是我们可以使用 tkinter 的扩展库 PySimpleGUI 来简化这个过程。以下是一个简单的示例代码,用于构建一个基本的 Tkinter 窗口界面,并显示一张图像:

import PySimpleGUI as sg

layout = [[sg.Image(filename='', key='-IMAGE-')],
           [sg.Button('Open'), sg.Button('Detect'), sg.Button('Exit')]]

window = sg.Window('YOLOv5 Detection', layout,
                   grab_anywhere=True, finalize=True)

while True:
    event, values = window.read()
    if event in ('Exit', None):
        break
    elif event == 'Open':
        # 打开文件对话框,获取用户选择的图像文件并显示在窗口
        filename = sg.popup_get_file('Select an image file', no_window=True)
        if filename:
            window['-IMAGE-'].update(filename=filename)
    elif event == 'Detect':
        # 调用 YOLOv5 模型进行目标检测,并显示检测结果在窗口中
        # TODO: 调用 YOLOv5 模型的代码
        pass

window.close()

使用 YOLOv5n6 进行目标检测

使用 PyTorch 载入训练好的 YOLOv5 模型,并对图像进行目标检测。以下是一个简单的示例代码,用于加载模型并进行检测:

import torch
from PIL import Image

def detect(image_path, model_path):
    # 加载模型
    model = torch.hub.load('ultralytics/yolov5', 'yolov5n6', pretrained=False)
    model.load_state_dict(torch.load(model_path))
    
    # 加载图像
    img = Image.open(image_path)
    
    # 目标检测
    results = model(img)
    
    # 返回检测结果
    return results

# 调用 detect 函数进行目标检测
image_path = 'path/to/image.jpg'
model_path = 'path/to/model.pt'
results = detect(image_path, model_path)

在界面中显示检测结果

将 YOLOv5 模型的检测结果显示在 Tkinter 窗口中,可以使用以下代码将检测结果绘制到图像上,并更新到界面中:

import cv2
import PySimpleGUI as sg

def draw_boxes(image, results):
    for bbox in results.xyxy[0]:
        x1, y1, x2, y2, conf, cls = bbox
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, f'{cls} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)

# 在 Tkinter 窗口中显示图像和检测结果
while True:
    event, values = window.read()
    if event in ('Exit', None):
        break
    elif event == 'Open':
        # 打开文件对话框,获取用户选择的图像文件并显示在窗口
        filename = sg.popup_get_file('Select an image file', no_window=True)
        if filename:
            window['-IMAGE-'].update(filename=filename)

    elif event == 'Detect':
        # 获取当前显示的图像路径
        filename = window.Element('-IMAGE-').Filename
      
        # 调用 YOLOv5 进行目标检测
        results = detect(filename, model_path)

        # 将检测结果绘制到图像上
        image = cv2.imread(filename)
        draw_boxes(image, results)

        # 保存绘制结果
        output_path = 'path/to/output.jpg'
        cv2.imwrite(output_path, image)

        # 显示绘制结果
        window['-IMAGE-'].update(filename=output_path)

window.close()

总结

本文介绍了如何使用 YOLOv5n6 和 Tkinter 构建一个检测模型的可视化界面。通过结合 YOLOv5n6 目标检测算法和 Tkinter 图形用户界面库,我们可以方便地查看图像中的目标检测结果,并进行交互式操作。希望本文能帮助你理解如何实现这个功能,并进行进一步的应用和扩展。

参考链接:


全部评论: 0

    我有话说: