如何将YOLOv8训练终端输出导入Tkinter ScrolledText组件？

我用Tkinter开发YOLOv8模型训练的GUI，想把YOLO训练时的终端输出（比如预训练权重迁移、数据集扫描、训练进度这些信息）显示到ScrolledText组件里。试过sys重定向、线程、输出捕获这些方法，但只有自定义的print内容能显示在GUI里，YOLO本身的输出还是只在VS Code终端显示。

YOLO训练的终端输出示例 #

Transferred 355/355 items from pretrained weights
train: Scanning C:\---\train\labels.cache... 1977 images, 81 
val: Scanning C:\---\valid\labels.cache... 145 images, 8 back
Plotting labels to runs\detect\train17\labels.jpg... 
optimizer: AdamW(lr=0.002, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)
Image sizes 960 train, 960 val
Using 0 dataloader workers
Logging results to runs\detect\train17
Starting training for 10 epochs...
Closing dataloader mosaic

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
  0%|          | 0/124 [00:02<?, ?it/s]

当前ScrolledText组件代码 #

output_text = scrolledtext.ScrolledText(left_frame, wrap=tk.WORD)
output_text.grid(row=0, column=0, padx=10, pady=10, sticky="nsew")

当前训练函数代码 #

def modelTrain(): #function to train model
    if trainModel == 'none' or trainModel == '': #conditions to check user has selected everything required for training
        tk.messagebox.showwarning("No model set", "You must set a model for training on (PyTorch .pt file)")
        return #if conditions not met we show a warning and exit function
    if trainDataset == 'none' or trainDataset == '':
        tk.messagebox.showwarning("No dataset set", "You must set a dataset for training (YOLO .yaml file)")
        return
    if epochs == 0:
        tk.messagebox.showwarning("Number of epochs not set", "You must set the amount of epochs to train for")
        return
    if imgsize == 0:
        tk.messagebox.showwarning("Input image size not set", "You must set the size of the images as exported for training")
        return
    
    model = YOLO(trainModel) #setting model to one selected by user
    dataset = trainDataset #setting dataset to one selected by user

    def train_in_thread(): #training model in seperate thread to avoid gui freezing from model training using up sys resources 
        model.train(data=dataset, epochs=epochs, pretrained=True, imgsz=imgsize)#initializing model training with other vars set by user
        model.val()

        output_text.insert(tk.END, "Training completed.\n")
        output_text.see(tk.END)  # Scroll to the end

    # Create a separate thread to run the training function
    train_thread = threading.Thread(target=train_in_thread)
    train_thread.start()

解决方法 #

YOLOv8的部分输出可能直接写入系统标准输出/错误流，而非完全通过Python的sys.stdout，同时Tkinter组件只能在主线程更新，所以需要结合流重定向和线程安全的GUI更新来实现：

1. 自定义输出捕获类

继承io.StringIO重写write方法，将捕获到的内容通过Tkinter的after方法在主线程更新ScrolledText，同时处理YOLO进度条的\r换行符：

import io
import sys

class RedirectOutput(io.StringIO):
    def __init__(self, text_widget):
        super().__init__()
        self.text_widget = text_widget
        self.buffer = ""

    def write(self, message):
        # 处理进度条的\r换行，避免内容覆盖
        if "\r" in message:
            message = message.replace("\r", "\n")
        self.buffer += message
        # 用after将更新操作放到主线程事件循环
        self.text_widget.after(0, self.update_text)

    def update_text(self):
        if self.buffer:
            self.text_widget.insert(tk.END, self.buffer)
            self.text_widget.see(tk.END)
            self.buffer = ""

def restore_stdout_stderr(original_stdout, original_stderr):
    sys.stdout = original_stdout
    sys.stderr = original_stderr

2. 修改训练线程逻辑

在训练前替换sys.stdout和sys.stderr为自定义捕获类，训练结束后恢复原始流：

def modelTrain():
    if trainModel == 'none' or trainModel == '':
        tk.messagebox.showwarning("No model set", "You must set a model for training on (PyTorch .pt file)")
        return
    if trainDataset == 'none' or trainDataset == '':
        tk.messagebox.showwarning("No dataset set", "You must set a dataset for training (YOLO .yaml file)")
        return
    if epochs == 0:
        tk.messagebox.showwarning("Number of epochs not set", "You must set the amount of epochs to train for")
        return
    if imgsize == 0:
        tk.messagebox.showwarning("Input image size not set", "You must set the size of the images as exported for training")
        return
    
    model = YOLO(trainModel)
    dataset = trainDataset

    def train_in_thread():
        # 保存原始输出流
        original_stdout = sys.stdout
        original_stderr = sys.stderr
        # 初始化重定向器
        redirector = RedirectOutput(output_text)
        sys.stdout = redirector
        sys.stderr = redirector

        try:
            model.train(data=dataset, epochs=epochs, pretrained=True, imgsz=imgsize)
            model.val()
            # 训练完成提示也用线程安全方式更新
            output_text.after(0, lambda: output_text.insert(tk.END, "Training completed.\n"))
            output_text.after(0, lambda: output_text.see(tk.END))
        finally:
            # 恢复原始输出流，避免影响后续程序输出
            restore_stdout_stderr(original_stdout, original_stderr)

    train_thread = threading.Thread(target=train_in_thread)
    train_thread.start()

内容的提问来源于stack exchange，提问作者Cooper