"""cui 数据集分割脚本 1. 将train目录下的图片按70:15:15分割到train/val/test 2. 对样本量过少的类别进行标记 """ import os import shutil import random from pathlib import Path from typing import List, Tuple def split_dataset( dataset_root: str, train_ratio: float = 0.70, val_ratio: float = 0.15, test_ratio: float = 0.15, min_samples_threshold: int = 10, mark_insufficient: bool = True, max_total: int = 120, seed: int = 42 ): """ 分割数据集 Args: dataset_root: 数据集根目录（包含train目录） train_ratio: 训练集比例（默认0.70） val_ratio: 验证集比例（默认0.15） test_ratio: 测试集比例（默认0.15） min_samples_threshold: 最小样本数阈值，低于此值的类别会被标记（默认10） mark_insufficient: 是否在文件夹名称后添加"样本数太少"标记（默认True） max_total: 如果总数超过100，限制到max_total张（默认120） seed: 随机种子（默认42） """ # 验证比例 if abs(train_ratio + val_ratio + test_ratio - 1.0) > 0.001: raise ValueError(f"比例总和必须为1.0，当前为: {train_ratio + val_ratio + test_ratio}") dataset_root = Path(dataset_root) train_dir = dataset_root / 'train' val_dir = dataset_root / 'val' test_dir = dataset_root / 'test' if not train_dir.exists(): raise ValueError(f"训练目录不存在: {train_dir}") # 创建val和test目录 val_dir.mkdir(parents=True, exist_ok=True) test_dir.mkdir(parents=True, exist_ok=True) # 设置随机种子 random.seed(seed) print("="*80) print(f"数据集分割工具") print("="*80) print(f"数据集根目录: {dataset_root}") print(f"分割比例: 训练集={train_ratio*100:.0f}%, 验证集={val_ratio*100:.0f}%, 测试集={test_ratio*100:.0f}%") print(f"最小样本数阈值: {min_samples_threshold}") print(f"样本不足标记: {'启用' if mark_insufficient else '禁用'}") print("="*80 + "\n") # 检查val和test目录是否已有数据 val_has_data = val_dir.exists() and any(val_dir.iterdir()) test_has_data = test_dir.exists() and any(test_dir.iterdir()) if val_has_data or test_has_data: print("⚠ 检测到验证集或测试集已有数据") print("将检查所有类别的比例，如果不符合70:15:15将重新分配\n") # 获取所有类别文件夹 class_folders = [f for f in train_dir.iterdir() if f.is_dir()] total_classes = len(class_folders) print(f"找到 {total_classes} 个类别\n") stats = { 'total_classes': total_classes, 'processed_classes': 0, 'skipped_classes': 0, 'sufficient_samples': 0, 'insufficient_samples': 0, 'total_images': 0, 'train_images': 0, 'val_images': 0, 'test_images': 0, 'marked_classes': [] } # 处理每个类别 for class_folder in sorted(class_folders): class_name = class_folder.name # 收集所有图片（从train/val/test） train_images = [] for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: train_images.extend(list(class_folder.glob(ext))) val_class_dir = val_dir / class_name val_images = [] if val_class_dir.exists(): for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: val_images.extend(list(val_class_dir.glob(ext))) test_class_dir = test_dir / class_name test_images = [] if test_class_dir.exists(): for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: test_images.extend(list(test_class_dir.glob(ext))) # 合并所有图片 all_images = train_images + val_images + test_images num_images = len(all_images) stats['total_images'] += num_images if num_images == 0: print(f"⚠ {class_name}: 0张图片，跳过") continue # 检查当前比例是否正确 current_train = len(train_images) current_val = len(val_images) current_test = len(test_images) # 计算期望的数量 expected_train = int(num_images * train_ratio) expected_val = int(num_images * val_ratio) expected_test = num_images - expected_train - expected_val # 检查比例是否接近（允许1-2张的误差） ratio_correct = ( abs(current_train - expected_train) <= 2 and abs(current_val - expected_val) <= 2 and abs(current_test - expected_test) <= 2 ) # 如果总数超过100，即使比例正确也需要限制 if ratio_correct and num_images > 0 and num_images <= 100: # 比例正确且总数<=100，跳过 stats['skipped_classes'] += 1 is_insufficient = num_images < min_samples_threshold if is_insufficient: stats['insufficient_samples'] += 1 else: stats['sufficient_samples'] += 1 print(f"✓ {class_name}: 总数={num_images}, 比例正确 (train={current_train}, val={current_val}, test={current_test})，跳过") continue # 如果比例正确但总数>100，需要限制 if ratio_correct and num_images > 100: print(f"处理 {class_name}: 总数={num_images} > 100, 比例正确但需要限制到{max_total}张") stats['processed_classes'] += 1 # 随机选择max_total张 random.shuffle(all_images) image_files = all_images[:max_total] num_images = len(image_files) # 重新计算期望数量 expected_train = int(num_images * train_ratio) expected_val = int(num_images * val_ratio) expected_test = num_images - expected_train - expected_val print(f" 将重新分配为: train={expected_train}, val={expected_val}, test={expected_test}") elif not ratio_correct: # 比例不正确，需要重新分配 stats['processed_classes'] += 1 print(f"处理 {class_name}: 总数={num_images}, 当前比例不正确 (train={current_train}, val={current_val}, test={current_test})") # 如果总数超过100，限制到max_total if num_images > 100: print(f" 总数={num_images} > 100, 限制到{max_total}张") random.shuffle(all_images) image_files = all_images[:max_total] num_images = len(image_files) # 重新计算期望数量 expected_train = int(num_images * train_ratio) expected_val = int(num_images * val_ratio) expected_test = num_images - expected_train - expected_val else: image_files = all_images print(f" 将重新分配为: train={expected_train}, val={expected_val}, test={expected_test}") # 检查样本数是否足够 is_insufficient = num_images < min_samples_threshold if is_insufficient: stats['insufficient_samples'] += 1 if mark_insufficient: # 标记样本不足的类别 marked_name = f"{class_name}_样本数太少" stats['marked_classes'].append((class_name, num_images, marked_name)) else: stats['sufficient_samples'] += 1 # 确定最终的类别名称（如果需要标记） final_class_name = f"{class_name}_样本数太少" if (is_insufficient and mark_insufficient) else class_name # 如果还没打乱（总数<=100的情况），现在打乱 if num_images == len(all_images): random.shuffle(image_files) # 计算分割点 n_train = int(num_images * train_ratio) n_val = int(num_images * val_ratio) # 剩余的全部给test，确保所有图片都被分配 # 分割图片（重新分配） new_train_images = image_files[:n_train] new_val_images = image_files[n_train:n_train + n_val] new_test_images = image_files[n_train + n_val:] # 创建目标目录 train_target = train_dir / final_class_name val_target = val_dir / final_class_name test_target = test_dir / final_class_name train_target.mkdir(exist_ok=True) val_target.mkdir(exist_ok=True) test_target.mkdir(exist_ok=True) # 先移动图片到正确的位置（在删除之前） moved_to_train = 0 for img in new_train_images: if not img.exists(): continue target = train_target / img.name if img.parent != train_target: if target.exists(): target.unlink() try: shutil.move(str(img), str(target)) moved_to_train += 1 except Exception as e: print(f" ⚠ 移动失败 {img}: {e}") moved_to_val = 0 for img in new_val_images: if not img.exists(): continue target = val_target / img.name if img.parent != val_target: if target.exists(): target.unlink() try: shutil.move(str(img), str(target)) moved_to_val += 1 except Exception as e: print(f" ⚠ 移动失败 {img}: {e}") moved_to_test = 0 for img in new_test_images: if not img.exists(): continue target = test_target / img.name if img.parent != test_target: if target.exists(): target.unlink() try: shutil.move(str(img), str(target)) moved_to_test += 1 except Exception as e: print(f" ⚠ 移动失败 {img}: {e}") # 然后清理所有目录（删除不在新分配列表中的图片） # 收集所有原始图片路径 original_images = set(train_images + val_images + test_images) new_images = set(new_train_images + new_val_images + new_test_images) images_to_delete = original_images - new_images for img in images_to_delete: if img.exists(): try: img.unlink() except Exception as e: pass # 处理样本不足标记（如果需要） if mark_insufficient and is_insufficient: # 需要重命名train目录的文件夹 if class_folder != train_target: # 如果原文件夹还存在且不为空，需要移动剩余文件 if class_folder.exists() and any(class_folder.iterdir()): for item in class_folder.iterdir(): if item.is_file(): shutil.move(str(item), train_target / item.name) # 删除原文件夹（如果为空） try: if class_folder.exists() and not any(class_folder.iterdir()): class_folder.rmdir() except Exception as e: print(f" 警告: 删除原文件夹失败: {e}") # 更新统计 stats['train_images'] += len(new_train_images) stats['val_images'] += len(new_val_images) stats['test_images'] += len(new_test_images) # 打印进度 status = "⚠样本不足" if is_insufficient else "✓" print(f"{status} {class_name}: 总计={num_images}, 训练={len(new_train_images)}, 验证={len(new_val_images)}, 测试={len(new_test_images)}") # 打印总结 print("\n" + "="*80) print("分割完成!") print("="*80) print(f"总类别数: {stats['total_classes']}") print(f"已处理: {stats['processed_classes']} 个类别（比例不正确，已重新分配）") print(f"已跳过: {stats['skipped_classes']} 个类别（比例正确）") print(f"样本充足: {stats['sufficient_samples']} 个类别") print(f"样本不足: {stats['insufficient_samples']} 个类别") print(f"\n总图片数: {stats['total_images']}") print(f"训练集: {stats['train_images']} 张 ({stats['train_images']/stats['total_images']*100:.1f}%)") print(f"验证集: {stats['val_images']} 张 ({stats['val_images']/stats['total_images']*100:.1f}%)") print(f"测试集: {stats['test_images']} 张 ({stats['test_images']/stats['total_images']*100:.1f}%)") if stats['marked_classes']: print(f"\n已标记的类别（样本数 < {min_samples_threshold}）:") for old_name, count, new_name in stats['marked_classes']: print(f" {old_name} ({count}张) -> {new_name}") print("\n✓ 完成!") if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='数据集分割工具') parser.add_argument('--dataset', type=str, default='dataset/WholeIngredientRecognition', help='数据集根目录（默认: dataset/WholeIngredientRecognition）') parser.add_argument('--train-ratio', type=float, default=0.70, help='训练集比例（默认: 0.70）') parser.add_argument('--val-ratio', type=float, default=0.15, help='验证集比例（默认: 0.15）') parser.add_argument('--test-ratio', type=float, default=0.15, help='测试集比例（默认: 0.15）') parser.add_argument('--min-samples', type=int, default=10, help='最小样本数阈值，低于此值会被标记（默认: 10）') parser.add_argument('--no-mark', action='store_true', help='不标记样本不足的类别') parser.add_argument('--seed', type=int, default=42, help='随机种子（默认: 42）') parser.add_argument('--max-total', type=int, default=120, help='如果总数超过100，限制到max_total张（默认: 120）') args = parser.parse_args() # 获取项目根目录 script_dir = Path(__file__).parent project_root = script_dir.parent # 处理相对路径 if not os.path.isabs(args.dataset): dataset_path = project_root / args.dataset else: dataset_path = Path(args.dataset) split_dataset( dataset_root=str(dataset_path), train_ratio=args.train_ratio, val_ratio=args.val_ratio, test_ratio=args.test_ratio, min_samples_threshold=args.min_samples, mark_insufficient=not args.no_mark, max_total=args.max_total, seed=args.seed )