实现步骤如下：
-
环境设置
-
下载 YOLOv9 和 SAM 的预训练模型权重
-
图像推理
-
可视化和分析
-
获取检测结果
-
使用 SAM 进行分割
环境设置
需要有 Google 帐户才能访问 Google Colab，这是一项免费云服务，为深度学习任务提供必要的计算资源，包括访问高达 16 GB 的 GPU 和 TPU。
GPU状态检查
首先，我们确保 GPU 的可用性和就绪性，用于处理和运行 YOLOv9+SAM 模型。
安装 Google 云盘
接下来，如果您已经下载了数据集，我们需要导航到存储数据集的文件夹，否则我们可以直接使用 Roboflow 加载数据集。
from google.colab import drivedrive.mount('/content/drive')
or
%cd {HOME}/!pip install -q roboflow
from roboflow import Roboflowrf = 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 osHOME = 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。
!mkdir -p {HOME}/data!wget -P {HOME}/data -q /content/drive/MyDrive/data/image9.jpegSOURCE_IMAGE_PATH = f"{HOME}/image9.jpeg"
使用自定义数据运行检测
之后，执行detect.py指定参数对图像进行目标检测，设置置信度阈值并保存检测结果。这将创建一个包含 class_ids、边界框坐标和置信度分数的文本文件，我们稍后将使用它。
!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
然后，我们利用 IPython 的显示和图像功能来展示指定路径图像中检测到的对象，并进行调整以获得最佳观看效果。
安装 Ultralytics 包以访问 YOLO 对象检测模型实现和实用程序，不要忘记导入 YOLO 类以执行对象检测任务。
!pip install ultralytics from ultralytics import YOLO
安装
Segment-Anything模型
现在让我们安装 Segment-Anything 库并下载 SAM 模型的权重文件，为高质量图像分割任务做好准备。
!pip install 'git+https://github.com/facebookresearch/segment-anything.git'
!wget https:
提取检测结果和置信度分数
我们需要将 YOLOv9 检测结果保存在上面的文本文件中来提取类 ID、置信度分数和边界框坐标。这里的坐标已经标准化，所以我们先将它们转换为图像比例，然后打印它们进行验证。
import cv2
image_path = '/content/drive/MyDrive/data/image9.jpeg'
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]]
cx *= image_width cy *= image_height w *= image_width h *= image_height
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)
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, SamPredictorsam_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 cv2image = cv2.cvtColor(cv2.imread('/content/drive/MyDrive/data/image9.jpeg'), cv2.COLOR_BGR2RGB)predictor.set_image(image)
结果可视化
为了可视化检测和分割结果，我们必须使用 SAM 将边界框转换为分割掩模。我们随机为类 ID 分配唯一的颜色，然后定义用于显示掩码、置信度分数和边界框的辅助函数。coco.yaml 文件用于将 class_ids 映射到类名。
import matplotlib.patches as patchesfrom matplotlib import pyplot as pltimport numpy as npimport 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]
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
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)
