【深度学习】使用 YOLOv9+SAM 进行动态物体检测和分割

YOLOv9+SAM：自定义对象检测新突破，RF100 Construction-Safety-2数据集显威力！

在本文中，我们大胆尝试，将领先的YOLOv9算法与SAM（分割注意力模块）技术完美融合，并在RF100 Construction-Safety-2数据集上进行了深入的自定义对象检测模型测试。这一前沿技术的集成，不仅显著提升了在不同图像中检测和分割对象的精准度和细致度，更为我们的应用场景打开了更为广阔的大门。

想象一下，从自动驾驶系统的安全性能增强，到医学成像中诊断过程的精确改进，这一组合技术的潜力无处不在。它仿佛一双智能的眼睛，能够精准捕捉到每一个关键细节，为我们的生活和工作带来前所未有的便利和安全性。

而这一切的背后，正是YOLOv9的高效检测能力与SAM以零样本方式分割对象能力的完美结合。这种强大的组合不仅最大限度地减少了大量再训练或数据注释的繁琐工作，更使得它成为一种既实用又易于扩展的解决方案。

在RF100 Construction-Safety-2数据集上的表现更是令人瞩目。它像一把锐利的剑，精准地切割出每一个目标对象，无论是大小、形状还是颜色，都能得到完美的呈现。这不仅证明了YOLOv9+SAM技术的强大实力，更为我们未来的研究和应用指明了方向。

综上所述，YOLOv9+SAM无疑是一种值得我们深入研究和广泛应用的前沿技术。相信在不久的将来，它将在更多领域展现出强大的应用潜力和价值！

YOLOv9简介

来自官方 YOLOv9 存储库的 YOLOv9 性能。

YOLOv9 模型图来自官方 YOLOv9 存储库。

YOLOv9：实时物体检测领域的革命性突破

在实时物体检测领域，YOLOv9无疑是一次重大进步。它凭借出色的效率、准确性和适应性，成为了业界的翘楚。而这一切，都得益于它结合了可编程梯度信息（PGI）和广义高效层聚合网络（GELAN）两大创新技术。在MS COCO数据集上的卓越表现，更是成为了其强大实力的有力证明。

YOLOv9不仅继承了开源社区的协作精神，更在Ultralytics YOLOv5的基础上进行了深入优化。它巧妙地利用信息瓶颈原理和可逆函数，成功解决了深度学习中的信息丢失问题，确保在各层之间能够保留重要的数据。这一创新策略不仅提高了模型的结构效率，更确保了精确的检测能力，即便是在轻量级模型中，也能保持对细节的敏锐捕捉。

更值得一提的是，YOLOv9的架构经过精心设计，有效减少了不必要的参数和计算需求。这使得它能够在各种模型大小（从紧凑的YOLOv9-S到更广泛的YOLOv9-E）中都能实现最佳性能。无论是追求速度的快速检测，还是追求高精度的细致识别，YOLOv9都能展现出速度与检测精度之间的和谐平衡。

作为计算机视觉领域的里程碑，YOLOv9不仅树立了新的标杆，更拓宽了人工智能在物体检测与分割领域的应用视野。它的出现，不仅彰显了该领域战略创新与协同努力的影响，更为我们带来了更多关于未来智能生活的无限遐想。

我们相信，随着技术的不断进步和应用的不断拓展，YOLOv9将会在未来继续发挥其强大的潜力，为我们带来更多的惊喜和突破。让我们拭目以待，共同见证这一领域的蓬勃发展！

分段任意模型 (SAM)简介

来自官方 META 存储库的SAM 模型图

任意分割模型 (SAM) 通过简化图像分割推动计算机视觉的发展，这对于从科学研究到创造性努力的一系列用途都至关重要。

SAM 利用迄今为止最大的 Segment Anything 1-Billion (SA-1B) 掩模数据集，通过减少对专业知识、强大计算能力和大量数据集注释的依赖，实现分割的民主化。

