快捷導(dǎo)航

python yolo混合文件xml和img整理/回顯yolo標(biāo)注文件方式

更新時(shí)間：2026年01月21日 09:30:59 作者：像風(fēng)一樣的男人@

本文介紹了如何根據(jù)XML文件中的標(biāo)簽順序進(jìn)行索引轉(zhuǎn)換,并提供了一種自定義標(biāo)簽轉(zhuǎn)格式的方法,以回顯YOLO標(biāo)注文件

按xml中提取的標(biāo)簽順序轉(zhuǎn)索引

import os
import random
import time
from pathlib import Path
import shutil
import tkinter as tk
from tkinter import filedialog
from loguru import logger
import xml.etree.ElementTree as ET


class AnalysisXML(object):
    '''清洗xml'''

    def __init__(self):
        root = tk.Tk()
        root.withdraw()
        root.attributes('-topmost', 1)
        self.directory = filedialog.askdirectory()  # 打開(kāi)目錄選擇器
        root.destroy()
        logger.warning(f'路徑選擇:【{self.directory}】')

    def xml_img_split(self):
        '''分割圖片和xml'''
        logger.info(f'---------------------------分割圖片和xml------------------------')
        self.images_path = Path(self.directory).parent.joinpath('images')
        self.xml_labels_path = Path(self.directory).parent.joinpath('xml_labels')

        self.images_path.mkdir(parents=True, exist_ok=True)
        self.xml_labels_path.mkdir(parents=True, exist_ok=True)

        for i in Path(self.directory).iterdir():
            if i.suffix == '.xml':
                new_path = self.xml_labels_path.joinpath(i.name)
                logger.debug(f'移動(dòng):【{i}】 -> 【{new_path}】')
                shutil.copy(str(i), str(new_path))

            if i.suffix in ('.jpg', '.png'):
                new_path = self.images_path.joinpath(i.name)
                logger.debug(f'移動(dòng):【{i}】 -> 【{new_path}】')
                shutil.copy(str(i), str(new_path))

    def xml_to_txt(self):
        '''xml轉(zhuǎn)txt'''
        logger.info(f'----------------------------正在將xml轉(zhuǎn)為txt-----------------------')
        self.txt_labels = self.xml_labels_path.joinpath('labels')  # 替換為實(shí)際的輸出TXT文件夾路徑
        os.makedirs(self.txt_labels, exist_ok=True)

        names_set = set()
        for filename in os.listdir(self.xml_labels_path):
            if filename.endswith('.xml'):
                tree = ET.parse(os.path.join(self.xml_labels_path, filename))
                root = tree.getroot()

                for obj in root.findall('object'):
                    name = obj.find('name').text
                    names_set.add(name)
        # 輸出所有的name
        categories = []
        for name in names_set:
            categories.append(name)
        logger.success(f'標(biāo)注的內(nèi)容names:【{categories}】')

        category_to_index = {category: index for index, category in enumerate(categories)}

        # 遍歷輸入文件夾中的所有XML文件
        for filename in os.listdir(self.xml_labels_path):
            if filename.endswith('.xml'):
                xml_path = os.path.join(self.xml_labels_path, filename)
                logger.warning(f'正在處理:【{xml_path}】')
                # 解析XML文件
                tree = ET.parse(xml_path)
                root = tree.getroot()
                # 提取圖像的尺寸
                size = root.find('size')
                width = int(size.find('width').text)
                height = int(size.find('height').text)
                # 存儲(chǔ)name和對(duì)應(yīng)的歸一化坐標(biāo)
                objects = []
                # 遍歷XML中的object標(biāo)簽
                for obj in root.findall('object'):
                    name = obj.find('name').text
                    if name in category_to_index:
                        category_index = category_to_index[name]
                    else:
                        continue  # 如果name不在指定類(lèi)別中，跳過(guò)該object
                    bndbox = obj.find('bndbox')
                    xmin = int(bndbox.find('xmin').text)
                    ymin = int(bndbox.find('ymin').text)
                    xmax = int(bndbox.find('xmax').text)
                    ymax = int(bndbox.find('ymax').text)
                    # 轉(zhuǎn)換為中心點(diǎn)坐標(biāo)和寬高
                    x_center = (xmin + xmax) / 2.0
                    y_center = (ymin + ymax) / 2.0
                    w = xmax - xmin
                    h = ymax - ymin
                    # 歸一化
                    x = x_center / width
                    y = y_center / height
                    w = w / width
                    h = h / height
                    objects.append(f"{category_index} {x:.6f} {y:.6f} {w:.6f} {h:.6f}")
                # 輸出結(jié)果到對(duì)應(yīng)的TXT文件
                txt_filename = os.path.splitext(filename)[0] + '.txt'
                txt_path = os.path.join(self.txt_labels, txt_filename)
                with open(txt_path, 'w') as f:
                    for obj in objects:
                        f.write(obj + '\n')

    def to_dataset(self, test_ratio):
        '''整理為dataset'''
        output_folder = os.path.join(os.path.dirname(self.directory), 'datasets')

        input_image_folder = self.images_path
        input_label_folder = self.txt_labels

        train_images_folder = os.path.join(output_folder, 'train', 'images')
        train_labels_folder = os.path.join(output_folder, 'train', 'labels')
        val_images_folder = os.path.join(output_folder, 'val', 'images')
        val_labels_folder = os.path.join(output_folder, 'val', 'labels')

        os.makedirs(train_images_folder, exist_ok=True)
        os.makedirs(train_labels_folder, exist_ok=True)
        os.makedirs(val_images_folder, exist_ok=True)
        os.makedirs(val_labels_folder, exist_ok=True)

        # 獲取所有圖像文件列表
        images = [f for f in os.listdir(input_image_folder) if f.endswith('.jpg') or f.endswith('.png')]

        # 隨機(jī)打亂圖像文件列表
        random.shuffle(images)

        # 計(jì)算驗(yàn)證集的數(shù)量
        val_size = int(len(images) * test_ratio)

        # 劃分驗(yàn)證集和訓(xùn)練集
        val_images = images[:val_size]
        train_images = images[val_size:]

        # 復(fù)制驗(yàn)證集圖像和標(biāo)簽
        for image in val_images:
            label = os.path.splitext(image)[0] + '.txt'
            if os.path.exists(os.path.join(input_label_folder, label)):
                shutil.copy(os.path.join(input_image_folder, image), os.path.join(val_images_folder, image))
                shutil.copy(os.path.join(input_label_folder, label), os.path.join(val_labels_folder, label))
                logger.debug(
                    f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(val_images_folder, image)}】')
                logger.success(
                    f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(val_labels_folder, label)}】')
            else:
                logger.error(f"Warning: Label file {label} not found for image {image}")

        # 復(fù)制訓(xùn)練集圖像和標(biāo)簽
        for image in train_images:
            label = os.path.splitext(image)[0] + '.txt'
            if os.path.exists(os.path.join(input_label_folder, label)):
                shutil.copy(os.path.join(input_image_folder, image), os.path.join(train_images_folder, image))
                shutil.copy(os.path.join(input_label_folder, label), os.path.join(train_labels_folder, label))
                logger.debug(
                    f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(train_images_folder, image)}】')
                logger.success(
                    f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(train_labels_folder, label)}】')
            else:
                logger.error(f"Warning: Label file {label} not found for image {image}")

    def start(self):
        '''啟動(dòng)'''
        time.sleep(1)
        self.xml_img_split()
        time.sleep(1)
        self.xml_to_txt()
        time.sleep(1)
        self.to_dataset(0.2)


