BimAnt

YOLOv8+SAM标注分类

admin -

在本指南中，我们将学习如何将计算机视觉基础模型链接在一起，以非常快速地自动用 Labelbox 中的类填充预标签。 我们将完成一个简单的语义分割任务：在图像中特定类的所有对象周围绘制掩模。

利用 Labelbox 中的链式基础模型将大大减少你绘制分割掩模所需的时间； 通过使用分类增强最近发布的 SAM 模型，将自动执行分配类别的任务。 你可以集中精力检查、验证和纠正 AI 模型绘制的标签，而不是执行繁琐的标记工作。

以下是我们将在下面逐步介绍的流程的高层次摘要：

  • 在图像上运行对象检测器以生成具有指定类别分类的边界框
  • 将边界框作为 Meta 的 Segment Anything 模型的输入，该模型将为每个边界框生成分割掩码
  • 将掩模预测作为预标签上传到 Labelbox
  • 打开图像编辑器并像平常一样检查或修改预标签

可以使用我们的 Colab 笔记本在你的图像上开箱即用地运行上述所有内容。 只需加载图像并在短短几分钟内自动获取带有类的分段蒙版。

在本指南中，我们将使用以下许多彩色椅子的图像：

1、在图像上运行 YOLOv8

YOLOv8 是 YOLO（You Only Look Once）模型系列的最新版本，它是一个对象检测器，可以围绕常见对象生成边界框和类。 YOLOv8 以其速度和准确性而闻名，拥有一些令人印象深刻的功能 - 使其成为广泛应用的宝贵工具。 在这里，我们使用 YOLOv8 自动检测并定位图像中的所有椅子。

# load the YOLOv8 model
model = YOLO(f'{HOME}/yolov8n.pt')

# run the model on the image
results = model.predict(source='chairs.jpg', conf=0.25)
predicted_boxes = results[0].boxes.xyxy


# read in the image for visualization
image_bgr = cv2.imread(IMAGE_PATH, cv2.IMREAD_COLOR)

# use cv2 to visualize the bounding boxes on the image
for box in predicted_boxes:
 cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
cv2.imshow("YOLOv8 predictions", image_bgr)

2、将边界框作为SAM 模型的输入

SAM（Segment Anything Model）——Meta AI 最近发布的一种先进的计算机视觉模型，旨在将图像和视频准确地分割成不同的对象。 使用先进的深度学习技术，SAM 能够识别和分割图像中的对象，使其成为适用于各种应用的强大工具。 SAM 模型能够根据提示（包括边界框提示）生成分段掩码，我们将在下面的代码中使用它。

# load the SAM model
sam = sam_model_registry["vit_h"](checkpoint="/sam_vit_h_4b8939.pth
").to(device=torch.device('cuda:0'))

mask_predictor = SamPredictor(sam)

# transform the YOLOv8 predicted boxes to match input format expected by SAM model
transformed_boxes = mask_predictor.transform.apply_boxes_torch(predicted_boxes, image_bgr.shape[:2])


# run SAM model on all the boxes
mask_predictor.set_image(image_bgr)
masks, scores, logits = mask_predictor.predict_torch(
   boxes = transformed_boxes,
   multimask_output=False,
   point_coords=None,
   point_labels=None
)

# combine all masks into one for easy visualization
final_mask = None
for i in range(len(masks) - 1):
  if final_mask is None:
    final_mask = np.bitwise_or(masks[i][0], masks[i+1][0])
  else:
    final_mask = np.bitwise_or(final_mask, masks[i+1][0])

# visualize the predicted masks
plt.figure(figsize=(10, 10))
plt.imshow(image_rgb)
plt.imshow(final_mask, cmap='gray', alpha=0.7)
plt.show()

3、将预测的掩模作为预标签上传到 Labelbox 上

预测的掩模可以通过我们的 SDK 轻松无缝地集成到 Labelbox 中。 上传只需几行代码，运行时间不到一分钟。

class_names = []
for c in results[0].boxes.cls:
 class_names.append(model.names[int(c)])

annotations = []
for mask in masks:
  # convert a 2D array to 3D array
  mask_data = lb_types.MaskData.from_2D_arr(np.asarray(mask[0], dtype="uint8"))
  mask_annotation = lb_types.ObjectAnnotation(
    name = class_names[idx], # assign class from Step 1
    value=lb_types.Mask(mask=mask_data, color=color),
  )
  annotations.append(mask_annotation)

labels = [
lb_types.Label(data=lb_types.ImageData(global_key="image_name"),annotations=annotations)
]
upload_job = lb.MALPredictionImport.create_from_objects(
   client=client,
   project_id=project.uid,
   name="mal_job" + str(uuid.uuid4()),
   predictions=labels
)
upload_job.wait_until_done()

print(f"Errors: {upload_job.errors}", )
print(f"Status of uploads: {upload_job.statuses}")

4、结束语

虽然 Meta 的 AI 的 SAM 在分割方面非常强大，但它忽略了分类的关键任务。 在本指南中，我们演示了如何使用 YOLOv8（或任何对象检测器）生成带有类的边界框，然后自动将这些类应用到 SAM 生成的掩模。 我们还展示了它如何与 Labelbox 模型辅助标签 SDK 无缝集成。

原文链接：Using Meta’s Segment Anything (SAM) model with YOLOv8 to automatically classify masks

BimAnt翻译整理，转载请标明出处