在 Apache 2.0 许可下，SAM 引入了基础模型框架，允许通过简单的提示轻松地适应任务，反映了自然语言处理方面的进步。

通过对超过 10 亿个不同掩模进行训练，SAM 理解了物体的广义概念，促进了不熟悉领域的零样本迁移，并增强了其在 AR/VR、创意艺术和科学研究等各个领域的实用性。

该模型的提示驱动灵活性和广泛的任务适应性标志着分割技术的重大飞跃，将 SAM 定位为研究界和应用程序开发人员的多功能且可访问的工具。

关于数据集

该项目利用Roboflow 的RF100施工数据集，特别是Construction-Safety-2子集，来展示集成模型的功能。

RF100 是英特尔赞助的一项计划，旨在为对象检测模型建立开源基准。它侧重于 Roboflow Universe 上可用数据集的通用性，促进可访问性并加快 AI 和深度学习的开发过程。

RF100施工数据集与Roboflow中的其他项目一样，是开源的，可以自由使用。

入门

为了在你的项目中充分利用 YOLOv9 和 SAM 的综合优势，请遵循以下步骤

• 环境设置

• 下载 YOLOv9 和 SAM 的预训练模型权重

• 图像推理

• 可视化与分析

• 获取检测结果

• 使用 SAM 进行分割

环境设置

需要有 Google 帐户才能访问 Google Colab，这是一项免费的云服务，为深度学习任务提供必要的计算资源，包括访问高达 16 GB 的 GPU 和 TPU。

GPU 状态检查

首先，我们确保 GPU 可用于处理和运行 YOLOv9+SAM 模型。

安装 Google Drive

接下来，如果您已经下载了数据集，我们需要导航到存储数据集的文件夹，否则我们可以直接使用 Roboflow 来加载数据集。

from google.colab import drive
drive.mount('/content/drive')


or


%cd {HOME}/
!pip install -q roboflow


from roboflow import Roboflow
rf = Roboflow(api_key="YOUR API KEY")
project = rf.workspace("roboflow-100").project("construction-safety-gsnvb")
dataset = project.version(2).download("yolov7")

设置 YOLOv9

数据集准备好后，克隆 YOLOv9 存储库，然后切换到 YOLOv9 目录并安装所需的依赖项，为对象检测和分割任务做准备。

!git clone https://github.com/SkalskiP/yolov9.git
%cd yolov9
!pip3 install -r requirements.txt -q

显示当前目录

将当前工作目录的路径存储在HOME变量中以供参考。

import os
HOME = os.getcwd()
print(HOME)

下载模型权重

让我们创建一个模型权重目录，并从 GitHub 上的发布页面下载特定的 YOLOv9 和 GELAN 模型权重，这对于使用预训练参数初始化模型至关重要

!mkdir -p {HOME}/weights 
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /yolov9-c.pt 
!wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /yolov9-e.pt !wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /gelan-c.pt !wget -P {HOME}/weights -q 
https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /gelan-e.pt

下载图像进行推理

为了使用 YOLOv9 权重进行推理，我们必须设置一个数据目录并下载一个示例图像进行处理，并在变量中设置该图像的路径SOURCE_IMAGE_PATH。

!python detect.py --weights {HOME}/weights/gelan-c.pt --conf 0.1 --source /content/drive/MyDrive/data/image9.jpeg --device 0 --save-txt --save-conf

使用自定义数据运行检测

此后，detect.py使用指定的参数执行对图像进行对象检测，设置置信度阈值并保存检测结果。这将创建一个文本文件，其中包含 class_id、边界框坐标和置信度分数，我们稍后会用到它。

!mkdir -p {HOME}/data
!wget -P {HOME}/data -q /content/drive/MyDrive/data/image9.jpeg
SOURCE_IMAGE_PATH = f"{HOME}/image9.jpeg"

显示检测结果

然后我们利用 IPython 的显示和图像函数来展示指定路径图像中检测到的对象，并进行调整以获得最佳观看效果。

import matplotlib.patches as patches
from matplotlib import pyplot as plt
import numpy as np
import yaml