if __name__ == '__main__':
    base_dir = os.path.dirname(__file__)
    log_path = os.path.join(base_dir, 'log.log')
    if os.path.exists(log_path):
        os.unlink(log_path)

    logger.add(log_path)
    print('...第一層文件夾')
    print('     -->第二層文件夾↓')
    print('       -->[xml和img混合文件夾]')
    print('\n')
    status = input('請(qǐng)確認(rèn)xml和圖片在同一個(gè)文件夾(99:確認(rèn))(任意值:取消):')
    if status in (99, '99'):
        a = AnalysisXML()
        a.start()
        logger.success('系統(tǒng)完成')
        for i in (3, 2, 1):
            time.sleep(1)
            logger.success(f'{i}/秒')
    else:
        logger.error('系統(tǒng)退出!')
        for i in (3, 2, 1):
            time.sleep(1)
            logger.error(f'{i}/秒')

根據(jù)自定義標(biāo)簽轉(zhuǎn)格式(推薦)

import os
import random
import time
from pathlib import Path
import shutil
import tkinter as tk
from tkinter import filedialog, simpledialog
from loguru import logger
import xml.etree.ElementTree as ET


class AnalysisXML(object):
    '''清洗xml并按自定義標(biāo)簽列表轉(zhuǎn)換為YOLO格式'''

    def __init__(self):
        root = tk.Tk()
        root.withdraw()
        root.attributes('-topmost', 1)
        self.directory = filedialog.askdirectory()  # 打開(kāi)目錄選擇器
        # 新增：輸入自定義標(biāo)簽列表（用英文逗號(hào)分隔，如 cat,dog,car）
        self.custom_categories = simpledialog.askstring(
            "請(qǐng)輸入",
            '自定義標(biāo)簽列表,(空格逗號(hào)分隔);           輸入示例: cat dog car person',
            parent=root
        )
        root.destroy()

        # 處理輸入的標(biāo)簽列表
        if self.custom_categories:
            # 去空格、拆分、去重
            self.categories = [cat.strip() for cat in self.custom_categories.split() if cat.strip()]
            logger.warning(f'自定義標(biāo)簽列表:【{self.categories}】')
        else:
            logger.error('未輸入標(biāo)簽列表，程序退出！')
            exit(1)

        logger.warning(f'路徑選擇:【{self.directory}】')

    def xml_img_split(self):
        '''分割圖片和xml'''
        logger.info(f'---------------------------分割圖片和xml------------------------')
        self.images_path = Path(self.directory).parent.joinpath('images')
        self.xml_labels_path = Path(self.directory).parent.joinpath('xml_labels')

        self.images_path.mkdir(parents=True, exist_ok=True)
        self.xml_labels_path.mkdir(parents=True, exist_ok=True)

        for i in Path(self.directory).iterdir():
            if i.suffix == '.xml':
                new_path = self.xml_labels_path.joinpath(i.name)
                logger.debug(f'移動(dòng):【{i}】 -> 【{new_path}】')
                shutil.copy(str(i), str(new_path))

            if i.suffix in ('.jpg', '.png'):
                new_path = self.images_path.joinpath(i.name)
                logger.debug(f'移動(dòng):【{i}】 -> 【{new_path}】')
                shutil.copy(str(i), str(new_path))

    def xml_to_txt(self):
        '''xml轉(zhuǎn)txt（按自定義標(biāo)簽列表映射索引）'''
        logger.info(f'----------------------------正在將xml轉(zhuǎn)為txt-----------------------')
        self.txt_labels = self.xml_labels_path.joinpath('labels')  # 替換為實(shí)際的輸出TXT文件夾路徑
        os.makedirs(self.txt_labels, exist_ok=True)

        # 核心修改：用自定義標(biāo)簽列表生成索引映射，而非自動(dòng)無(wú)序生成
        category_to_index = {category: index for index, category in enumerate(self.categories)}
        logger.success(f'標(biāo)簽->索引映射:【{category_to_index}】')

        # 遍歷輸入文件夾中的所有XML文件
        for filename in os.listdir(self.xml_labels_path):
            if filename.endswith('.xml'):
                xml_path = os.path.join(self.xml_labels_path, filename)
                logger.warning(f'正在處理:【{xml_path}】')
                # 解析XML文件
                tree = ET.parse(xml_path)
                root = tree.getroot()
                # 提取圖像的尺寸
                size = root.find('size')
                width = int(size.find('width').text)
                height = int(size.find('height').text)
                # 存儲(chǔ)name和對(duì)應(yīng)的歸一化坐標(biāo)
                objects = []
                # 遍歷XML中的object標(biāo)簽
                for obj in root.findall('object'):
                    name = obj.find('name').text
                    # 檢查標(biāo)簽是否在自定義列表中，不在則跳過(guò)并警告
                    if name not in category_to_index:
                        logger.warning(f'XML【{filename}】中發(fā)現(xiàn)未定義標(biāo)簽【{name}】，已跳過(guò)！')
                        continue
                    category_index = category_to_index[name]

                    bndbox = obj.find('bndbox')
                    xmin = int(bndbox.find('xmin').text)
                    ymin = int(bndbox.find('ymin').text)
                    xmax = int(bndbox.find('xmax').text)
                    ymax = int(bndbox.find('ymax').text)
                    # 轉(zhuǎn)換為中心點(diǎn)坐標(biāo)和寬高
                    x_center = (xmin + xmax) / 2.0
                    y_center = (ymin + ymax) / 2.0
                    w = xmax - xmin
                    h = ymax - ymin
                    # 歸一化
                    x = x_center / width
                    y = y_center / height
                    w = w / width
                    h = h / height
                    objects.append(f"{category_index} {x:.6f} {y:.6f} {w:.6f} {h:.6f}")
                # 輸出結(jié)果到對(duì)應(yīng)的TXT文件
                txt_filename = os.path.splitext(filename)[0] + '.txt'
                txt_path = os.path.join(self.txt_labels, txt_filename)
                with open(txt_path, 'w') as f:
                    for obj in objects:
                        f.write(obj + '\n')

    def to_dataset(self, test_ratio):
        '''整理為dataset'''
        output_folder = os.path.join(os.path.dirname(self.directory), 'datasets')

        input_image_folder = self.images_path
        input_label_folder = self.txt_labels

        train_images_folder = os.path.join(output_folder, 'train', 'images')
        train_labels_folder = os.path.join(output_folder, 'train', 'labels')
        val_images_folder = os.path.join(output_folder, 'val', 'images')
        val_labels_folder = os.path.join(output_folder, 'val', 'labels')

        os.makedirs(train_images_folder, exist_ok=True)
        os.makedirs(train_labels_folder, exist_ok=True)
        os.makedirs(val_images_folder, exist_ok=True)
        os.makedirs(val_labels_folder, exist_ok=True)

        # 獲取所有圖像文件列表
        images = [f for f in os.listdir(input_image_folder) if f.endswith('.jpg') or f.endswith('.png')]

        # 隨機(jī)打亂圖像文件列表
        random.shuffle(images)

        # 計(jì)算驗(yàn)證集的數(shù)量
        val_size = int(len(images) * test_ratio)

        # 劃分驗(yàn)證集和訓(xùn)練集
        val_images = images[:val_size]
        train_images = images[val_size:]

        # 復(fù)制驗(yàn)證集圖像和標(biāo)簽
        for image in val_images:
            label = os.path.splitext(image)[0] + '.txt'
            if os.path.exists(os.path.join(input_label_folder, label)):
                shutil.copy(os.path.join(input_image_folder, image), os.path.join(val_images_folder, image))
                shutil.copy(os.path.join(input_label_folder, label), os.path.join(val_labels_folder, label))
                logger.debug(
                    f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(val_images_folder, image)}】')
                logger.success(
                    f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(val_labels_folder, label)}】')
            else:
                logger.error(f"Warning: Label file {label} not found for image {image}")

        # 復(fù)制訓(xùn)練集圖像和標(biāo)簽
        for image in train_images:
            label = os.path.splitext(image)[0] + '.txt'
            if os.path.exists(os.path.join(input_label_folder, label)):
                shutil.copy(os.path.join(input_image_folder, image), os.path.join(train_images_folder, image))
                shutil.copy(os.path.join(input_label_folder, label), os.path.join(train_labels_folder, label))
                logger.debug(
                    f'【{os.path.join(input_image_folder, image)}】 --> 【{os.path.join(train_images_folder, image)}】')
                logger.success(
                    f'【{os.path.join(input_label_folder, label)}】 --> 【{os.path.join(train_labels_folder, label)}】')
            else:
                logger.error(f"Warning: Label file {label} not found for image {image}")

    def start(self):
        '''啟動(dòng)'''
        time.sleep(1)
        self.xml_img_split()
        time.sleep(1)
        self.xml_to_txt()
        time.sleep(1)
        self.to_dataset(0.2)


