YOLO自动标注工具软件 选择模型检测数据集,支持导出为不同格式的标签(txt,json,xml)附加功能如:调节标注阈值,日志记录检测结果预览
智慧标注:-YOLO自动标注工具软件
,功能全,响应快,操作简单!
选择模型检测你的数据集,支持导出为不同格式的标签(txt,json,xml)。
其他附加功能如:调节标注阈值,日志记录,检测结果预览等。
配合labelimg等标注工具能够快速的帮助你进行数据标注,节省人力成本。
1
1
这个 YOLO 自动标注工具,利用训练好的 YOLO 模型对未标注的图片进行推理,然后将推理得到的边界框坐标转换成标注文件(如 YOLO 格式的 .txt),从而替代人工进行初步标注。
界面搭建(PyQt5)、模型推理 和 标签转换 三个关键部分。
核心代码实现
环境依赖
你需要安装以下库来运行代码:
pip install PyQt5 ultralytics opencv-python torch numpy
完整代码 (auto_labeling_tool.py)
这是一个精简但功能完整的单文件版本,涵盖了你截图中的主要功能:模型加载、图片浏览、参数调节、自动推理和 YOLO 格式保存。
import sys
import os
from PyQt5.QtWidgets import(QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
QPushButton, QLabel, QSlider, QFileDialog, QProgressBar,
QMessageBox, QGroupBox, QFormLayout)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QThread, pyqtSignal
import cv2
import numpy as np
from ultralytics import YOLO
# --- 后台推理线程 ---
classInferenceThread(QThread):
progress = pyqtSignal(int,str)# 进度值, 日志信息
finished = pyqtSignal()
image_updated = pyqtSignal(np.ndarray)# 发送处理后的图像用于预览
def__init__(self):
super().__init__()
self.image_paths =[]
self.model =None
self.conf_thres =0.3
self.iou_thres =0.45
self.save_dir =""
self.is_running =False
defrun(self):
self.is_running =True
total =len(self.image_paths)
for i, img_path inenumerate(self.image_paths):
ifnot self.is_running:
break
# 1. 推理
results = self.model(img_path, conf=self.conf_thres, iou=self.iou_thres)
result = results[0]
# 2. 生成 YOLO 格式标签 (.txt)
img_name = os.path.basename(img_path)
name_without_ext = os.path.splitext(img_name)[0]
txt_path = os.path.join(self.save_dir,f"{name_without_ext}.txt")
h, w = result.orig_img.shape[:2]
withopen(txt_path,'w')as f:
for box in result.boxes:
cls_id =int(box.cls[0])
conf = box.conf[0]
# xyxy 转 xywh (归一化)
x1, y1, x2, y2 = box.xyxy[0]
x_center =((x1 + x2)/2)/ w
y_center =((y1 + y2)/2)/ h
width =(x2 - x1)/ w
height =(y2 - y1)/ h
f.write(f"{cls_id}{x_center:.6f}{y_center:.6f}{width:.6f}{height:.6f}\n")
# 3. 绘制图像用于预览 (在图片上画框)
annotated_img = result.plot()
# 发送信号
self.progress.emit(int((i +1)/ total *100),f"已处理: {img_name}")
if i ==0:# 只预览第一张或最后一张,避免界面卡顿
self.image_updated.emit(annotated_img)
self.finished.emit()
defstop(self):
self.is_running =False
# --- 主窗口界面 ---
classAutoLabelingTool(QMainWindow):
def__init__(self):
super().__init__()
self.setWindowTitle("YOLO自动标注工具")
self.setGeometry(100,100,1200,800)
self.model =None
self.image_paths =[]
self.current_idx =0
self.init_ui()
self.thread = InferenceThread()
definit_ui(self):
main_layout = QHBoxLayout()
# --- 左侧控制面板 ---
left_panel = QWidget()
left_layout = QVBoxLayout()
# 1. 模型设置
model_group = QGroupBox("模型设置")
model_layout = QFormLayout()
self.btn_load_model = QPushButton("加载模型")
self.btn_load_model.clicked.connect(self.load_model)
model_layout.addRow(self.btn_load_model)
model_group.setLayout(model_layout)
# 2. 参数设置
param_group = QGroupBox("检测参数")
param_layout = QFormLayout()
self.slider_conf = QSlider(Qt.Horizontal)
self.slider_conf.setRange(0,100)
self.slider_conf.setValue(30)
self.lbl_conf = QLabel("0.30")
self.slider_conf.valueChanged.connect(lambda: self.lbl_conf.setText(f"{self.slider_conf.value()/100:.2f}"))
self.slider_iou = QSlider(Qt.Horizontal)
self.slider_iou.setRange(0,100)
self.slider_iou.setValue(45)
self.lbl_iou = QLabel("0.45")
self.slider_iou.valueChanged.connect(lambda: self.lbl_iou.setText(f"{self.slider_iou.value()/100:.2f}"))
param_layout.addRow(QLabel("置信度:"), self.slider_conf)
param_layout.addRow(self.lbl_conf, QLabel(""))
param_layout.addRow(QLabel("IOU阈值:"), self.slider_iou)
param_layout.addRow(self.lbl_iou, QLabel(""))
param_group.setLayout(param_layout)
# 3. 文件设置
file_group = QGroupBox("文件设置")
file_layout = QFormLayout()
self.btn_select_input = QPushButton("选择输入文件夹")
self.btn_select_input.clicked.connect(self.select_input)
self.btn_select_output = QPushButton("选择输出文件夹")
self.btn_select_output.clicked.connect(self.select_output)
file_layout.addRow(self.btn_select_input)
file_layout.addRow(self.btn_select_output)
file_group.setLayout(file_layout)
# 4. 操作按钮
self.btn_start = QPushButton("开始标注")
self.btn_start.setStyleSheet("background-color: #4CAF50; color: white; height: 40px;")
self.btn_start.clicked.connect(self.start_labeling)
self.btn_stop = QPushButton("停止标注")
self.btn_stop.setStyleSheet("background-color: #F44336; color: white; height: 40px;")
self.btn_stop.clicked.connect(self.stop_labeling)
self.btn_stop.setEnabled(False)
# 5. 进度条
self.progress_bar = QProgressBar()
self.progress_bar.setValue(0)
left_layout.addWidget(model_group)
left_layout.addWidget(param_group)
left_layout.addWidget(file_group)
left_layout.addWidget(self.btn_start)
left_layout.addWidget(self.btn_stop)
left_layout.addWidget(QLabel("处理进度:"))
left_layout.addWidget(self.progress_bar)
left_panel.setLayout(left_layout)
# --- 右侧预览区域 ---
right_panel = QWidget()
right_layout = QVBoxLayout()
self.lbl_image = QLabel("预览区域")
self.lbl_image.setAlignment(Qt.AlignCenter)
self.lbl_image.setStyleSheet("background-color: #222; color: #fff;")
right_layout.addWidget(self.lbl_image)
right_panel.setLayout(right_layout)
main_layout.addWidget(left_panel,1)
main_layout.addWidget(right_panel,3)
container = QWidget()
container.setLayout(main_layout)
self.setCentralWidget(container)
# 连接线程信号
self.thread.progress.connect(self.update_progress)
self.thread.finished.connect(self.labeling_finished)
self.thread.image_updated.connect(self.show_image)
defload_model(self):
path, _ = QFileDialog.getOpenFileName(self,"选择YOLO模型文件","","Model Files (*.pt)")
if path:
try:
self.model = YOLO(path)
QMessageBox.information(self,"成功","模型加载成功!")
except Exception as e:
QMessageBox.critical(self,"错误",f"加载失败: {str(e)}")
defselect_input(self):
folder = QFileDialog.getExistingDirectory(self,"选择图片文件夹")
if folder:
# 支持常见的图片格式
self.image_paths =[os.path.join(folder, f)for f in os.listdir(folder)
if f.lower().endswith(('.png','.jpg','.jpeg','.bmp'))]
QMessageBox.information(self,"提示",f"找到 {len(self.image_paths)} 张图片")
defselect_output(self):
folder = QFileDialog.getExistingDirectory(self,"选择保存标签的文件夹")
if folder:
self.thread.save_dir = folder
defstart_labeling(self):
ifnot self.model:
QMessageBox.warning(self,"警告","请先加载模型")
return
ifnot self.image_paths:
QMessageBox.warning(self,"警告","请先选择输入图片文件夹")
return
ifnot self.thread.save_dir:
QMessageBox.warning(self,"警告","请先选择输出文件夹")
return
self.btn_start.setEnabled(False)
self.btn_stop.setEnabled(True)
# 配置线程参数
self.thread.image_paths = self.image_paths
self.thread.conf_thres = self.slider_conf.value()/100
self.thread.iou_thres = self.slider_iou.value()/100
self.thread.model = self.model
self.thread.start()
defstop_labeling(self):
self.thread.stop()
self.btn_stop.setEnabled(False)
defupdate_progress(self, val, log_msg):
self.progress_bar.setValue(val)
# 这里可以添加一个日志文本框来显示 log_msg
deflabeling_finished(self):
self.btn_start.setEnabled(True)
self.btn_stop.setEnabled(False)
QMessageBox.information(self,"完成","所有图片处理完毕!")
defshow_image(self, img_bgr):
# BGR 转 RGB
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
h, w, ch = img_rgb.shape
bytes_per_line = ch * w
qt_image = QImage(img_rgb.data, w, h, bytes_per_line, QImage.Format_RGB888)
# 自动缩放以适应标签大小
pixmap = QPixmap.fromImage(qt_image).scaled(self.lbl_image.size(), Qt.KeepAspectRatio)
self.lbl_image.setPixmap(pixmap)
if __name__ =='__main__':
app = QApplication(sys.argv)
window = AutoLabelingTool()
window.show()
sys.exit(app.exec_())
代码核心功能解析
-
多线程处理:
-
使用 QThread进行后台推理,防止在检测大量图片时 GUI 界面卡死(无响应)。 InferenceThread
类负责循环读取图片、推理、保存标签。 -
YOLO 格式转换:
-
代码中的核心算法部分是将 result.boxes.xyxy(像素坐标)转换为 YOLO 标准的归一化xywh(中心点+宽高):x c e n t e r = x 1 + x 2 2 × W i m g 
-
参数动态调节:
-
通过 QSlider实时获取置信度和 IOU 阈值,并在点击“开始”时传递给推理线程。 -
可视化预览:
-
利用 result.plot()获取绘制好边框的图像,转换为QImage显示在 PyQt 界面上,让你直观看到自动标注的效果。
如何运行
-
将上述代码保存为 main.py。 -
准备一个 .pt模型文件(如yolov8n.pt或你自己训练的模型)。 -
运行 python main.py。 -
在界面中加载模型,选择包含图片的文件夹和输出文件夹,点击“开始标注”即可。
夜雨聆风