使用Blender进行分割缩放:自动创建数据集的指南

2025年01月23日 由 alex 发表 685 0

如果你曾经为新项目训练过细分模型,那么你可能知道,问题不在于模型,而在于数据。


收集图像通常很简单;主要挑战通常在于标记。对图像进行注释以进行分割非常耗时。即使使用Meta 的 SAM2等高级工具,创建完全注释、稳健且多样化的数据集仍需要相当长的时间。


在本文中,我们将探讨另一种通常较少被探索的选项:使用 3D 工具,例如Blender。事实上,3D 引擎越来越强大和逼真。此外,它们还具有一个引人注目的优势:能够在创建数据集时自动生成标签,无需手动注释。


在本文中,我们将概述创建手部分割模型的完整解决方案,分为以下几个关键部分:

  • 使用 Blender 生成手部以及如何获得手势、位置和肤色的多样性
  • 使用OpenCV使用生成的 Blender 图像和选定的背景图像生成数据集
  • 使用 PyTorch 训练和评估模型


生成手部图像

为了生成手部图像,让我们使用Blender。我虽然不是这类工具的专家,但它为我们提供了一些非常有用的功能:

  • 它是免费的——无需商业许可,任何人都可以立即下载并使用
  • 有一个庞大的社区,网上可以找到许多模型,有些免费,有些收费
  • 最后,它包含Python API,能够自动化生成具有各种特征的图像


如我们所见,这些功能相当有用,并且将使我们能够相当容易地制作合成数据。为了确保足够的多样性,我们将探索如何在生成的手部图像中自动随机化以下参数:

  • 手指位置:我们想要获得许多不同位置的手部图像
  • 相机位置:我们想要从各种视角获得手部图像
  • 肤色:我们想要肤色的多样性,以使模型足够稳健


注:此处提出的方法并非完全无潜在肤色偏见,也不声称无偏见。基于此方法的任何产品都必须针对任何伦理偏见进行仔细评估。


在深入这些步骤之前,我们需要一个手部的3D模型。像Turbosquid这样的网站上有许多模型,但我使用了一个可以在此处免费获取的手部模型。如果你用Blender打开这个文件,你会得到类似下面的截图。


7


如图所示,该模型不仅包括手部的形状和纹理,还包括骨骼结构,从而能够模拟手部的运动。让我们以此为基础,通过调整手指位置、肤色和相机位置,来获取一组多样化的手部图像。


修改手指位置

第一步是确保获得一组多样化且逼真的手指位置。为了不涉及过多细节(因为这更多与Blender本身相关),我们需要为运动创建控制器,并对允许的运动施加约束。基本上,我们不希望手指向后折叠或向不真实的方向弯曲。


一旦Blender文件设置了正确的约束,我们就可以使用Python脚本来自动化任何手指位置:


def randomize_fingers(armature, mini: float=-0.25, maxi: float=0.05):
    bpy.context.view_layer.objects.active = armature
    bpy.ops.object.mode_set(mode='POSE')
    for bone in armature.pose.bones:
        if "controller" in bone.name.lower(): # Adjust this condition to match your finger bones naming
            new_location = Vector((random.uniform(mini, maxi), random.uniform(mini, maxi), random.uniform(mini, maxi)))
            bone.location = new_location


如我们所见,我们所做的只是随机更新控制器的位置,从而在约束条件下移动手指。在正确的约束条件下,我们得到的手指位置如下所示:


8


这生成了逼真且多样的手指位置,最终使得能够生成一组多样化的手部图像。现在,让我们来调整肤色。


修改肤色

在创建包含人物的新图像数据集时,最具挑战性的方面之一可能是实现足够广泛的肤色表示。确保模型在所有肤色上都能无偏见地高效工作是一个至关重要的优先事项。虽然我并不声称能够消除任何偏见,但我在此提出的方法允许通过自动改变肤色来提供一个权宜之计。


注:此方法并不声称能够使模型免受任何伦理偏见的影响。任何用于生产的模型都必须经过公平性评估的仔细测试。可以参考谷歌为其人脸检测模型所做的工作作为示例。


我在这里所做的是对图像进行纯粹的图像处理计算。想法很简单:给定一个目标颜色和渲染出的手的平均颜色,我只需计算这两种颜色之间的差异。然后,我将这个差异应用到渲染出的手上,以得到新的肤色:


