实践教程｜使用Stable Diffusion图像修复来生成自己的目的检测数据集

本文给出了一种网络场景数据的方法。 

深度学习模型需要大量的数据才能得到很好的结果，目的检测模型也是一样。
要训练一个YOLOv5的模型来自动检测你最喜欢的玩具，你需要拍几千张你的玩具在差别上下文中的照片，对于每张图，你需要标注玩具在图中的位置。
如许是非常耗时的。
本文提出了使用图像分割和stable diffusion来自动生成目的检测数据集的方法。

生成自定义数据集的pipeline
生成目的检测数据集的pipeline包罗4个步骤：

• 找一个和你要识别的物体属于相同实例的数据集（比如狗数据集）。
  
• 使用图像分割生成狗的mask。
  
• 微调图像修复Stable Diffusion模型。
  
• 使用Stable Diffusion图像修复模型和生成的mask来生成数据。
  

图像分割：生成mask图像

Stable Diffusion图像修复pipeline需要输入一个提示，一张图像和一张mask图像，这个模型会只从mask图像中的白色像素部分上去生成新的图像。
PixelLib这个库帮助我们来做图像分割，只用几行代码就可以，在这个例子里，我们会使用PointRend模型来检测狗，下面是图像分割的代码。

1. import pixellib
2. from pixellib.torchbackend.instance import instanceSegmentation
3. ​
4. ins = instanceSegmentation()
5. ins.load_model("pointrend_resnet50.pkl")
6. target_classes = ins.select_target_classes(dog=True)
7. results, output = ins.segmentImage(
8.   "dog.jpg",
9.   show_bboxes=True,
10.   segment_target_classes=target_classes,
11.   output_image_name="mask_image.jpg"
12. )

复制代码

使用pixellib来做图像分割
segmentImage 函数返回一个tuple:

• results : 是一个字典，包罗了 'boxes', 'class_ids', 'class_names', 'object_counts', 'scores', 'masks', 'extracted_objects'这些字段。
  
• output : 原始的图像和mask图像举行了混淆，如果show_bboxes 设置为True，还会有包围框。
  

生成mask图像

我们生成的mask只包罗白色和玄色的像素，我们的mask会比原来图中的狗略大一些，如许可以给Stable Diffusion充足的空间来举行修复。为了做到这种结果，我们将mask向左、右、上、下分别平移了10个像素。

1. from PIL import Image
2. import numpy as np
3. ​
4. width, height = 512, 512
5. image=Image.open("dog.jpg")
6. ​
7. # Store the mask of dogs found by the pointrend model
8. mask_image = np.zeros(image.size)
9. for idx, mask in enumerate(results["masks"].transpose()):
10.   if results["class_names"][idx] == "dog":
11.    mask_image += mask
12. ​
13. ​
14. # Create a mask image bigger than the original segmented image
15. mask_image += np.roll(mask_image, 10, axis=[0, 0]) # Translate the mask 10 pixels to the left
16. mask_image += np.roll(mask_image, -10, axis=[0, 0]) # Translate the mask 10 pixels to the right
17. mask_image += np.roll(mask_image, 10, axis=[1, 1]) # Translate the mask 10 pixels to the bottom
18. mask_image += np.roll(mask_image, -10, axis=[1, 1]) # Translate the mask 10 pixels to the top
19. ​
20. ​
21. # Set non black pixels to white pixels
22. mask_image = np.clip(mask_image, 0, 1).transpose() * 255
23. # Save the mask image
24. mask_image = Image.fromarray(np.uint8(mask_image)).resize((width, height))
25. mask_image.save("mask_image.jpg")

复制代码

从pixellib的输出生成图像的mask
现在，我们有了狗图像的原始图和其对应的mask。

使用pixellib基于狗的图像生成mask
微调Stable Diffusion图像修复pipeline

Dreambooth是微调Stable Diffusion的一种技术，我们可以使用很少的几张照片将新的概念教给模型，我们准备使用这种技术来微调图像修复模型。train_dreambooth_inpaint.py这个脚本中展示了怎样在你自己的数据集上微调Stable Diffusion模型。
微调需要的硬件资源

在单个24GB的GPU上可以使用gradient_checkpointing和mixed_precision来微调模型，如果要使用更大的batch_size 和更快的训练，需要使用至少30GB的GPU。
安装依靠

在运行脚本之前，确保安装了这些依靠:

1. pip install git+https://github.com/huggingface/diffusers.git
2. pip install -U -r requirements.txt

复制代码

并初始化加快环境：

1. accelerate config

复制代码

你需要注册Hugging Face Hub的用户，你还需要token来使用这些代码，运行下面的下令来授权你的token：

1. huggingface-cli login

复制代码

微调样本

在运行这些计算量很大的训练的时间，超参数微调很关键，需要在你跑训练的机器上尝试差别的参数，我保举的参数如下：

1. $ accelerate launch train_dreambooth_inpaint.py \
2.   --pretrained_model_name_or_path="runwayml/stable-diffusion-inpainting"  \
3.   --instance_data_dir="dog_images" \
4.   --output_dir="stable-diffusion-inpainting-toy-cat" \
5.   --instance_prompt="a photo of a toy cat" \
6.   --resolution=512 \
7.   --train_batch_size=1 \
8.   --learning_rate=5e-6 \  
9.   --lr_scheduler="constant" \  
10.   --lr_warmup_steps=0 \  
11.   --max_train_steps=400 \
12.   --gradient_accumulation_steps=2 \
13.   --gradient_checkpointing \
14.   --train_text_encoder

复制代码

运行Stable Diffusion图像修复pipeline

Stable Diffusion图像修复是一个text2image的扩散模型，使用一张带mask的图像和文本输入来生成真实的图像。使用GitHub - huggingface/diffusers:

继续阅读请点击广告