172 lines
6.6 KiB
Python
172 lines
6.6 KiB
Python
"""
|
||
统计每个类别的图片数量（包含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)
|