with open('/content/yolov9/data/coco.yaml', 'r') as file:
    coco_data = yaml.safe_load(file)
    class_names = coco_data['names']


for class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):
    class_name = class_names[class_id]
# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')


color_map = {}
for class_id in class_ids:
    color_map[class_id] = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)


def show_mask(mask, ax, color):
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * np.array(color).reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_box(box, label, conf_score, color, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    rect = plt.Rectangle((x0, y0), w, h, edgecolor=color, facecolor='none', lw=2)
    ax.add_patch(rect)


    label_offset = 10


# Construct the label with the class name and confidence score
    label_text = f'{label} {conf_score:.2f}'


    ax.text(x0, y0 - label_offset, label_text, color='black', fontsize=10, va='top', ha='left',
            bbox=dict(facecolor=color, alpha=0.7, edgecolor='none', boxstyle='square,pad=0.4'))


plt.figure(figsize=(10, 10))
ax = plt.gca()
plt.imshow(image)


# Display and process each bounding box with the corresponding mask
for class_id, bbox in zip(class_ids, bboxes):
    class_name = class_names[class_id]
# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')


    color = color_map[class_id]
    input_box = np.array(bbox)


# Generate the mask for the current bounding box
    masks, _, _ = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_box,
        multimask_output=False,
    )


    show_mask(masks[0], ax, color=color)
    show_box(bbox, class_name, conf, color, ax)


# Show the final plot
plt.axis('off')
plt.show()

安装 Ultralytics

安装 Ultralytics 包以访问 YOLO 对象检测模型实现和实用程序，不要忘记导入 YOLO 类用于对象检测任务。

!pip install ultralytics
from ultralytics import YOLO

安装 Segment-Anything 模型

现在让我们安装 Segment-Anything 库并下载 SAM 模型的权重文件，为高质量的图像分割任务做好准备。

!pip 安装'git+https://github.com/facebookresearch/segment-anything.git'


 !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

提取检测结果和置信度分数

我们需要上面保存在文本文件中的 YOLOv9 检测结果来提取类别 ID、置信度分数和边界框坐标。这里的坐标是标准化的，因此让我们先将它们转换为图像比例，然后打印它们以进行验证。

import cv2


# Specify the path to your image
image_path = '/content/drive/MyDrive/data/image9.jpeg'


# Read the image to get its dimensions
image = cv2.imread(image_path)
image_height, image_width, _ = image.shape


detections_path = '/content/yolov9/runs/detect/exp/labels/image9.txt'


bboxes = []
class_ids = []
conf_scores = []


with open(detections_path, 'r') as file:
for line in file:
components = line.split()
class_id = int(components[0])
confidence = float(components[5])
cx, cy, w, h = [float(x) for x in components[1:5]]


        # Convert from normalized [0, 1] to image scale
cx *= image_width
cy *= image_height
w *= image_width
h *= image_height


        # Convert the center x, y, width, and height to xmin, ymin, xmax, ymax
xmin = cx - w / 2
ymin = cy - h / 2
xmax = cx + w / 2
ymax = cy + h / 2


class_ids.append(class_id)
bboxes.append((xmin, ymin, xmax, ymax))
conf_scores.append(confidence)


# Display the results
for class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):
print(f'Class ID: {class_id}, Confidence: {conf:.2f}, BBox coordinates: {bbox}')

初始化 SAM 进行图像分割

一旦使用指定的预训练权重初始化 SAM，我们就会继续从 SAM 模型注册表中选择模型类型来生成分割蒙版。

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
sam_checkpoint = "/content/yolov9/sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
predictor = SamPredictor(sam)

加载图像进行分割

使用 OpenCV 库，我们加载图像以便用 SAM 进行处理，为分割做好准备。

import cv2
image = cv2.cvtColor(cv2.imread('/content/drive/MyDrive/data/image9.jpeg'), cv2.COLOR_BGR2RGB)
predictor.set_image(image)

结果可视化

