""" 统计每个类别的图片数量（包含train/val/test），并可选择删除样本不足的类别 """ import os import shutil from pathlib import Path from collections import defaultdict def count_images_by_class(dataset_root: str, output_file: str = None, delete_threshold: int = None, auto_confirm: bool = False): """ 统计每个类别的图片数量（包含train/val/test），并可选择删除样本不足的类别 Args: dataset_root: 数据集根目录（包含train/val/test目录） output_file: 输出文件路径（可选，如果提供则保存到文件） delete_threshold: 删除阈值（可选，如果提供则删除总数少于该值的类别） auto_confirm: 是否自动确认删除（默认False，需要用户确认） """ 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(): print(f"错误: 训练目录不存在: {train_dir}") return print("="*80) print("每个类别的图片数量统计（包含train/val/test）") print("="*80) class_counts = {} class_counts_detail = {} # 详细统计：train/val/test分别的数量 total_images = 0 total_train = 0 total_val = 0 total_test = 0 # 获取所有类别文件夹（从train目录） class_folders = sorted([f for f in train_dir.iterdir() if f.is_dir()]) # 也检查val和test目录中是否有train中没有的类别 all_class_names = set() for folder in class_folders: all_class_names.add(folder.name) if val_dir.exists(): for folder in val_dir.iterdir(): if folder.is_dir(): all_class_names.add(folder.name) if test_dir.exists(): for folder in test_dir.iterdir(): if folder.is_dir(): all_class_names.add(folder.name) # 统计每个类别的图片数量 for class_name in sorted(all_class_names): train_count = 0 val_count = 0 test_count = 0 # 统计train目录 train_class_dir = train_dir / class_name if train_class_dir.exists(): for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: train_count += len(list(train_class_dir.glob(ext))) # 统计val目录 val_class_dir = val_dir / class_name if val_class_dir.exists(): for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: val_count += len(list(val_class_dir.glob(ext))) # 统计test目录 test_class_dir = test_dir / class_name if test_class_dir.exists(): for ext in ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG']: test_count += len(list(test_class_dir.glob(ext))) total = train_count + val_count + test_count class_counts[class_name] = total class_counts_detail[class_name] = { 'train': train_count, 'val': val_count, 'test': test_count, 'total': total } total_images += total total_train += train_count total_val += val_count total_test += test_count # 按数量排序 sorted_classes = sorted(class_counts.items(), key=lambda x: x[1], reverse=True) # 修复文件开头的注释 if output_file: # 确保sorted_classes已定义 pass # 打印结果 print(f"\n总类别数: {len(class_counts)}") print(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})") print(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n") print("-"*100) print(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}") print("-"*100) # 统计样本不足的类别（假设阈值是10） min_threshold = 10 insufficient_count = 0 for class_name, count in sorted_classes: detail = class_counts_detail[class_name] status = "⚠样本不足" if count < min_threshold else "✓" if count < min_threshold: insufficient_count += 1 print(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}") print("-"*100) print(f"\n样本不足的类别数（<{min_threshold}张）: {insufficient_count}") print(f"样本充足的类别数（>={min_threshold}张）: {len(class_counts) - insufficient_count}") # 删除样本不足的类别（如果指定了删除阈值） if delete_threshold is not None: classes_to_delete = [class_name for class_name, count in class_counts.items() if count < delete_threshold] if classes_to_delete: print("\n" + "="*100) print(f"发现 {len(classes_to_delete)} 个类别的总图片数少于 {delete_threshold} 张:") print("="*100) for class_name in classes_to_delete: detail = class_counts_detail[class_name] print(f" - {class_name}: {detail['total']} 张 (训练:{detail['train']}, 验证:{detail['val']}, 测试:{detail['test']})") # 确认删除 if not auto_confirm: confirm = input(f"\n是否删除这 {len(classes_to_delete)} 个类别? (yes/no): ").strip().lower() if confirm not in ['yes', 'y']: print("取消删除操作") return # 执行删除 deleted_count = 0 for class_name in classes_to_delete: try: # 删除train目录中的类别文件夹 train_class_dir = train_dir / class_name if train_class_dir.exists(): shutil.rmtree(train_class_dir) print(f"✓ 已删除: {train_class_dir}") # 删除val目录中的类别文件夹 val_class_dir = val_dir / class_name if val_class_dir.exists(): shutil.rmtree(val_class_dir) print(f"✓ 已删除: {val_class_dir}") # 删除test目录中的类别文件夹 test_class_dir = test_dir / class_name if test_class_dir.exists(): shutil.rmtree(test_class_dir) print(f"✓ 已删除: {test_class_dir}") deleted_count += 1 except Exception as e: print(f"✗ 删除失败 {class_name}: {e}") print(f"\n删除完成: 成功删除 {deleted_count}/{len(classes_to_delete)} 个类别") else: print(f"\n没有找到总数少于 {delete_threshold} 张的类别") # 保存到文件（如果指定） if output_file: output_path = Path(output_file) with open(output_path, 'w', encoding='utf-8') as f: f.write("="*100 + "\n") f.write("每个类别的图片数量统计（包含train/val/test）\n") f.write("="*100 + "\n\n") f.write(f"总类别数: {len(class_counts)}\n") f.write(f"总图片数: {total_images} (训练集: {total_train}, 验证集: {total_val}, 测试集: {total_test})\n") f.write(f"平均每类: {total_images // len(class_counts) if len(class_counts) > 0 else 0} 张\n\n") f.write("-"*100 + "\n") f.write(f"{'类别名称':<50} {'总计':<8} {'训练':<8} {'验证':<8} {'测试':<8} {'状态':<10}\n") f.write("-"*100 + "\n") for class_name, count in sorted_classes: detail = class_counts_detail[class_name] status = "⚠样本不足" if count < min_threshold else "✓" f.write(f"{class_name:<50} {detail['total']:<8} {detail['train']:<8} {detail['val']:<8} {detail['test']:<8} {status:<10}\n") f.write("-"*100 + "\n") f.write(f"\n样本不足的类别数（<{min_threshold}张）: {insufficient_count}\n") f.write(f"样本充足的类别数（>={min_threshold}张）: {len(class_counts) - insufficient_count}\n") print(f"\n✓ 统计结果已保存到: {output_path}") if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='统计每个类别的图片数量，并可选择删除样本不足的类别') parser.add_argument('--dataset', type=str, default='dataset/WholeIngredientRecognition', help='数据集根目录（默认: dataset/WholeIngredientRecognition）') parser.add_argument('--output', type=str, default=None, help='输出文件路径（可选，保存统计结果到文件）') parser.add_argument('--delete', type=int, default=None, help='删除阈值（可选，删除总数少于该值的类别，例如: --delete 20）') parser.add_argument('--yes', action='store_true', help='自动确认删除，不需要手动输入yes（谨慎使用）') 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) count_images_by_class(str(dataset_path), args.output, args.delete, args.yes)