我自己的原文哦~ https://blog.51cto.com/whaosoft/13602490
一、其他YOLO
1.1、YOLO 和 SAM 检测和分割物体
目前有很多分割模型,一般来说,它们的工作原理相似。分割模型的典型流程包括收集数据,然后配置模型,最后训练分割模型。最后,您可以分割对象,但如果我告诉您不必做所有这些事情呢?使用 SAM,您只需向模型提供某种坐标,它就会为您完成所有分割。
在本文中,我将向您展示如何使用 YOLO 模型检测物体,以及如何通过使用检测信息使用 SAM 分割这些物体。

在继续之前,您可能会想,我们为什么要使用 YOLO 模型?答案很简单:SAM 不提供标签输出。它只根据位置输入进行分割。
通过使用对象检测模型,可以找到对象的位置和标签。然后,根据位置数据进行分割,最后,您将得到一个分割良好且标记良好的对象。

与其他分割模型不同,您需要向 SAM 提供一些位置信息,您可以提供 3 种不同类型的信息作为输入:
- 单点:一个点对(x,y)作为输入
- 边界框:边界框坐标 (x1, y1, x2, y2) 作为输入
- 多点:正点和负点作为输入
因为我们使用的是 YOLO 模型,所以边界框方法似乎最适合我们的目的。YOLO 模型将边界框坐标作为输出,我们可以直接使用此数据进行分割。看图;这正是我们要构建的。
Yolov8 物体检测
我已经有一篇关于如何训练自定义 YOLO 模型的文章,你可以阅读它(链接),但现在我将使用预训练yolov8n.pt模型。如果你决定使用预训练模型,你可以直接使用相同的代码,它会yolov8n直接下载模型。
pip install ultralytics
from ultralytics import YOLO
# Load a model
model = YOLO("yolov8n.pt") # pretrained YOLO11n model
# Run batched inference on a list of images
results = model([r"ball.jpg"]) # return a list of Results objects
# Process results list
for result in results:
boxes = result.boxes # Boxes object for bounding box outputs
masks = result.masks # Masks object for segmentation masks outputs
keypoints = result.keypoints # Keypoints object for pose outputs
probs = result.probs # Probs object for classification outputs
obb = result.obb # Oriented boxes object for OBB outputs
result.show() # display to screen
result.save(filename="result.jpg") # save to disk

使用 SAM 进行检测 + 分割
首先,你需要下载 SAM 模型。你可以从 Hugging Face (链接) 下载。
https://medium.com/@siromermer/detecting-chess-pieces-pipeline-for-training-custom-yolov8-models-765027097f15
1. 安装必要的库
%pip install git+https://github.com/facebookresearch/segment-anything.git
%pip install opencv-python mediapipe ultralytics numpy torch matplotlib
2. 导入库
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from ultralytics import YOLO
from segment_anything import sam_model_registry, SamPredictor
3. 加载模型
# Load YOLO model
model = YOLO("yolov8n.pt")
# Load SAM model
sam_checkpoint = "sam_vit_b_01ec64.pth" # Replace with your model path
model_type = "vit_b"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
4. 检测+分割
# Load image
image_path = r"ball.jpg"
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB for SAM
# Run YOLO inference
results = model(image, cnotallow=0.3)
# Loop through detections
for result in results:
# Get bounding boxes
for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
x1, y1, x2, y2 = map(int, box) # Convert to integers
# get ID
class_id = int(cls) # Class ID
# Get class label
class_label = model.names[class_id]
# Prepare SAM
predictor.set_image(image_rgb)
# Define a box prompt for SAM
input_box = np.array([[x1, y1, x2, y2]])
# Get SAM mask
masks, _, _ = predictor.predict(box=input_box, multimask_output=False)
# Create a copy of the original image for overlaying the mask
highlighted_image = image_rgb.copy()
# Apply the mask with semi-transparent blue color to the image
mask = masks[0]
# Create a blank image
blue_overlay = np.zeros_like(image_rgb, dtype=np.uint8)
# Blue color for the segmented area (RGB)
blue_overlay[mask == 1] = [0, 0, 255]
# Blend the blue overlay with the original image using transparency
alpha = 0.7 # Transparency level for the overlay
highlighted_image = cv2.addWeighted(highlighted_image, 1 - alpha, blue_overlay, alpha, 0)
# Add label (class name) on top of the bounding box
font = cv2.FONT_HERSHEY_SIMPLEX
label = f"{class_label}" # Label is the class name
cv2.putText(highlighted_image, label, (x1, y1 - 10), font, 2, (255, 255, 0), 2, cv2.LINE_AA)
# Optional: Save the image with the bounding box and highlighted segment
output_filename = f"highlighted_output.png"
cv2.imwrite(output_filename, cv2.cvtColor(highlighted_image, cv2.COLOR_RGB2BGR))
