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