为了可视化检测和分割结果，我们必须使用 SAM 将边界框转换为分割掩码。我们随机为类 ID 分配唯一颜色，然后定义用于显示掩码、置信度分数和边界框的辅助函数。coco.yaml 文件用于将 class_id 映射到类名。

import matplotlib.patches as patchesfrom matplotlib import pyplot as pltimport numpy as npimport yamlwith open('/content/yolov9/data/coco.yaml', 'r') as file:    coco_data = yaml.safe_load(file)    class_names = coco_data['names']for class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):    class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')color_map = {}for class_id in class_ids:    color_map[class_id] = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)def show_mask(mask, ax, color):    h, w = mask.shape[-2:]    mask_image = mask.reshape(h, w, 1) * np.array(color).reshape(1, 1, -1)    ax.imshow(mask_image)def show_box(box, label, conf_score, color, ax):    x0, y0 = box[0], box[1]    w, h = box[2] - box[0], box[3] - box[1]    rect = plt.Rectangle((x0, y0), w, h, edgecolor=color, facecolor='none', lw=2)    ax.add_patch(rect)    label_offset = 10# Construct the label with the class name and confidence score    label_text = f'{label} {conf_score:.2f}'    ax.text(x0, y0 - label_offset, label_text, color='black', fontsize=10, va='top', ha='left',            bbox=dict(facecolor=color, alpha=0.7, edgecolor='none', boxstyle='square,pad=0.4'))plt.figure(figsize=(10, 10))ax = plt.gca()plt.imshow(image)# Display and process each bounding box with the corresponding maskfor class_id, bbox in zip(class_ids, bboxes):    class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')    color = color_map[class_id]    input_box = np.array(bbox)# Generate the mask for the current bounding box    masks, _, _ = predictor.predict(        point_coords=None,        point_labels=None,        box=input_box,        multimask_output=False,    )    show_mask(masks[0], ax, color=color)    show_box(bbox, class_name, conf, color, ax)# Show the final plotplt.axis('off')plt.show()

提取蒙版

最后，我们生成一个合成图像，在白色背景上突出显示检测到的物体，从分割蒙版创建聚合蒙版，并将其应用到原始图像与白色背景混合，以增强可视化。

aggregate_mask = np.zeros(image.shape[:2], dtype=np.uint8)


# Generate and accumulate masks for all bounding boxes
for bbox in bboxes:
input_box = np.array(bbox).reshape(1, 4)
masks, _, _ = predictor.predict(
point_coords=None,
point_labels=None,
box=input_box,
multimask_output=False,
)
aggregate_mask = np.where(masks[0] > 0.5, 1, aggregate_mask)


# Convert the aggregate segmentation mask to a binary mask
binary_mask = np.where(aggregate_mask == 1, 1, 0)


# Create a white background with the same size as the image
white_background = np.ones_like(image) * 255


# Apply the binary mask to the original image
# Where the binary mask is 0 (background), use white_background; otherwise, use the original image
new_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]


# Display the new image with the detections and white background
plt.figure(figsize=(10, 10))
plt.imshow(new_image.astype(np.uint8))
plt.axis('off')
plt.show()

结论

我希望本指南清晰易懂，并且您已成功训练模型。通过利用 Roboflow 的 RF100 构造数据集，我们了解了如何使用模型通过自定义数据集识别和分割对象，展示了将 YOLOv9 和 SAM 结合用于高级计算机视觉任务的实际应用。

往期精彩回顾




适合初学者入门人工智能的路线及资料下载
(图文+视频)机器学习入门系列下载
机器学习及深度学习笔记等资料打印
《统计学习方法》的代码复现专辑

• 交流群

欢迎加入机器学习爱好者微信群一起和同行交流，目前有机器学习交流群、博士群、博士申报交流、CV、NLP等微信群，请扫描下面的微信号加群，备注：”昵称-学校/公司-研究方向“，例如：”张小明-浙大-CV“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~（也可以加入机器学习交流qq群772479961）