【YOLO脚本】yolo格式数据集删除缺失删除图片和标签

两个变量记录图片和标签地址。统计数据集，并且删除没有标签的图片和没有图片的标签。

import os
import shutil
from collections import defaultdict

def analyze_and_clean_yolo_dataset(image_dir, label_dir, dry_run=True):
    """
    分析并清理YOLO格式数据集，删除无标签的图片和无图片的标签
    
    参数:
    image_dir (str): 图片文件夹路径
    label_dir (str): 标签文件夹路径
    dry_run (bool): 是否执行实际删除操作，True表示只显示统计信息不删除
    """
    # 定义支持的图片格式
    IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.JPG', '.JPEG', '.PNG', '.BMP'}
    
    # 收集图片和标签文件
    image_files = {}
    label_files = {}
    
    # 收集图片文件（不包含扩展名的文件名作为键）
    for filename in os.listdir(image_dir):
        name, ext = os.path.splitext(filename)
        if ext in IMAGE_EXTS:
            image_files[name] = filename
    
    # 收集标签文件（不包含扩展名的文件名作为键）
    for filename in os.listdir(label_dir):
        name, ext = os.path.splitext(filename)
        if ext == '.txt':
            label_files[name] = filename
    
    # 统计分析
    stats = defaultdict(int)
    stats['total_images'] = len(image_files)
    stats['total_labels'] = len(label_files)
    
    # 找出无标签的图片
    images_without_labels = []
    for name in image_files:
        if name not in label_files:
            images_without_labels.append(image_files[name])
    stats['images_without_labels'] = len(images_without_labels)
    
    # 找出无图片的标签
    labels_without_images = []
    for name in label_files:
        if name not in image_files:
            labels_without_images.append(label_files[name])
    stats['labels_without_images'] = len(labels_without_images)
    
    # 找出空标签文件
    empty_labels = []
    for name, filename in label_files.items():
        if name in image_files:  # 只检查有对应图片的标签
            label_path = os.path.join(label_dir, filename)
            try:
                if os.path.getsize(label_path) == 0:
                    empty_labels.append(filename)
            except OSError:
                continue
    stats['empty_labels'] = len(empty_labels)
    
    # 打印统计结果
    print("===== 数据集统计 =====")
    print(f"总图片数: {stats['total_images']}")
    print(f"总标签数: {stats['total_labels']}")
    print(f"无标签的图片数: {stats['images_without_labels']}")
    print(f"无图片的标签数: {stats['labels_without_images']}")
    print(f"空标签文件数: {stats['empty_labels']}")
    
    # 执行清理操作
    if not dry_run:
        print("\n===== 执行清理 =====")
        
        # 删除无标签的图片
        for filename in images_without_labels:
            file_path = os.path.join(image_dir, filename)
            try:
                os.remove(file_path)
                print(f"已删除图片: {filename}")
            except Exception as e:
                print(f"删除图片 {filename} 失败: {e}")
        
        # 删除无图片的标签
        for filename in labels_without_images:
            file_path = os.path.join(label_dir, filename)
            try:
                os.remove(file_path)
                print(f"已删除标签: {filename}")
            except Exception as e:
                print(f"删除标签 {filename} 失败: {e}")
        
        # 删除空标签文件
        for filename in empty_labels:
            file_path = os.path.join(label_dir, filename)
            try:
                os.remove(file_path)
                print(f"已删除空标签: {filename}")
            except Exception as e:
                print(f"删除空标签 {filename} 失败: {e}")
        
        print("\n清理完成!")
    else:
        print("\n=== 干运行模式，未执行实际删除操作 ===")
        print("若要执行清理，请将 dry_run 参数设置为 False")

# 使用示例
if __name__ == "__main__":
    # 请修改为你的图片和标签路径
    IMAGE_DIR = "images"
    LABEL_DIR = "labels"
    
    # 先进行干运行，查看统计信息
    analyze_and_clean_yolo_dataset(IMAGE_DIR, LABEL_DIR, dry_run=True)
    
    # 如果统计结果符合预期，可以取消注释下面的行执行实际清理
    # analyze_and_clean_yolo_dataset(IMAGE_DIR, LABEL_DIR, dry_run=False)