Skip to content

Liwx1014/AutoDataSetsCompositor

BranchesTags

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
assets/kites_ai
assets/kites_ai
 
 
config
config
 
 
demo
demo
 
 
output
output
 
 
scripts
scripts
 
 
templates
templates
 
 
.gitignore
.gitignore
 
 
README.md
README.md
 
 
logo.jpeg
logo.jpeg
 
 
process.jpeg
process.jpeg
 
 
requirements.txt
requirements.txt
 
 
run.py
run.py
 
 

Repository files navigation

AutoDataSetsCompositor

可扩展图像合成数据集生成框架 - 支持任意类别配置的目标检测数据集扩充工具，通过配置文件和 Jinja2 模板支持任意类别的素材生成与合成。

效果演示

参考图 生成图
参考图 生成图(增加山火)
参考图 生成图(夜视)

特性

  • 任意类别支持 - 通过配置文件定义新类别，无需修改代码
  • 多种素材来源 - AI 生成、爬虫下载、手动绘制三种方式
  • 智能光照匹配 - 自动分析背景图光照条件，调整素材亮度/色温
  • 灵活合成模式 - CV 精确控制 + AI 边缘优化的混合模式
  • 提示词模板化 - 使用 Jinja2 模板管理 AI 提示词
  • 全局配置管理 - API 密钥、模型参数统一配置

快速开始

安装

# 克隆项目
git clone https://github.com/your-repo/AutoDataSetsCompositor.git
cd AutoDataSetsCompositor

# 安装依赖
pip install -r requirements.txt

# 配置 API Key，推荐使用火山引擎（用于 AI 生成模式）
export ARK_API_KEY="your-api-key"

基础用法

# 1. 生成素材（三选一）

# 方式A：手动绘制（快速测试，无需 API）
python run.py manual_assets --category kite --output-dir ./assets/kite

# 方式B：AI 生成（高质量）
# 使用模板: templates/{category}_generation.jinja2
# 变量来源: config/categories/{category}.json → generation 字段
# --count: 每种类型生成数量（默认 1）
python run.py ai_assets --category kite --output-dir ./assets/kite --count 3

# 方式C：爬虫下载（批量获取）
# --max: 每个关键词最大爬取数量（默认 20）
python run.py crawl_assets --category kite --output-dir ./assets/kite_crawled --max 50

# 2 支持三种模式生成数据

##  批量合成到背景图（CV 模式，默认）
# 处理所有背景图×素材组合
python run.py batch \
  --category kite \
  --background-dir ./backgrounds \
  --asset-dir ./assets/kite \
  --output-dir ./output

## 单张合成 - 两张参考图生图（hybrid 模式）
##    参考图1: 素材图  参考图2: 背景图
##    AI 将素材自然融合到背景中
##    模板: templates/{category}_composite_dual_ref.jinja2
##    --count: 生成数量（1-4，默认 1）
python run.py composite \
  --category kite \
  --background ./backgrounds/scene.jpg \
  --asset ./assets/kite/kite_red_diamond.png \
  --output ./output/result_hybrid.jpg \
  --mode hybrid \
  --count 3

## 单张合成 - 一张参考图生图（prompt 模式）
##    仅使用背景图作为参考，AI 根据提示词生成素材
##    模板: templates/{category}_composite_prompt.jinja2
##    --count: 生成数量（1-4，默认 1）
python run.py composite \
  --category kite \
  --background ./backgrounds/scene.jpg \
  --output ./output/result_prompt.jpg \
  --mode prompt \
  --count 3

## 单张合成 - 一张参考图 + 素材形状参考（prompt 模式增强）
##   背景图作为主参考，素材图仅作为形状/颜色参考
##   模板: templates/{category}_composite_prompt.jinja2
python run.py composite \
  --category kite \
  --background ./backgrounds/scene.jpg \
  --asset ./assets/kite/kite_red_diamond.png \
  --output ./output/result_prompt_ref.jpg \
  --mode prompt

## 批量合成 - 处理多张背景图（hybrid/prompt 模式）
##    使用 --background-dir 指定背景图目录
##    使用 --output-dir 指定输出目录
python run.py composite \
  --category kite \
  --background-dir ./backgrounds \
  --asset ./assets/kite/kite_red_diamond.png # hybrid 模式\  
  --output-dir ./output \
  --mode {hybrid or prompt} \
  --count 1

命令参数说明

素材生成命令

命令 参数 类型 默认值 说明
manual_assets --category str kite 类别名称
--output-dir str ./assets/{category} 输出目录
--size int 512 图像尺寸（像素）
ai_assets --category str kite 类别名称
--output-dir str ./assets/{category} 输出目录
--count int 1 每种类型生成数量
--api-key str 环境变量 API Key
crawl_assets --category str kite 类别名称
--output-dir str ./assets/{category}_crawled 输出目录
--max int 20 每个关键词最大爬取数
--keywords str 配置文件 自定义搜索关键词

图像合成命令

命令 参数 类型 默认值 说明
batch --background-dir str 必填 背景图目录
--asset-dir str 必填 素材目录
--output-dir str 必填 输出目录
--category str kite 类别名称
--mode str cv 合成模式：cv/hybrid
--ratio-min float 配置文件 素材最小占比
--ratio-max float 配置文件 素材最大占比
--api-key str 环境变量 API Key（hybrid 模式）
composite --background str - 背景图路径（单张模式）
--background-dir str - 背景图目录（批量模式）
--asset str - 素材路径（prompt 模式可选）
--output str - 输出路径（单张模式）
--output-dir str - 输出目录（批量模式）
--category str kite 类别名称
--mode str cv 合成模式：cv/hybrid/prompt
--ratio float 随机 素材占画面比例
--position str 随机 素材位置 x,y
--count int 1 生成数量（hybrid/prompt）
--api-key str 环境变量 API Key
--web-search flag 启用 启用联网搜索
--no-web-search flag - 禁用联网搜索

