345 lines
16 KiB
Python
345 lines
16 KiB
Python
"""
|
||
统计每个类别的图片数量（包含train/val/test）
|
||
"""
|
||
import os
|
||
from pathlib import Path
|
||
from collections import defaultdict
|
||
from datetime import datetime
|
||
|
||
|
||
def export_to_excel(sorted_classes, class_counts_detail, class_counts,
|
||
total_images, total_train, total_val, total_test,
|
||
min_threshold, dataset_root: Path, output_dir: Path = None):
|
||
"""
|
||
将统计结果导出为 Excel 文件（带格式，方便粘贴到飞书）
|
||
|
||
Args:
|
||
sorted_classes: 按数量排序的类别列表
|
||
class_counts_detail: 每类的详细统计
|
||
class_counts: 每类的总数
|
||
total_images/train/val/test: 各分区汇总
|
||
min_threshold: 样本不足阈值
|
||
dataset_root: 数据集根目录（用于生成文件名）
|
||
output_dir: Excel 输出目录（默认和脚本同目录）
|
||
"""
|
||
try:
|
||
from openpyxl import Workbook
|
||
from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
|
||
from openpyxl.utils import get_column_letter
|
||
except ImportError:
|
||
print("❌ 未找到 openpyxl，请运行: pip install openpyxl")
|
||
return
|
||
|
||
wb = Workbook()
|
||
ws = wb.active
|
||
ws.title = "图片数量统计"
|
||
|
||
# ── 颜色 & 样式定义 ──────────────────────────────────────────
|
||
COLOR_HEADER_BG = "2563EB" # 深蓝色表头背景
|
||
COLOR_HEADER_FONT = "FFFFFF" # 白色表头字
|
||
COLOR_WARN_BG = "FEF08A" # 淡黄色：样本不足行
|
||
COLOR_SUMMARY_BG = "F1F5F9" # 浅灰：汇总行
|
||
COLOR_TOTAL_BG = "DBEAFE" # 浅蓝：总计行
|
||
|
||
def make_font(bold=False, color="000000", size=11):
|
||
return Font(bold=bold, color=color, size=size, name="微软雅黑")
|
||
|
||
def make_fill(hex_color):
|
||
return PatternFill("solid", fgColor=hex_color)
|
||
|
||
def make_border():
|
||
thin = Side(style="thin", color="CBD5E1")
|
||
return Border(left=thin, right=thin, top=thin, bottom=thin)
|
||
|
||
center = Alignment(horizontal="center", vertical="center")
|
||
left = Alignment(horizontal="left", vertical="center")
|
||
|
||
# ── 第1行：大标题 ────────────────────────────────────────────
|
||
ws.merge_cells("A1:F1")
|
||
title_cell = ws["A1"]
|
||
title_cell.value = f"📊 图片数量统计 — {dataset_root.name}"
|
||
title_cell.font = make_font(bold=True, color="1E3A5F", size=14)
|
||
title_cell.alignment = center
|
||
title_cell.fill = make_fill("EFF6FF")
|
||
ws.row_dimensions[1].height = 32
|
||
|
||
# ── 第2行：生成时间 ──────────────────────────────────────────
|
||
ws.merge_cells("A2:F2")
|
||
time_cell = ws["A2"]
|
||
time_cell.value = f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} 数据集: {dataset_root}"
|
||
time_cell.font = make_font(color="64748B", size=9)
|
||
time_cell.alignment = center
|
||
time_cell.fill = make_fill("EFF6FF")
|
||
ws.row_dimensions[2].height = 18
|
||
|
||
# ── 第3行：空行分隔 ──────────────────────────────────────────
|
||
ws.row_dimensions[3].height = 8
|
||
|
||
# ── 第4行：汇总信息 ──────────────────────────────────────────
|
||
summary_labels = [
|
||
("总类别数", len(class_counts)),
|
||
("总图片数", total_images),
|
||
("训练集", total_train),
|
||
("验证集", total_val),
|
||
("测试集", total_test),
|
||
("平均每类", total_images // len(class_counts) if len(class_counts) > 0 else 0),
|
||
]
|
||
insufficient_count = sum(1 for _, cnt in sorted_classes if cnt < min_threshold)
|
||
|
||
# 汇总用两行：标签行 + 数值行
|
||
for col_idx, (label, _) in enumerate(summary_labels, start=1):
|
||
cell = ws.cell(row=4, column=col_idx, value=label)
|
||
cell.font = make_font(bold=True, color="475569", size=10)
|
||
cell.alignment = center
|
||
cell.fill = make_fill(COLOR_SUMMARY_BG)
|
||
cell.border = make_border()
|
||
ws.row_dimensions[4].height = 20
|
||
|
||
for col_idx, (_, value) in enumerate(summary_labels, start=1):
|
||
cell = ws.cell(row=5, column=col_idx, value=value)
|
||
cell.font = make_font(bold=True, size=12)
|
||
cell.alignment = center
|
||
cell.fill = make_fill(COLOR_TOTAL_BG)
|
||
cell.border = make_border()
|
||
ws.row_dimensions[5].height = 22
|
||
|
||
# ── 第6行：空行 ───────────────────────────────────────────────
|
||
ws.row_dimensions[6].height = 8
|
||
|
||
# ── 第7行：列表头 ────────────────────────────────────────────
|
||
headers = ["类别名称", "总计", "训练集", "验证集", "测试集", "状态"]
|
||
for col_idx, header in enumerate(headers, start=1):
|
||
cell = ws.cell(row=7, column=col_idx, value=header)
|
||
cell.font = make_font(bold=True, color=COLOR_HEADER_FONT, size=11)
|
||
cell.fill = make_fill(COLOR_HEADER_BG)
|
||
cell.alignment = center
|
||
cell.border = make_border()
|
||
ws.row_dimensions[7].height = 24
|
||
|
||
# ── 数据行（从第8行开始）────────────────────────────────────
|
||
for row_offset, (class_name, count) in enumerate(sorted_classes):
|
||
row = 8 + row_offset
|
||
detail = class_counts_detail[class_name]
|
||
is_warn = count < min_threshold
|
||
status = "⚠ 样本不足" if is_warn else "✓"
|
||
row_data = [class_name, detail['total'], detail['train'],
|
||
detail['val'], detail['test'], status]
|
||
|
||
for col_idx, value in enumerate(row_data, start=1):
|
||
cell = ws.cell(row=row, column=col_idx, value=value)
|
||
cell.border = make_border()
|
||
cell.alignment = left if col_idx == 1 else center
|
||
cell.font = make_font(size=10)
|
||
if is_warn:
|
||
cell.fill = make_fill(COLOR_WARN_BG)
|
||
elif row_offset % 2 == 1:
|
||
cell.fill = make_fill("F8FAFC") # 斑马条纹
|
||
|
||
ws.row_dimensions[row].height = 18
|
||
|
||
# ── 最后一行：样本汇总说明 ───────────────────────────────────
|
||
last_row = 8 + len(sorted_classes)
|
||
ws.merge_cells(f"A{last_row}:F{last_row}")
|
||
note_cell = ws[f"A{last_row}"]
|
||
note_cell.value = (f"样本不足（< {min_threshold} 张）: {insufficient_count} 类 ｜ "
|
||
f"样本充足（≥ {min_threshold} 张）: {len(class_counts) - insufficient_count} 类")
|
||
note_cell.font = make_font(bold=True, color="475569", size=10)
|
||
note_cell.alignment = center
|
||
note_cell.fill = make_fill(COLOR_SUMMARY_BG)
|
||
note_cell.border = make_border()
|
||
ws.row_dimensions[last_row].height = 22
|
||
|
||
# ── 备注行：采样目标说明 ─────────────────────────────────────
|
||
remark_row = last_row + 1
|
||
ws.merge_cells(f"A{remark_row}:F{remark_row}")
|
||
remark_cell = ws[f"A{remark_row}"]
|
||
remark_cell.value = "现阶段目标：每个类别至少 100 张样本。后期可根据实际情况调整阈值。高质量样本越多，模型泛化能力越强。"
|
||
remark_cell.font = make_font(color="64748B", size=9)
|
||
remark_cell.alignment = center
|
||
remark_cell.fill = make_fill("EFF6FF")
|
||
remark_cell.border = make_border()
|
||
ws.row_dimensions[remark_row].height = 18
|
||
|
||
# ── 列宽自适应 ───────────────────────────────────────────────
|
||
col_widths = [40, 8, 8, 8, 8, 12]
|
||
for col_idx, width in enumerate(col_widths, start=1):
|
||
ws.column_dimensions[get_column_letter(col_idx)].width = width
|
||
|
||
# ── 冻结表头（冻结到第8行，滚动时列头始终可见）──────────────
|
||
ws.freeze_panes = "A8"
|
||
|
||
# ── 生成文件名 & 保存 ────────────────────────────────────────
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = f"image_count_{dataset_root.name}_{timestamp}.xlsx"
|
||
save_dir = output_dir if output_dir else Path(__file__).parent
|
||
save_path = save_dir / filename
|
||
|
||
wb.save(save_path)
|
||
print(f"\n✅ Excel 已导出: {save_path}")
|
||
print(f" （直接打开后全选复制，粘贴到飞书云文档即可保留格式）")
|
||
return save_path
|
||
|
||
|
||
def count_images_by_class(dataset_root: str, output_file: str = None,
|
||
export_excel: bool = False):
|
||
"""
|
||
统计每个类别的图片数量（包含train/val/test）
|
||
|
||
Args:
|
||
dataset_root: 数据集根目录（包含train/val/test目录）
|
||
output_file: 输出文件路径（可选，如果提供则保存到 txt 文件）
|
||
export_excel: 是否同时导出 Excel 文件（默认 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 = {}
|
||
total_images = total_train = total_val = total_test = 0
|
||
|
||
# 获取所有类别（以 train 为主，兼容 val/test 独有类别）
|
||
all_class_names = set()
|
||
for folder in train_dir.iterdir():
|
||
if folder.is_dir():
|
||
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)
|
||
|
||
image_exts = {'.jpg', '.jpeg', '.png'}
|
||
|
||
for class_name in sorted(all_class_names):
|
||
def _count(base_dir):
|
||
d = base_dir / class_name
|
||
if not d.exists():
|
||
return 0
|
||
return sum(1 for f in d.iterdir()
|
||
if f.is_file() and f.suffix.lower() in image_exts)
|
||
|
||
train_count = _count(train_dir)
|
||
val_count = _count(val_dir)
|
||
test_count = _count(test_dir)
|
||
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)
|
||
min_threshold = 100
|
||
insufficient_count = 0
|
||
|
||
# ── 控制台输出 ──────────────────────────────────────────────
|
||
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)
|
||
|
||
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} "
|
||
f"{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}")
|
||
|
||
# ── 可选：保存到 txt 文件 ────────────────────────────────────
|
||
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} "
|
||
f"{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}")
|
||
|
||
# ── 可选：导出 Excel ─────────────────────────────────────────
|
||
if export_excel:
|
||
export_to_excel(
|
||
sorted_classes = sorted_classes,
|
||
class_counts_detail = class_counts_detail,
|
||
class_counts = class_counts,
|
||
total_images = total_images,
|
||
total_train = total_train,
|
||
total_val = total_val,
|
||
total_test = total_test,
|
||
min_threshold = min_threshold,
|
||
dataset_root = dataset_root,
|
||
)
|
||
|
||
|
||
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='输出 txt 文件路径（可选）'
|
||
)
|
||
parser.add_argument(
|
||
'--excel', action='store_true',default= '--excel',
|
||
help='是否同时导出 Excel 文件（可直接复制粘贴到飞书云文档）'
|
||
)
|
||
|
||
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, export_excel=args.excel)
|