if __name__ == '__main__':
    base_dir = os.path.dirname(__file__)
    log_path = os.path.join(base_dir, 'log.log')
    if os.path.exists(log_path):
        os.unlink(log_path)

    logger.add(log_path)
    print('...第一層文件夾')
    print('     -->第二層文件夾↓')
    print('       -->[xml和img混合文件夾]')
    print('\n')
    status = input('請(qǐng)確認(rèn)xml和圖片在同一個(gè)文件夾(99:確認(rèn))(任意值:取消):')
    if status in (99, '99'):
        a = AnalysisXML()
        a.start()
        logger.success('系統(tǒng)完成')
        for i in (3, 2, 1):
            time.sleep(1)
            logger.success(f'{i}/秒')
    else:
        logger.error('系統(tǒng)退出!')
        for i in (3, 2, 1):
            time.sleep(1)
            logger.error(f'{i}/秒')

回顯yolo標(biāo)注文件

import cv2
import os


def draw_yolo_annotation(image_path, txt_path, save_path=None):
    """
    在圖片上繪制YOLO標(biāo)注框和類(lèi)別索引
    :param image_path: 原圖路徑（如 test.jpg）
    :param txt_path: 對(duì)應(yīng)的YOLO標(biāo)注txt文件路徑（如 test.txt）
    :param save_path: 繪制后的圖片保存路徑，None則直接顯示
    """
    # 1. 讀取圖片
    img = cv2.imread(image_path)
    if img is None:
        print(f"錯(cuò)誤：無(wú)法讀取圖片 {image_path}")
        return
    h, w = img.shape[:2]  # 獲取圖片高、寬

    # 2. 讀取并解析YOLO標(biāo)注文件
    if not os.path.exists(txt_path):
        print(f"警告：標(biāo)注文件 {txt_path} 不存在，直接返回原圖")
        if save_path:
            cv2.imwrite(save_path, img)
        else:
            cv2.imshow("No Annotation", img)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        return

    with open(txt_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    # 3. 遍歷每個(gè)標(biāo)注框，轉(zhuǎn)換坐標(biāo)并繪制
    for line in lines:
        line = line.strip()
        if not line:
            continue

        # 解析標(biāo)注行：類(lèi)別索引、cx、cy、bw、bh
        parts = line.split()
        if len(parts) < 5:
            print(f"無(wú)效標(biāo)注行：{line}")
            continue

        class_idx = int(parts[0])
        cx = float(parts[1]) * w  # 歸一化→像素坐標(biāo)
        cy = float(parts[2]) * h
        bw = float(parts[3]) * w
        bh = float(parts[4]) * h

        # 計(jì)算框的左上角、右下角坐標(biāo)（YOLO中心點(diǎn)→OpenCV矩形框）
        x1 = int(cx - bw / 2)
        y1 = int(cy - bh / 2)
        x2 = int(cx + bw / 2)
        y2 = int(cy + bh / 2)

        # 4. 繪制矩形框+類(lèi)別索引文本
        # 框的樣式：綠色（BGR）、線(xiàn)寬2
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        # 文本背景（黑色半透明），避免文字被遮擋
        text = f"Class: {class_idx}"
        text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
        text_x = x1
        text_y = y1 - 10 if y1 - 10 > 10 else y1 + 20  # 避免文字超出圖片
        cv2.rectangle(img, (text_x, text_y - text_size[1] - 5),
                      (text_x + text_size[0] + 5, text_y + 5), (0, 0, 0), -1)
        # 繪制文本：白色、字體大小0.6、線(xiàn)寬2
        cv2.putText(img, text, (text_x + 2, text_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

    # 5. 保存/顯示結(jié)果
    if save_path:
        # 確保保存目錄存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        cv2.imwrite(save_path, img)
        print(f"標(biāo)注可視化完成，保存至：{save_path}")
    else:
        img = cv2.resize(img, (1500, 800))
        cv2.imshow("YOLO Annotation", img)
        cv2.waitKey(0)  # 按任意鍵關(guān)閉窗口
        cv2.destroyAllWindows()


# ===================== 測(cè)試調(diào)用 =====================
if __name__ == "__main__":
    # 替換為你的圖片和標(biāo)注文件路徑
    IMAGE_PATH = r"C:\Users\123\Desktop\22\控制室東_2026-01-16 15-22-23.249.jpg"  # 原圖路徑
    TXT_PATH = r"C:\Users\123\Desktop\22\控制室東_2026-01-16 15-22-23.249.txt"  # 對(duì)應(yīng)的YOLO標(biāo)注txt
    # SAVE_PATH = "test_anno.jpg"  # 繪制后的保存路徑（可選）

    # 方式1：直接顯示標(biāo)注后的圖片
    draw_yolo_annotation(IMAGE_PATH, TXT_PATH)

    # 方式2：保存標(biāo)注后的圖片（不顯示）
    # draw_yolo_annotation(IMAGE_PATH, TXT_PATH, SAVE_PATH)