ControlNet自定义数据集制作：标注工具与格式规范-优快云博客

ControlNet自定义数据集制作：标注工具与格式规范

【免费下载链接】ControlNet Let us control diffusion models! 项目地址: https://gitcode.com/gh_mirrors/co/ControlNet

在使用ControlNet控制Stable Diffusion（SD）模型时，高质量的自定义数据集是训练效果的关键。本文将详细介绍如何制作符合ControlNet要求的标注数据集，包括标注工具选择、数据格式规范及验证流程，帮助你快速掌握从原始图像到训练数据的完整转换过程。

标注工具选择与使用

ControlNet支持多种控制类型（如边缘、关键点、分割等），项目内置了多种标注工具可直接用于数据预处理。这些工具位于annotator/目录下，涵盖从边缘检测到人体姿态估计的完整功能链。

核心标注模块

标注类型	工具路径	应用场景
边缘检测	annotator/canny/	提取物体轮廓边缘
人体姿态	annotator/openpose/	检测人体关键点
深度估计	annotator/midas/	生成深度图控制
语义分割	annotator/uniformer/	像素级区域标注
线条检测	annotator/mlsd/	检测直线和曲线结构

以边缘检测为例，Gradio交互工具gradio_canny2image.py展示了完整的Canny边缘检测流程：

from annotator.canny import CannyDetector
detector = CannyDetector()
control_image = detector(img, low_threshold=100, high_threshold=200)

上述代码片段来自项目内置的交互演示，通过调整阈值参数可控制边缘检测的敏感度，生成符合训练要求的控制图像。

可视化标注流程

标注工具的典型工作流程包括图像加载、特征提取和结果保存三个步骤。以下是使用HED边缘检测器处理图像的示例（代码改编自gradio_hed2image.py）：

from annotator.hed import HEDdetector
from annotator.util import resize_image, HWC3

# 加载并预处理图像
img = HWC3(cv2.imread("input.jpg"))
img = resize_image(img, 512)

# 生成边缘控制图
detector = HEDdetector()
control_img = detector(img)

# 保存结果
cv2.imwrite("control_edge.jpg", control_img)

处理前后的效果对比可参考项目测试图像：

原始图像：test_imgs/room.png
边缘标注结果：test_imgs/house_line.png

数据集格式规范

ControlNet训练要求特定的数据组织结构和文件格式。遵循规范的数据集能确保训练脚本正确解析控制信号与目标图像的对应关系。

文件组织结构

标准数据集应包含三个核心部分：控制图像（source）、目标图像（target）和提示文本（prompt）。推荐的目录结构如下（以Fill50K数据集为例）：

training/
└── fill50k/
    ├── prompt.json       # 包含所有样本的元数据
    ├── source/           # 控制图像目录
    │   ├── 00001.png
    │   └── ...
    └── target/           # 目标图像目录
        ├── 00001.png
        └── ...

这种结构与docs/train.md中推荐的组织方式一致，便于训练脚本./tutorial_train.py直接加载数据。

JSON元数据格式

prompt.json文件采用JSON Lines格式，每行包含一个样本的元数据：

{"source": "source/00001.png", "target": "target/00001.png", "prompt": "a red car in street"}
{"source": "source/00002.png", "target": "target/00002.png", "prompt": "a cat sitting on sofa"}

每个样本必须包含三个字段：

source: 控制图像路径（相对路径）
target: 目标图像路径（相对路径）
prompt: 文本描述（与目标图像内容匹配）

图像格式要求

属性	要求	说明
尺寸	512×512像素	可通过annotator/util.py中的`resize_image`函数统一处理
格式	PNG/JPG	推荐使用PNG保持图像质量
色彩模式	RGB	训练脚本会自动处理BGR转RGB（见tutorial_dataset.py第29-30行）
数据范围	控制图[0,1]，目标图[-1,1]	归一化代码见tutorial_dataset.py第33-36行

数据集制作实战

