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

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

import os
import shutil
from collections import defaultdictdef 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:continuestats['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)