def change_color(img, mask, target_color):
    # get average color of hand
    average_color = cv2.mean(img, mask=mask)[:3]
    # compute difference colors and make into an image the same size as input
    diff_color = (target_color - average_color).astype(np.int32)
    diff_color = np.full_like(img, diff_color, dtype=np.int32)
    # shift input image color
    output_img = (img + diff_color).clip(0, 255).astype(np.uint8)
    # antialias mask, convert to float in range 0 to 1 and make 3-channels
    facemask = cv2.GaussianBlur(mask, (0, 0), sigmaX=3, sigmaY=3, borderType=cv2.BORDER_DEFAULT)
    facemask = skimage.exposure.rescale_intensity(
        facemask, in_range=(100, 150), out_range=(0, 1)
    ).astype(np.float32)
    facemask = cv2.merge([facemask, facemask, facemask])
    # combine img and output_img using mask
    result = (img * (1 - facemask) + output_img * facemask)
    result = result.clip(0, 255).astype(np.uint8)
    result = np.concatenate([result, np.expand_dims(mask, -1)], -1)
    return result


结果,它生成了以下手部图像:


9


虽然结果并不完美,但通过简单的图像处理,它们生成了具有多样肤色的相当逼真的图像。为了获得足够多样化的图像集,只剩最后一步:渲染视角。


修改相机位置

最后,让我们调整相机位置,以从多个视角捕捉手部。为此,相机被放置在以手部为中心的一个球体上的随机点上。这可以通过简单地调整球面坐标的两个角度来轻松实现。在以下代码中,我生成了球体上的一个随机位置:


def random_point_on_sphere(radius: float):
    theta = random.uniform(0, 2*math.pi)
    phi = random.uniform(0, math.pi)
    x = radius * math.sin(phi) * math.cos(theta)
    y = radius * math.sin(phi) * math.sin(theta)
    z = radius * math.cos(phi)
    return Vector((x, y, z))


然后,利用这一点,并在球面位置上添加一些约束,我就可以在Blender中更新手部周围的相机位置:


def randomize_camera(camera, target, radius):
    camera.location = target.location + random_point_on_sphere(radius)
    # Add a track_to constraint to the camera
    if "Track To" not in camera.constraints:
        constraint = camera.constraints.new(type='TRACK_TO')
        constraint.target = target
        constraint.track_axis = 'TRACK_NEGATIVE_Z'
        constraint.up_axis = 'UP_Y'
    else:
        camera.constraints["Track To"].target = target


结果,我们现在得到了以下图像样本:


10


现在,我们拥有了手指位置、肤色各异,且从不同视角拍摄的手部图像。在训练分割模型之前,下一步实际上是生成各种背景和情境中的手部图像。


生成训练数据

为了生成足够多样且逼真的图像,我们将把我们生成的手部与一组选定的背景图像进行融合。


我从Unsplash上选取了一些无版权的图像作为背景图像。我确保这些图像中没有手。然后,我将随机把这些Blender生成的手部添加到这些背景图像中:


def generate_composite_image(hand_folder: str, background_folder: str, hand_target_size: int, background_resize_scale: float):
    """Generate a composite image with a random hand on a random background."""
    # Load random hand and background images
    hand_image = load_random_image(hand_folder, target_size=hand_target_size)
    hand_mask = np.expand_dims(hand_image[:,:,3].astype(bool), -1)
    hand_image = (hand_image[:, :, :3]*255).astype(np.uint8)
    background_image = load_random_image(background_folder)
    resize_scale = random.uniform(background_resize_scale, 1.)
    resize_x = max(hand_image.shape[1], int(resize_scale*background_image.shape[1]))
    resize_y = max(hand_image.shape[0], int(resize_scale*background_image.shape[0]))
    background_image = cv2.resize(background_image, (resize_x, resize_y))
    # Generate random position for hand placement
    max_x = background_image.shape[1] - hand_image.shape[1]
    max_y = background_image.shape[0] - hand_image.shape[0]
    x = random.randint(0, max_x)
    y = random.randint(0, max_y)
    # Create output mask
    mask = np.zeros(background_image.shape[:2], dtype=np.uint8)
    mask[y:y + hand_image.shape[0], x:x + hand_image.shape[1]] = hand_mask.squeeze(-1)
    # Create blended image
    blended_hand = hand_mask * hand_image + (1 - hand_mask)*background_image[y:y + hand_image.shape[0], x:x + hand_image.shape[1]]
    output = background_image.copy()
    output[y:y + hand_image.shape[0], x:x + hand_image.shape[1]] = blended_hand
    return output, mask


这个函数虽然很长,但执行的操作很简单:

  • 加载一张随机的手部图像和掩码
  • 加载一张随机背景图像
  • 调整背景图像的大小
  • 在背景图像中选择一个随机位置来放置手部
  • 计算新的掩码
  • 计算背景图像和手部图像的融合图像


