在计算机视觉领域,对象检测和实例分割是两项关键任务,它们使得机器能够理解并交互视觉数据。准确识别和隔离图像中的对象对于从自动驾驶汽车到医学成像等众多实际应用至关重要。
在这篇博客文章中,将探讨如何将边界框转换为分割掩码,并使用Jupyter笔记本借助Roboflow和Ultralytics YOLOv8移除图像背景。
想象一下,有一组包含感兴趣对象的图像数据集,每张图像都标注了边界框。虽然边界框提供了对象的位置信息,但它们缺乏进行更高级计算机视觉任务(如实例分割或背景移除)所需的细节。将边界框转换为分割掩码使能够提取准确的对象边界并将它们与背景分离,为分析和操作开辟了新的机会。
为了应对将边界框转换为分割掩码的挑战,将在Jupyter笔记本环境中使用Roboflow和Ultralytics库。Roboflow简化了数据准备和注释,而Ultralytics提供了最先进的对象检测模型和实用工具。
# 设置GPU加速的设备
device = "cuda"
# 检查Ultralytics版本和设置完成情况
ultralytics.checks()
# 设置首次运行标志为False
first_run = False
if first_run:
!{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'
!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
接下来,使用Roboflow API加载数据集,以访问和管理数据集。以下代码片段展示了如何从特定项目加载数据集版本:
from roboflow import Roboflow
rf = Roboflow(api_key="YOUR_API_KEY")
project = rf.workspace("vkr-v2").project("vkrrr")
dataset = project.version(5).download("yolov8")
要使用YOLOv8进行对象检测,运行以下代码:
from ultralytics import YOLO
# 加载YOLOv8模型
model = YOLO('yolov8n.pt')
# 对图像进行对象检测
results = model.predict(source='PATH_TO_IMAGE', conf=0.25)
一旦有了YOLOv8的结果,可以提取检测到的对象的边界框坐标:
for result in results:
boxes = result.boxes
bbox = boxes.xyxy.tolist()[0]
print(bbox)
[746.568603515625, 40.80133056640625, 1142.08056640625, 712.3660888671875]
让加载SAM模型并为其设置推理:
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
image = cv2.cvtColor(cv2.imread('PATH_TO_IMAGE'), cv2.COLOR_BGR2RGB)
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
predictor.set_image(image)
接下来,将使用SAM模型将边界框坐标转换为分割掩码:
input_box = np.array(bbox)
masks, _, _ = predictor.predict(
point_coords=None,
point_labels=None,
box=bbox[None, :],
multimask_output=False,
)
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[0], plt.gca())
show_box(input_box, plt.gca())
plt.axis('off')
plt.show()
最后,可以使用分割掩码从图像中移除背景。以下代码片段展示了该过程:
segmentation_mask = masks[0]
# 将分割掩码转换为二进制掩码
binary_mask = np.where(segmentation_mask > 0.5, 1, 0)
white_background = np.ones_like(image) * 255
# 应用二进制掩码
new_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]
plt.imshow(new_image.astype(np.uint8))
plt.axis('off')
plt.show()
结果图像随后显示,没有背景。