import os import cv2 import sys import xml.etree.ElementTree as ET from xml.dom import minidom from ultralytics import YOLO from PyQt5.QtWidgets import ( QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QLineEdit, QLabel, QFileDialog, QProgressBar, QTextEdit, QSlider, QCheckBox, QScrollArea, QGroupBox, QGridLayout ) from PyQt5.QtCore import Qt, QThread, pyqtSignal # ========== 后台任务线程 ========== class DetectionThread(QThread): update_rgb_progress = pyqtSignal(int) update_model_progress = pyqtSignal(int) update_xml_progress = pyqtSignal(int) log_message = pyqtSignal(str) finished_signal = pyqtSignal() def __init__(self, model_path, input_dir, output_xml_dir, temp_dir, conf, classes): super().__init__() self.model_path = model_path self.input_dir = input_dir self.output_xml_dir = output_xml_dir self.temp_dir = temp_dir self.conf = conf self.classes = classes # 类别映射 self.class_mapping = { 0: "EM14", 1: "EM18", 2: "EM17", 3: "EM170", 4: "EM19", 5: "EM190", 6: "EM20", 7: "EM200", 8: "EM201", 9: "EM202", 10: "EM203", 11: "EM180", 12: "EM181" } def run(self): try: os.makedirs(self.output_xml_dir, exist_ok=True) os.makedirs(self.temp_dir, exist_ok=True) # 加载模型 self.log_message.emit(f"加载模型：{self.model_path}") model = YOLO(self.model_path) self.log_message.emit("模型加载完成 ✅") # 读取图片 imgs = [f for f in os.listdir(self.input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))] total_imgs = len(imgs) if total_imgs == 0: self.log_message.emit("输入目录中未找到图片 ❌") self.finished_signal.emit() return # ========= 灰度转RGB阶段 ========= self.log_message.emit("开始灰度图像转RGB...此过程时间较长，请等待") for i, img_name in enumerate(imgs, 1): img_path = os.path.join(self.input_dir, img_name) img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) if img is None: continue if len(img.shape) == 2 or img.shape[2] == 1: img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) cv2.imwrite(os.path.join(self.temp_dir, img_name), img) self.update_rgb_progress.emit(int(i / total_imgs * 100)) self.log_message.emit("RGB 转换完成 ✅") # ========= 模型检测阶段 ========= self.log_message.emit("开始执行YOLO检测...") results = [] for i, file in enumerate(imgs, 1): temp_path = os.path.join(self.temp_dir, file) result = model.predict( source=temp_path, conf=self.conf, save=False, classes=self.classes if self.classes else None ) results.extend(result) self.update_model_progress.emit(int(i / total_imgs * 100)) self.log_message.emit(f"[检测完成] {file}") # ========= XML生成阶段 ========= self.log_message.emit("开始生成VOC格式XML...") total_results = len(results) for j, result in enumerate(results, 1): file_name = os.path.basename(result.path) image_path = os.path.join(self.input_dir, file_name) img = cv2.imread(image_path) if img is None: continue height, width, depth = img.shape annotation = ET.Element("annotation") folder = ET.SubElement(annotation, "folder") folder.text = os.path.basename(os.path.dirname(image_path)) filename = ET.SubElement(annotation, "filename") filename.text = file_name path = ET.SubElement(annotation, "path") path.text = os.path.abspath(image_path) source = ET.SubElement(annotation, "source") database = ET.SubElement(source, "database") database.text = "Unknown" size = ET.SubElement(annotation, "size") ET.SubElement(size, "width").text = str(width) ET.SubElement(size, "height").text = str(height) ET.SubElement(size, "depth").text = str(depth) segmented = ET.SubElement(annotation, "segmented") segmented.text = "0" for box in result.boxes: cls = int(box.cls[0]) xyxy = box.xyxy[0].tolist() xmin, ymin, xmax, ymax = map(int, xyxy) obj = ET.SubElement(annotation, "object") ET.SubElement(obj, "name").text = self.class_mapping.get(cls, str(cls)) ET.SubElement(obj, "pose").text = "Unspecified" ET.SubElement(obj, "truncated").text = "0" ET.SubElement(obj, "difficult").text = "0" bndbox = ET.SubElement(obj, "bndbox") ET.SubElement(bndbox, "xmin").text = str(xmin) ET.SubElement(bndbox, "ymin").text = str(ymin) ET.SubElement(bndbox, "xmax").text = str(xmax) ET.SubElement(bndbox, "ymax").text = str(ymax) xml_str = minidom.parseString(ET.tostring(annotation, 'utf-8')).toprettyxml(indent="\t", encoding="utf-8").decode("utf-8") xml_str = "\n".join([line for line in xml_str.splitlines() if line.strip()]) xml_path = os.path.join(self.output_xml_dir, os.path.splitext(file_name)[0] + ".xml") with open(xml_path, "w", encoding="utf-8") as f: f.write(xml_str) self.update_xml_progress.emit(int(j / total_results * 100)) self.log_message.emit(f"[生成完成] {xml_path}") self.log_message.emit("✅ 所有任务完成！") except Exception as e: self.log_message.emit(f"❌ 错误: {e}") finally: self.finished_signal.emit() # ========== 主界面 ========== class MainWindow(QWidget): def __init__(self): super().__init__() self.setWindowTitle("YOLO 半自动标注工具（华东专用）") self.resize(750, 750) layout = QVBoxLayout() # 路径部分 self.model_edit = self._add_path_selector(layout, "模型路径：") self.input_edit = self._add_path_selector(layout, "输入图片文件夹：") self.output_edit = self._add_path_selector(layout, "输出XML文件夹：") self.temp_edit = self._add_path_selector(layout, "临时RGB图片文件夹：") # 置信度 conf_layout = QHBoxLayout() conf_layout.addWidget(QLabel("置信度阈值：")) self.conf_slider = QSlider(Qt.Horizontal) self.conf_slider.setRange(0, 100) self.conf_slider.setValue(20) self.conf_value = QLabel("0.2") self.conf_slider.valueChanged.connect(lambda v: self.conf_value.setText(str(v / 100))) conf_layout.addWidget(self.conf_slider) conf_layout.addWidget(self.conf_value) layout.addLayout(conf_layout) # 类别选择区域 self.class_mapping = { 0: "EM14", 1: "EM18", 2: "EM17", 3: "EM170", 4: "EM19", 5: "EM190", 6: "EM20", 7: "EM200", 8: "EM201", 9: "EM202", 10: "EM203", 11: "EM180", 12: "EM181" } layout.addWidget(QLabel("选择要检测的类别（可多选）：")) self.class_checkboxes = {} class_group = QGroupBox() grid = QGridLayout() for idx, (key, name) in enumerate(self.class_mapping.items()): cb = QCheckBox(f"{key}: {name}") self.class_checkboxes[key] = cb grid.addWidget(cb, idx // 4, idx % 4) class_group.setLayout(grid) scroll = QScrollArea() scroll.setWidget(class_group) scroll.setWidgetResizable(True) scroll.setFixedHeight(120) layout.addWidget(scroll) # 控制按钮 btn_layout = QHBoxLayout() self.btn_load = QPushButton("加载模型") self.btn_start = QPushButton("开始检测") btn_layout.addWidget(self.btn_load) btn_layout.addWidget(self.btn_start) layout.addLayout(btn_layout) # 三个进度条 self.progress_rgb = QProgressBar() self.progress_rgb.setFormat("RGB转换进度：%p%") self.progress_model = QProgressBar() self.progress_model.setFormat("模型检测进度：%p%") self.progress_xml = QProgressBar() self.progress_xml.setFormat("XML生成进度：%p%") layout.addWidget(self.progress_rgb) layout.addWidget(self.progress_model) layout.addWidget(self.progress_xml) # 日志 layout.addWidget(QLabel("日志输出：")) self.log_output = QTextEdit() self.log_output.setReadOnly(True) layout.addWidget(self.log_output) self.setLayout(layout) # 信号绑定 self.btn_load.clicked.connect(self.load_model) self.btn_start.clicked.connect(self.start_detection) self.model_loaded = False def _add_path_selector(self, layout, label_text): hlayout = QHBoxLayout() label = QLabel(label_text) edit = QLineEdit() btn = QPushButton("浏览") hlayout.addWidget(label) hlayout.addWidget(edit) hlayout.addWidget(btn) layout.addLayout(hlayout) btn.clicked.connect(lambda: self._select_path(edit, "选择文件夹" if "文件夹" in label_text else "选择模型文件")) return edit def _select_path(self, edit, title): if "模型" in title: path, _ = QFileDialog.getOpenFileName(self, title, "", "Model Files (*.pt *.onnx)") else: path = QFileDialog.getExistingDirectory(self, title) if path: edit.setText(path) def load_model(self): path = self.model_edit.text().strip() if not os.path.exists(path): self._log("❌ 模型文件不存在") return self.model_loaded = True self._log("✅ 模型路径加载成功，可开始检测") def start_detection(self): if not self.model_loaded: self._log("⚠️ 请先加载模型路径") return model_path = self.model_edit.text().strip() input_dir = self.input_edit.text().strip() output_xml_dir = self.output_edit.text().strip() temp_dir = self.temp_edit.text().strip() conf = float(self.conf_slider.value() / 100) selected_classes = [cid for cid, cb in self.class_checkboxes.items() if cb.isChecked()] self._log(f"选中类别：{selected_classes if selected_classes else '全部'}") if not os.path.exists(input_dir): self._log("❌ 输入文件夹不存在") return if not temp_dir: self._log("⚠️ 请选择临时RGB图片文件夹") return self.progress_rgb.setValue(0) self.progress_model.setValue(0) self.progress_xml.setValue(0) self.thread = DetectionThread(model_path, input_dir, output_xml_dir, temp_dir, conf, selected_classes) self.thread.update_rgb_progress.connect(self.progress_rgb.setValue) self.thread.update_model_progress.connect(self.progress_model.setValue) self.thread.update_xml_progress.connect(self.progress_xml.setValue) self.thread.log_message.connect(self._log) self.thread.finished_signal.connect(self._finish) self.thread.start() self.btn_start.setEnabled(False) self._log("开始执行检测任务...") def _log(self, msg): self.log_output.append(msg) self.log_output.verticalScrollBar().setValue(self.log_output.verticalScrollBar().maximum()) def _finish(self): self.btn_start.setEnabled(True) self._log("任务完成 ✅") # ========== 主程序入口 ========== if __name__ == "__main__": app = QApplication(sys.argv) window = MainWindow() window.show() sys.exit(app.exec_())