From f60c4d98832be282d23636f900ec27b77cd52add Mon Sep 17 00:00:00 2001 From: good Date: Fri, 6 Feb 2026 02:56:26 -0500 Subject: [PATCH] =?UTF-8?q?=E8=AE=A1=E7=AE=97=E7=85=A7=E7=89=87=E6=80=BB?= =?UTF-8?q?=E6=95=B0=E4=BB=A5=E5=8F=8A=E6=95=B0=E6=8D=AE=E5=BA=93=E5=88=86?= =?UTF-8?q?=E7=B1=BB=E6=83=85=E5=86=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_management/count_images.py | 171 ++++++++++++++++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 data_management/count_images.py diff --git a/data_management/count_images.py b/data_management/count_images.py new file mode 100644 index 0000000..a64927d --- /dev/null +++ b/data_management/count_images.py @@ -0,0 +1,171 @@ +""" +统计每个类别的图片数量（包含train/val/test） +""" +import os +from pathlib import Path +from collections import defaultdict + + +def count_images_by_class(dataset_root: str, output_file: str = None): + """ + 统计每个类别的图片数量（包含train/val/test） + + Args: + dataset_root: 数据集根目录（包含train/val/test目录） + output_file: 输出文件路径（可选，如果提供则保存到文件） + """ + 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 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='输出文件路径（可选，保存统计结果到文件）') + + 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)