结果,生成数百甚至数千张带有分割任务标签的图像变得相当容易。以下是生成的图像样本:


11


有了这些生成的图像和掩码,我们现在可以进入下一步:训练分割模型。


训练和评估分割模型

既然我们已经正确地生成了数据,让我们在此基础上训练一个分割模型。首先,让我们讨论一下训练流程,然后评估使用这些生成数据的好处。


训练模型

我们将使用PyTorch来训练模型,同时使用Segmentation Models Pytorch库,该库允许轻松训练多种分割模型。


以下代码片段允许模型训练:


def train_model(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # Create datasets and dataloaders
    train_dataset = CustomDataset(os.path.join(args.dataset, 'train'), transform=get_train_transform())
    valid_dataset = CustomDataset(os.path.join(args.dataset, 'valid'), transform=get_valid_transform())
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
    valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
    # Create model
    model = smp.Unet(
        encoder_name=args.encoder,
        encoder_weights="imagenet",
        in_channels=3,
        classes=1,
        activation='sigmoid',
    )
    model = model.to(device)
    # Define loss function and optimizer
    criterion = nn.BCELoss()  
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
    # Train the model
    trained_model = train_model(model, train_loader, valid_loader, criterion, optimizer, args.num_epochs, device)
    # Save the trained model
    torch.save(trained_model.state_dict(), 'trained_model.pth')


这段代码执行了模型训练的典型步骤:

  • 实例化训练集和验证集,以及数据加载器
  • 实例化模型本身
  • 定义损失函数和优化器
  • 训练模型并保存


模型本身接受几个输入参数:

  • 编码器,从已实现的模型列表中选择,比如我这里使用的MobileNetV3
  • 在ImageNet数据集上的初始化权重
  • 输入通道数,这里为3(来自RGB,因为我们使用的是彩色图像)
  • 输出通道数,这里为1(因为只有一个类别)
  • 输出激活函数:这里为sigmoid(同样因为只有一个类别)


评估模型

为了评估模型以及融合图像带来的改进,让我们进行以下比较:

  1. 在Ego Hands数据集上训练并评估模型
  2. 在Ego Hands数据集上训练并评估相同的模型,同时将我们融合生成的数据添加到训练集中


在这两种情况下,我都将在Ego Hands数据集的相同子集上评估模型。作为评估指标,我将使用交并比(IoU)(也称为Jaccard指数)。以下是结果:

  • 仅在Ego Hands数据集上,经过20个epoch后:IoU = 0.72
  • 在Ego Hands数据集 + Blender生成的图像上,经过20个epoch后:IoU = 0.76


如我们所见,得益于Blender生成的图像数据集,IoU从0.72显著提高到0.76。


测试模型

对于任何希望在自己的计算机上尝试此模型的人,我还在GitHub上添加了一个脚本,以便它可以在网络摄像头馈送上实时运行。


def run_demo(model_path: str):
    # Load the trained model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = smp.Unet(
        encoder_name="timm-mobilenetv3_large_100",
        encoder_weights=None,
        in_channels=3,
        classes=1,
        activation='sigmoid',
    )
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()
    # Open webcam
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        # Preprocess the frame
        input_tensor = preprocess_image(frame)
        input_tensor = input_tensor.to(device)
        # Run inference
        with torch.no_grad():
            mask = model(input_tensor)
            mask = mask.squeeze().cpu().numpy()
        # Resize mask to match the original frame size
        mask_resized = cv2.resize(mask, (frame.shape[1], frame.shape[0]))
        # Create a colored overlay
        overlay = np.zeros_like(frame)
        overlay[mask_resized > 0.5] = [0, 255, 0]  # Green color for segmented areas
        # Blend the original frame with the overlay
        alpha = 0.2
        cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)
        # Display the result
        cv2.imshow('Hand segmentation', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


由于我训练的是一个相对较小的模型(MobileNetV3 Large 100),大多数现代笔记本电脑应该都能有效地运行这段代码。


结论

让我们以几个关键要点来总结这篇文章:

  • Blender是一个强大的工具,可以让你在不同条件下(如光线、相机位置、变形等)生成逼真的图像。
  • 利用Blender生成合成数据最初可能需要一些时间,但可以通过Python API实现完全自动化。
  • 使用生成的数据提高了语义分割任务中模型的性能:将交并比(IoU)从0.72提高到了0.76。
  • 为了获得更多样化的数据集,可以使用更多的Blender手部模型:更多的手部形状和纹理可以帮助分割模型更好地泛化。
文章来源:https://medium.com/towards-data-science/scaling-segmentation-with-blender-how-to-automate-dataset-creation-73aa38967599
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消