项目结构

AutoDataSetsCompositor/
├── config/
│   ├── global.json              # 全局配置（API、模型参数）
│   └── categories/              # 类别配置目录
│       └── {category}.json      # 类别配置文件
├── templates/                   # Jinja2 提示词模板
│   ├── {category}_generation.jinja2        # 素材生成提示词模板
│   ├── {category}_composite_dual_ref.jinja2  # 两张参考图生图模板
│   └── {category}_composite_prompt.jinja2    # 一张参考图生图模板
├── scripts/
│   ├── config_loader.py         # 配置加载器
│   ├── template_renderer.py     # 模板渲染器
│   ├── generate_assets.py       # AI 素材生成
│   ├── composite.py             # 图像合成核心
│   ├── batch_composite.py       # 批量合成
│   ├── crawl_assets.py          # 爬虫下载
│   ├── create_manual_assets.py  # 手动绘制
│   └── seedream_client.py       # Seedream API 客户端
├── assets/                      # 素材存放目录
├── output/                      # 输出目录
├── requirements.txt
└── run.py                       # 统一入口脚本

配置说明

全局配置 config/global.json

{
  "api": {
    "api_base": "https://ark.cn-beijing.volces.com/api/v3",
    "model": "doubao-seedream-5-0-260128",
    "api_key_env": "ARK_API_KEY"
  },
  "default_category": "kite",
  "templates_dir": "templates"
}

类别配置 config/categories/{category}.json

{
  "name": "kite",
  "display_name": "风筝",
  "description": "输电线路风筝异物检测数据集",
  "generation": {
    "colors": [
      {"name": "dark_red", "desc": "暗红色、褪色的红色"}
    ],
    "shapes": [
      {"name": "diamond", "desc": "传统菱形风筝"}
    ],
    "damage_states": [
      {"name": "light", "desc": "轻微褶皱"}
    ]
  },
  "composite": {
    "ratio_min": 0.003,
    "ratio_max": 0.008,
    "position_y_range": [0.10, 0.40],
    "position_x_range": [0.20, 0.80]
  },
  "templates": {
    "generation": "kite_generation.jinja2",
    "composite_dual_ref": "kite_composite_dual_ref.jinja2"
  },
  "crawl_keywords": [
    "kite isolated white background",
    "风筝 白色背景 素材"
  ]
}

类别配置参数说明

参数 类型 必须 默认值 说明
name string ✅ - 类别标识符，必须与文件名一致
display_name string ❌ {name} 显示名称，用于日志输出
description string ❌ "" 类别描述信息
generation object ❌ {} AI 素材生成配置
├─ colors array ❌ [{"name": "default"}] 颜色列表，每项含 name/desc
├─ shapes array ❌ [{"name": "default"}] 形状列表，每项含 name/desc
└─ damage_states array ❌ [{"name": "default"}] 状态列表，每项含 name/desc
composite object ❌ {} 图像合成配置
├─ ratio_min float ❌ 0.003 素材占背景面积最小比例
├─ ratio_max float ❌ 0.008 素材占背景面积最大比例
├─ position_y_range array ❌ [0.10, 0.40] Y轴位置范围 [min, max]
└─ position_x_range array ❌ [0.20, 0.80] X轴位置范围 [min, max]
templates object ❌ {} 提示词模板配置
├─ generation string ❌ {name}_generation.jinja2 素材生成模板
├─ composite_dual_ref string ❌ {name}_composite_dual_ref.jinja2 双参考图合成模板
└─ composite_prompt_driven string ❌ {name}_composite_prompt.jinja2 单参考图合成模板
crawl_keywords array ❌* [] 爬虫搜索关键词（*使用爬虫功能时必须配置）
lighting_adjustment object ❌ {} 光照调整参数（高级）

最小配置示例：只需提供 name 字段，其他均可缺省

{
  "name": "balloon"
}

添加新类别指南

  1. 创建类别配置文件 config/categories/balloon.json
  2. 创建对应的提示词模板 templates/balloon_*.jinja2
  3. 使用新类别：
python run.py ai_assets --category balloon --output-dir ./assets/balloon

合成模式参考

模式 说明 API 需求
cv 传统 CV 图像处理，精确控制尺寸和位置 不需要
hybrid CV 预处理 + AI 边缘融合优化 需要
prompt 纯提示词驱动，AI 完全控制 需要

工作流程

工作流程图

如何贡献

欢迎贡献代码、报告问题或提出新功能建议！

贡献流程

  1. Fork 本仓库 - 点击右上角 Fork 按钮
  2. 创建特性分支 - git checkout -b feature/your-feature-name
  3. 提交更改 - git commit -m 'Add some feature'
  4. 推送分支 - git push origin feature/your-feature-name
  5. 创建 Pull Request - 在 GitHub 上发起 PR

添加新类别模板

如果你想为项目贡献新的类别模板（如气球、塑料袋等）：

  1. 在 config/categories/ 下创建 {category}.json 配置文件
  2. 在 templates/ 下创建对应的 Jinja2 模板文件
  3. 提交 PR 时请附带模板效果示例

代码规范

  • 遵循 PEP 8 编码规范
  • 新功能请添加必要的注释
  • 确保现有测试通过

问题反馈

  • 使用 GitHub Issues 报告 Bug
  • 提供复现步骤和环境信息

License

MIT License

致谢

About

通用图像合成数据集生成框架 - 支持任意类别配置的目标检测数据集扩充工具

Resources

Readme
Activity

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages