如果你曾经为新项目训练过细分模型,那么你可能知道,问题不在于模型,而在于数据。
收集图像通常很简单;主要挑战通常在于标记。对图像进行注释以进行分割非常耗时。即使使用Meta 的 SAM2等高级工具,创建完全注释、稳健且多样化的数据集仍需要相当长的时间。
在本文中,我们将探讨另一种通常较少被探索的选项:使用 3D 工具,例如Blender。事实上,3D 引擎越来越强大和逼真。此外,它们还具有一个引人注目的优势:能够在创建数据集时自动生成标签,无需手动注释。
在本文中,我们将概述创建手部分割模型的完整解决方案,分为以下几个关键部分:
生成手部图像
为了生成手部图像,让我们使用Blender。我虽然不是这类工具的专家,但它为我们提供了一些非常有用的功能:
如我们所见,这些功能相当有用,并且将使我们能够相当容易地制作合成数据。为了确保足够的多样性,我们将探索如何在生成的手部图像中自动随机化以下参数:
注:此处提出的方法并非完全无潜在肤色偏见,也不声称无偏见。基于此方法的任何产品都必须针对任何伦理偏见进行仔细评估。
在深入这些步骤之前,我们需要一个手部的3D模型。像Turbosquid这样的网站上有许多模型,但我使用了一个可以在此处免费获取的手部模型。如果你用Blender打开这个文件,你会得到类似下面的截图。
如图所示,该模型不仅包括手部的形状和纹理,还包括骨骼结构,从而能够模拟手部的运动。让我们以此为基础,通过调整手指位置、肤色和相机位置,来获取一组多样化的手部图像。
修改手指位置
第一步是确保获得一组多样化且逼真的手指位置。为了不涉及过多细节(因为这更多与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
如我们所见,我们所做的只是随机更新控制器的位置,从而在约束条件下移动手指。在正确的约束条件下,我们得到的手指位置如下所示:
这生成了逼真且多样的手指位置,最终使得能够生成一组多样化的手部图像。现在,让我们来调整肤色。
修改肤色
在创建包含人物的新图像数据集时,最具挑战性的方面之一可能是实现足够广泛的肤色表示。确保模型在所有肤色上都能无偏见地高效工作是一个至关重要的优先事项。虽然我并不声称能够消除任何偏见,但我在此提出的方法允许通过自动改变肤色来提供一个权宜之计。
注:此方法并不声称能够使模型免受任何伦理偏见的影响。任何用于生产的模型都必须经过公平性评估的仔细测试。可以参考谷歌为其人脸检测模型所做的工作作为示例。
我在这里所做的是对图像进行纯粹的图像处理计算。想法很简单:给定一个目标颜色和渲染出的手的平均颜色,我只需计算这两种颜色之间的差异。然后,我将这个差异应用到渲染出的手上,以得到新的肤色:
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
结果,它生成了以下手部图像:
虽然结果并不完美,但通过简单的图像处理,它们生成了具有多样肤色的相当逼真的图像。为了获得足够多样化的图像集,只剩最后一步:渲染视角。
修改相机位置
最后,让我们调整相机位置,以从多个视角捕捉手部。为此,相机被放置在以手部为中心的一个球体上的随机点上。这可以通过简单地调整球面坐标的两个角度来轻松实现。在以下代码中,我生成了球体上的一个随机位置:
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
结果,我们现在得到了以下图像样本:
现在,我们拥有了手指位置、肤色各异,且从不同视角拍摄的手部图像。在训练分割模型之前,下一步实际上是生成各种背景和情境中的手部图像。
生成训练数据
为了生成足够多样且逼真的图像,我们将把我们生成的手部与一组选定的背景图像进行融合。
我从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
这个函数虽然很长,但执行的操作很简单:
结果,生成数百甚至数千张带有分割任务标签的图像变得相当容易。以下是生成的图像样本:
有了这些生成的图像和掩码,我们现在可以进入下一步:训练分割模型。
训练和评估分割模型
既然我们已经正确地生成了数据,让我们在此基础上训练一个分割模型。首先,让我们讨论一下训练流程,然后评估使用这些生成数据的好处。
训练模型
我们将使用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')
这段代码执行了模型训练的典型步骤:
模型本身接受几个输入参数:
评估模型
为了评估模型以及融合图像带来的改进,让我们进行以下比较:
在这两种情况下,我都将在Ego Hands数据集的相同子集上评估模型。作为评估指标,我将使用交并比(IoU)(也称为Jaccard指数)。以下是结果:
如我们所见,得益于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),大多数现代笔记本电脑应该都能有效地运行这段代码。
结论
让我们以几个关键要点来总结这篇文章: