检查yolo系列的分割数据,如果怀疑数据集有问题,可以使用此代码,dry_run 先设置为True,执行后会打印有问题的数据,观察打印的数据符合预期,然后,再改为False,删除异常数据。代码如下:
csharp
import os
import glob
def check_and_clean_segment_labels(label_dir, img_dir=None, num_classes=None, dry_run=False):
"""
检查 YOLO 分割标签格式,并删除异常标签及其对应的图像。
Args:
label_dir (str): 标签文件目录(如 'labels/train')
img_dir (str or None): 图像文件目录。若为 None,则自动推断(替换 'labels' 为 'images')
num_classes (int or None): 类别总数,用于检查 class_id 范围
dry_run (bool): 若为 True,仅打印将要删除的文件,不实际删除
"""
if img_dir is None:
# 自动推断 images 路径:将 'labels' 替换为 'images'
if 'labels' in label_dir:
img_dir = label_dir.replace('labels', 'images')
else:
raise ValueError("无法自动推断 img_dir,请显式传入 img_dir 参数")
label_paths = glob.glob(os.path.join(label_dir, "*.txt"))
if not label_paths:
print(f"⚠️ 警告:在 {label_dir} 中未找到任何 .txt 标签文件!")
return
# 支持的图像扩展名
IMG_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
invalid_files = []
for path in label_paths:
with open(path, 'r', encoding='utf-8') as f:
lines = [line.strip() for line in f.readlines() if line.strip()]
if not lines:
invalid_files.append((path, "空文件或无有效行"))
continue
valid_file = True
for i, line in enumerate(lines):
parts = line.split()
if len(parts) < 7:
invalid_files.append(
(path, f"第 {i + 1} 行:少于7个数值(当前{len(parts)}个),至少需要1个类别+3个点(7个数)"))
valid_file = False
break
try:
values = list(map(float, parts))
except ValueError:
invalid_files.append((path, f"第 {i + 1} 行:包含非数字字符"))
valid_file = False
break
class_id = int(values[0])
if num_classes is not None and (class_id < 0 or class_id >= num_classes):
invalid_files.append((path, f"第 {i + 1} 行:类别ID {class_id} 超出范围 [0, {num_classes - 1}]"))
valid_file = False
break
coords = values[1:]
if len(coords) % 2 != 0:
invalid_files.append((path, f"第 {i + 1} 行:坐标数量为奇数({len(coords)}),必须为偶数"))
valid_file = False
break
if not all(0 <= c <= 1 for c in coords):
invalid_files.append((path, f"第 {i + 1} 行:存在坐标超出 [0,1] 范围"))
valid_file = False
break
# 删除或预览异常文件
if invalid_files:
print(f"\n❌ 在 {label_dir} 中发现 {len(invalid_files)} 个异常标签文件:\n")
for txt_path, reason in invalid_files:
print(f" - {txt_path} → {reason}")
# 查找对应的图像文件
base_name = os.path.splitext(os.path.basename(txt_path))[0]
img_path = None
for ext in IMG_EXTS:
candidate = os.path.join(img_dir, base_name + ext)
if os.path.exists(candidate):
img_path = candidate
break
if dry_run:
print(f" 🗑️ [DRY RUN] 将删除: {txt_path}")
if img_path:
print(f" 🗑️ [DRY RUN] 将删除: {img_path}")
else:
print(f" ⚠️ 未找到对应图像文件(检查 {img_dir})")
else:
# 实际删除
try:
os.remove(txt_path)
print(f" ✅ 已删除标签: {txt_path}")
except Exception as e:
print(f" ❌ 删除失败 {txt_path}: {e}")
if img_path:
try:
os.remove(img_path)
print(f" ✅ 已删除图像: {img_path}")
except Exception as e:
print(f" ❌ 删除失败 {img_path}: {e}")
else:
print(f" ⚠️ 未找到对应图像文件(检查 {img_dir})")
else:
print(f"✅ {label_dir} 中所有标签文件格式合法!")
if __name__ == "__main__":
# ====== 配置你的路径 ======
train_label_dir = "datasets/Seg/labels/train"
val_label_dir = "datasets/Seg/labels/val"
num_classes = 1 # 你的类别数
# 设置 dry_run=True 先预览,确认无误后再设为 False 执行删除
dry_run = True # 👈 先设为 True 预览!确认后再改为 False
print("=== 检查并清理训练集 ===")
check_and_clean_segment_labels(
label_dir=train_label_dir,
num_classes=num_classes,
dry_run=dry_run
)
print("\n=== 检查并清理验证集 ===")
check_and_clean_segment_labels(
label_dir=val_label_dir,
num_classes=num_classes,
dry_run=dry_run
)