基于上述规范，我们可以通过三个步骤完成自定义数据集的制作：数据采集、自动标注和格式转换。

1. 数据采集与准备

首先收集原始图像对（如草图与真实图像、线稿与上色结果等），确保每对图像尺寸一致且内容对应。推荐使用项目测试图像集中的样本作为模板：

姿态控制样本：test_imgs/pose1.png
线稿控制样本：test_imgs/human_line.png

2. 批量标注脚本

以下是批量处理图像的示例脚本，使用Canny边缘检测器生成控制图像并构建数据集：

import os
import cv2
import json
from annotator.canny import CannyDetector
from annotator.util import HWC3, resize_image

# 初始化检测器
detector = CannyDetector()
input_dir = "raw_images/"
output_dir = "training/my_dataset/"

# 创建输出目录
os.makedirs(f"{output_dir}/source", exist_ok=True)
os.makedirs(f"{output_dir}/target", exist_ok=True)

# 处理所有图像
metadata = []
for i, filename in enumerate(os.listdir(input_dir)):
    if not filename.endswith((".png", ".jpg")):
        continue
        
    # 加载目标图像
    target_path = os.path.join(input_dir, filename)
    target_img = cv2.imread(target_path)
    
    # 生成控制图像
    source_img = HWC3(target_img)
    source_img = resize_image(source_img, 512)
    source_img = detector(source_img, 100, 200)  # 边缘检测
    
    # 保存结果
    source_filename = f"source/{i:05d}.png"
    target_filename = f"target/{i:05d}.png"
    cv2.imwrite(f"{output_dir}/{source_filename}", source_img)
    cv2.imwrite(f"{output_dir}/{target_filename}", target_img)
    
    # 添加元数据
    metadata.append({
        "source": source_filename,
        "target": target_filename,
        "prompt": f"a photo of {filename.split('.')[0]}"
    })

# 保存元数据文件
with open(f"{output_dir}/prompt.json", "w") as f:
    for item in metadata:
        f.write(json.dumps(item) + "\n")

3. 数据集验证

使用项目提供的测试脚本tutorial_dataset_test.py验证数据集格式正确性：

python tutorial_dataset_test.py

预期输出应包含：

50000  # 数据集样本数量
burly wood circle with orange background  # 随机样本的提示文本
(512, 512, 3)  # 目标图像尺寸
(512, 512, 3)  # 控制图像尺寸

若出现尺寸不匹配或文件找不到错误，需检查：

prompt.json中的路径是否与实际文件对应
图像文件是否统一调整为512×512像素
目录结构是否符合training/[dataset_name]/[source/target]规范

高级优化与最佳实践

标注质量提升技巧

参数调优：不同类型图像需调整检测参数，如MLSD线条检测可修改gradio_hough2image.py中的线段阈值
人工校对：关键样本建议人工检查，可参考test_imgs/pose2.png的姿态标注质量
数据增强：在MyDataset类中添加随机旋转、翻转等增强代码提升模型鲁棒性

性能优化建议

对于大规模数据集（>10k样本），建议：

使用annotator/midas/api.py中的批量处理接口
启用多线程预处理（修改tutorial_dataset.py中的DataLoader参数）
预处理结果缓存为.npy文件减少重复计算

总结与扩展

本文详细介绍了ControlNet数据集的制作流程，包括标注工具使用、格式规范和实战技巧。遵循这些指南，你可以构建高质量的自定义数据集，为训练特定控制类型的ControlNet模型奠定基础。

进阶探索方向：

尝试混合标注类型（如边缘+深度图）
研究docs/low_vram.md优化大数据集训练
探索tutorial_train_sd21.py支持的SD2.1模型适配

完整的训练流程请参考官方文档docs/train.md，祝你在ControlNet自定义训练中取得成功！

【免费下载链接】ControlNet Let us control diffusion models! 项目地址: https://gitcode.com/gh_mirrors/co/ControlNet

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考