文件夹图像自动标注工具：YOLO/CSV/TFRecord格式一键生成

1. 项目概述：文件夹图像自动标注工具

在计算机视觉项目中，最耗时的工作往往不是模型训练，而是数据准备阶段——特别是当你有成千上万张按文件夹分类的图片需要标注时。我最近开发了一个自动化脚本工具，能够根据文件夹结构自动为图像生成标签文件，将原本需要数天的手工标注工作缩短到几分钟完成。

这个工具特别适合以下场景：

• 已经按类别分好文件夹的图像数据集（如猫狗分类中所有猫图片在/cats文件夹）
• 需要转换为YOLO/TFRecord等格式的标准化标注文件
• 快速验证模型效果前的数据预处理阶段
• 处理监控摄像头按日期分类的素材归档

2. 核心功能设计思路

2.1 文件夹命名即标签

工具的核心逻辑建立在"文件夹名称=类别标签"的约定上。例如：

code复制dataset/
    ├── cat/
    │   ├── 001.jpg
    │   └── 002.jpg
    ├── dog/
        ├── 003.jpg
        └── 004.jpg

系统会扫描目录结构，自动生成对应的标签映射：

python复制{'cat': 0, 'dog': 1}

2.2 多格式输出支持

考虑到不同框架的需求差异，我们实现了三种主流标注格式：

1. CSV格式：兼容大多数数据分析工具

   csv复制filename,label
   /dataset/cat/001.jpg,cat
   

2. YOLO格式：生成txt标注文件与classes.txt

   code复制# classes.txt
   cat
   dog
   

3. TFRecord格式：直接生成TensorFlow标准数据集

3. 技术实现细节

3.1 目录遍历算法

采用递归方式处理嵌套文件夹结构：

python复制def scan_directory(root_path):
    for dirpath, dirnames, filenames in os.walk(root_path):
        if dirpath == root_path:  # 跳过根目录
            continue
        label = os.path.basename(dirpath)
        for filename in filenames:
            if filename.lower().endswith(('.png','.jpg')):
                yield label, os.path.join(dirpath, filename)

3.2 并行处理优化

当处理超过10万张图片时，采用多进程池加速：

python复制with Pool(processes=cpu_count()) as pool:
    results = pool.imap(process_image, file_paths)

4. 高级功能扩展

4.1 自动生成验证集

通过参数即可划分训练集/验证集：

bash复制python auto_label.py --split 0.2  # 20%作为验证集

4.2 图像预处理流水线

集成OpenCV实现自动处理：

python复制def preprocess(image_path):
    img = cv2.imread(image_path)
    img = cv2.resize(img, (256,256))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img

5. 实际应用案例

5.1 植物种类识别项目

处理来自科研机构的10万张植物图像：

code复制原始结构：
plant_dataset/
    ├── Rosaceae/
    ├── Liliaceae/
    └── Gramineae/

运行命令：
python auto_label.py --input plant_dataset --format yolo

5.2 工业质检应用

处理工厂传送带拍摄的产品缺陷图片：

python复制# 自定义标签过滤规则
def is_valid_label(label):
    return label in ['scratch', 'dent', 'normal']

6. 性能优化技巧

1. 内存映射技术：处理超大型数据集时使用numpy.memmap
2. 延迟加载：生成器替代列表存储文件路径
3. 增量处理：支持断点续处理功能

7. 常见问题解决方案

7.1 混合格式处理

当文件夹包含多种图像格式时：

python复制VALID_EXTENSIONS = ('.jpg','.jpeg','.png','.bmp','.tiff')

7.2 特殊字符处理

自动清理不合规的文件夹名称：

python复制import re
def clean_label(label):
    return re.sub(r'[^\w-]', '', label)

8. 部署与集成

8.1 Docker化部署

dockerfile复制FROM python:3.8
COPY auto_label.py /app/
ENTRYPOINT ["python", "/app/auto_label.py"]

8.2 作为Python包调用

python复制from auto_label import FolderLabeler
labeler = FolderLabeler(input_path='dataset')
labeler.generate_labels()

这个工具在实际项目中帮我节省了数百小时的人工标注时间。特别是在处理时间序列的监控数据时，只需按日期整理的文件夹结构，就能自动生成带时间戳的标注文件。对于需要快速验证模型效果的场景，这种自动化流程的价值更加凸显。