如何解决无训练数据问题:一种更为智能化的解决方案
手动标注数据真的很费时间,而且买数据集又贵得要命,还不一定能完全符合你的需求。但这里有个令人兴奋的好消息,为啥不用 AI 来解决这个问题呢?
别再依赖传统方法了,你可以用像 LLM(大型语言模型)和图像生成器这样的 AI 工具,为你的特定目标创建合成训练数据。如今有那么多开源和商业的 AI 模型可供选择,你可以根据自己的需求随意搭配,无论是想控制预算、提高效率,还是追求高质量的结果,都能轻松搞定。这对研究和商业来说,简直就是一场变革!
这篇博客展示了我如何为训练目标检测模型生成定制的合成数据集,帮助开发者少花点时间在繁琐的数据准备上,多花点时间在真正解决问题上。
我的系统就像一条流水线,每一步都在前一步的基础上构建,最终生成完整的训练数据集。下面是一个高层次的概览,展示它是如何工作的:
- 首先,是我们的基础提示模板,描述你想检测的对象以及它可能出现的地方。
- 接下来提示模版多样化,利用 LLM根据你的基础模板发挥创意,生成大量不同的情景,用于你的目标检测模型。
- 这些详细的描述随后会被发送到像 Stability.ai 和 DALL-E 这样的图像生成模型。这些模型会将所有这些文本描述转化为实际的图像,用于训练。
- 我们可以使用 Hugging Face 或其他来源的预训练视觉模型来验证这些合成图像的质量。
- 一旦我们有了合成图像,像 Grounded-SAM 这样的工具就会自动对它们进行标注,在我们想要检测的对象周围添加边界框。
- 然后,我们使用预训练的视觉模型来验证这些标注的质量,检查边界框是否正确放置。
- 最终结果?一个完整的训练数据集,每张图像都附带了匹配的边界框标注,随时可以用来训练目标检测模型。
文章目录
- 导入库
- 基础提示模板
- 提示扩展
- 1\. 使用 LLMs API
- 2\. 使用网页界面
- 图像生成
- 验证生成的图像
- 标注图像
- 可视化标注图像
- 验证标注图像
- 结论
- 更上一层楼!
导入库
接下来,我们导入所需的库:
import os # 用于与操作系统交互
import math # 用于数学运算
import io # 用于文件输入和输出操作
import ast # 用于解析和计算 Python 表达式
import base64 # 用于 base64 编码和解码
from io import BytesIO # 用于在内存中读写文件import numpy as np # 用于数值运算
import pandas as pd # 用于数据操作和分析
import matplotlib.pyplot as plt # 用于绘图和可视化import cv2 # OpenCV 库,用于计算机视觉任务
from PIL import ImageDraw # 用于图像处理和绘制图形import torch # PyTorch,用于深度学习
from diffusers import StableDiffusionPipeline # 用于使用 Stable Diffusion 进行文本到图像的生成
from autodistill.detection import CaptionOntology # 用于目标检测中的标注/注释任务
from autodistill_grounding_dino import GroundingDINO # 用于定位和检测任务
from openai import OpenAI # OpenAI API,用于 AI 聊天
基础提示模板
我们在流水线中的第一步是创建一个基础提示模板。这是你描述你想检测的对象以及它可能出现的地方的地方。下面是一个用于检测不同环境中棕熊的基础提示模板示例:
# 定义每个生成提示中必须出现的重要对象。
important_objects = "brown bear" # 如果有多个对象,用逗号隔开,例如:"不同种类的熊,瓶子,……等。"# 指定要生成的提示数量。
number_of_prompts = 50 # 定义要为图像生成任务生成的提示数量。# 提供一个简短的描述,说明你希望这些提示描绘出什么样的图像。
description_of_prompt = "brown bear in different environments" # 描述图像生成的情景或上下文。# 生成一个格式化的指令集,用于生成图像生成提示。
# 这个格式化的字符串将有助于为计算机视觉模型创建详细且多样的提示。base_prompt = f'''
# 重要对象列表:
# 列在这里的对象必须出现在每个生成的提示中。
必须出现在每个提示中的重要对象:
{important_objects}# 输入详情:
# 任务是根据提供的描述生成特定数量的相关提示。
输入:
根据 {description_of_prompt} 生成 {number_of_prompts} 个逼真的图像生成提示。# 提示生成指令:
# - 每个提示都应该描绘出涉及这些对象的真实行为和情景。
# - 所有重要对象都必须出现在每个提示中。
# - 确保这些对象在相机的不同距离处都被捕捉到:
# - 从非常近的特写镜头到远处背景中的对象。
# - 提示应该多样化且详细,以涵盖广泛的用例。# 输出格式:
# - 输出应该是一个包含所有生成提示的 Python 列表。
# - 每个提示都应该用引号括起来,并在列表中用逗号隔开。
输出:
返回一个包含这些提示的 Python 列表,稍后用于训练计算机视觉模型。
[prompt1, prompt2, ...]
'''# (可选) - 打印用于生成提示的格式化指令集。
print(base_prompt)
提示扩展
接下来,我们使用文本生成模型扩展基础提示模板,为我们的目标检测模型提供各种不同的情景。主要有两种方法:
- 使用预训练的文本生成模型 API/本地,例如 deepseek,kimi,元宝,qwen等。
- 使用文本生成模型的网页界面,如 ChatGPT、Gemini、Claude 等。
1. 使用 LLMs API
要使用LLM 的 API,你需要注册一个 API 密钥。你可以利用一切可用的deepseek,kimi,豆包,元宝等等等。一旦你有了 API 密钥,就可以使用以下代码来扩展你的提示:
# 使用你的 API 密钥初始化 OpenAI API 客户端。
openai_chat = OpenAI(api_key="YOUR_OPENAI_API_KEY"
)# 使用 OpenAI API 生成图像生成提示。
completion = openai_chat.chat.completions.create(model="gpt-4o-mini",messages=[{"role": "system", "content": base_prompt+ "\n\n your response: [prompt1, prompt2, ...] and do not say anything else and i will be be using ast.literal_eval to convert the string to a list"}]
)# 从 API 响应中提取生成的提示。
response = completion.choices[0].message.content
由于 API 的输出是一个字符串,我们需要将其转换为字符串列表。
# 提取包含变量定义的部分
variable_definition = response.strip()# 修复格式问题,确保字符串是一个有效的 Python 列表
if variable_definition.endswith(","):variable_definition = variable_definition[:-1] + "]"# 使用 ast.literal_eval 安全地计算变量定义
prompts = ast.literal_eval(variable_definition)# 打印前几个提示,验证输出
print(prompts[0:5])#### 输出 ####
['A brown bear in a dense forest, ...','A brown bear walking alone in a ...',"A close-up shot of a brown bear's face ...",...
]
2. 使用网页界面
如果你更喜欢使用网页界面,你可以使用豆包、kimi/元宝、DeepSEEK等工具。这些工具提供了用户友好的界面,你可以输入你的提示并获得扩展后的输出。下面是一个示例,展示如何通过提供基础提示模板来使用 ChatGPT 扩展你的提示:
一旦你有了扩展后的提示,我们可以通过简单地重复这些提示来进一步增加数据集的大小。下面是一个示例,展示如何增加提示的数量:
# 通过将现有提示重复一次来增加提示的数量
prompts = prompts * 2# 打乱提示以确保随机性
import random
random.shuffle(prompts)# 打印重复和打乱后提示的总数
len(prompts)#### 输出 ####
100
图像生成
现在我们有了扩展后的提示,我们可以使用图像生成模型将这些文本描述转化为实际的图像。有许多图像生成模型可供选择,例如 DALL-E、Stability.ai 等。下面是一个示例,展示如何使用稳定扩散模型根据提示生成图像:
# 定义用于稳定扩散流水线的模型名称和设备。
model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"# 使用指定的模型和设备加载稳定扩散流水线。
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)
现在我们已经加载了图像生成模型,我们可以使用它根据我们的提示生成图像。下面是一个示例,展示如何根据提示生成图像:
# 提取前 5 个提示用于生成图像。(示例数据)
sample_prompts = prompts[:10]# 使用稳定扩散流水线根据示例提示生成图像。
images = pipe(sample_prompts