概述
在當前的計算機視覺領域，目標分割技術正變得越來越重要。市面上有許多分割模型，它們的工作原理大致相似，通常包括收集數據、配置模型以及訓練分割模型等步驟。最終目標是實現精確的目標分割。而隨著 SAM（Segment Anything Model）的出現，這一過程變得更加高效。SAM 的獨特之處在于，它只需要用戶向模型提供某種坐標信息，就能自動完成所有分割工作，極大地簡化了操作流程。
在深入探討之前，可能會提出這樣一個問題：為何選擇 YOLO 模型作為我們的工具？ 答案十分明確：SAM（Segment Anything Model）本身并不具備直接輸出目標類別標簽的功能。 它僅能依據用戶提供的位置信息來執行分割任務。而 YOLO 模型的引入，恰到好處地填補了這一空白。借助目標檢測模型，我們能夠精準地定位目標的位置，并獲取其對應的類別標簽。隨后，利用這些位置數據引導 SAM 進行分割操作，從而最終實現目標的清晰分割，并為分割結果賦予明確的類別標注。
與傳統分割模型不同，SAM（Segment Anything Model）在執行分割任務時，需要用戶主動提供目標的位置信息。這些位置信息有三種主要類型：
- 單點輸入：僅提供一個坐標點（x，y），用于指示目標的大致位置。
- 邊界框輸入：提供一個邊界框的坐標（x1，y1，x2，y2），明確指定目標的區域范圍。
- 多點輸入：同時輸入多個正點和負點，以更精細地引導模型進行分割。
鑒于我們采用的是 YOLO 模型，邊界框方法無疑是最佳選擇。YOLO 模型能夠直接輸出目標的邊界框坐標，這與 SAM 的輸入需求完美契合。因此，我們可以無縫地將 YOLO 模型的輸出作為 SAM 的輸入，從而實現高效的目標檢測與分割。在接下來的內容中，我們將詳細闡述如何將 YOLO 與 SAM 結合，以實現這一目標。
YOLOv8 目標檢測
在正式展開本文內容之前，若對如何從零開始訓練自定義的 YOLO 模型抱有興趣，不妨參考我過往撰寫的一篇文章《YOLOV8目標識別——詳細記錄從環境配置、自定義數據、模型訓練到模型推理部署》。不過，在本文中，為了便于快速上手和聚焦于核心問題，我將直接調用預訓練好的 yolov8n.pt 模型。若您選擇采用這一預訓練模型，僅需運行以下代碼，系統便會自動為您完成模型的下載工作。
conda create -n yolo_sam python==3.10
conda activate yolo_sam
pip install ultralytics
from ultralytics import YOLO# 加載模型
model = YOLO("yolov8n.pt") # 預訓練的 YOLOv8n 模型# 對一系列圖像進行批量推理
results = model([r"ball.jpg"]) # 返回一個 Results 對象列表# 處理結果列表
for result in results:boxes = result.boxes # 用于邊界框輸出的 Boxes 對象masks = result.masks # 用于分割掩碼輸出的 Masks 對象keypoints = result.keypoints # 用于姿態輸出的 Keypoints 對象probs = result.probs # 用于分類輸出的 Probs 對象obb = result.obb # 用于 OBB 輸出的 Oriented boxes 對象result.show() # 顯示到屏幕result.save(filename="result.jpg") # 保存到磁盤
檢測 + 分割與 SAM
在開始之前，你需要下載 SAM 模型。你可以從 Hugging Face 下載它（鏈接）。
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. 加載模型
# 加載 YOLO 模型
model = YOLO("yolov8n.pt") # 加載 SAM 模型
sam_checkpoint = "sam_vit_b_01ec64.pth" # 替換為你的模型路徑
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. 檢測 + 分割
# 加載圖像
image_path = r"ball.jpg"
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 轉換為 RGB 格式以供 SAM 使用# 運行 YOLO 推理
results = model(image, conf=0.3)# 遍歷檢測結果
for result in results:# 獲取邊界框for box, cls in zip(result.boxes.xyxy, result.boxes.cls): x1, y1, x2, y2 = map(int, box) # 轉換為整數# 獲取 IDclass_id = int(cls) # 類別 ID# 獲取類別標簽class_label = model.names[class_id] # 準備 SAMpredictor.set_image(image_rgb)# 定義一個邊界框提示供 SAM 使用input_box = np.array([[x1, y1, x2, y2]])# 獲取 SAM 掩碼masks, _, _ = predictor.predict(box=input_box, multimask_output=False)# 創建原始圖像的副本以疊加掩碼highlighted_image = image_rgb.copy()# 將掩碼以半透明藍色應用于圖像mask = masks[0]# 創建一個空白圖像blue_overlay = np.zeros_like(image_rgb, dtype=np.uint8) # 藍色用于分割區域（RGB）blue_overlay[mask == 1] = [0, 0, 255] # 使用透明度將藍色疊加層與原始圖像混合alpha = 0.7 # 疊加層的透明度highlighted_image = cv2.addWeighted(highlighted_image, 1 - alpha, blue_overlay, alpha, 0)# 在邊界框上方添加標簽（類別名稱）font = cv2.FONT_HERSHEY_SIMPLEXlabel = f"{class_label}" # 標簽為類別名稱cv2.putText(highlighted_image, label, (x1, y1 - 10), font, 2, (255, 255, 0), 2, cv2.LINE_AA) # 可選：保存帶有邊界框和突出顯示的分割結果的圖像output_filename = f"highlighted_output.png"cv2.imwrite(output_filename, cv2.cvtColor(highlighted_image, cv2.COLOR_RGB2BGR))
交互加分割
使用 OpenCV 鼠標事件來畫框交互，并將框的坐標傳入 SAM 進行分割，是一種非常直觀且靈活的方法。
1. 代碼實現
import cv2
import numpy as np
import torch
from segment_anything import sam_model_registry, SamPredictor# 初始化全局變量
start_point = None
end_point = None
drawing = False# 鼠標回調函數
def draw_rectangle(event, x, y, flags, param):global start_point, end_point, drawingif event == cv2.EVENT_LBUTTONDOWN:start_point = (x, y)drawing = Trueelif event == cv2.EVENT_MOUSEMOVE:if drawing:end_point = (x, y)elif event == cv2.EVENT_LBUTTONUP:drawing = Falseend_point = (x, y)cv2.rectangle(image, start_point, end_point, (0, 255, 0), 2)cv2.imshow("Image", image)# 加載 SAM 模型
sam_checkpoint = "sam_vit_b_01ec64.pth" # 替換為你的模型路徑
model_type = "vit_b"
device = "cuda" if torch.cuda.is_available() else "cpu"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)# 加載圖像
image_path = "example.jpg"
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 轉換為 RGB 格式# 創建窗口并綁定鼠標回調函數
cv2.namedWindow("Image")
cv2.setMouseCallback("Image", draw_rectangle)while True:cv2.imshow("Image", image)if cv2.waitKey(1) & 0xFF == ord('q'):break# 如果框已經畫好，進行分割if start_point and end_point:# 轉換為 SAM 需要的邊界框格式input_box = np.array([start_point[0], start_point[1], end_point[0], end_point[1]])input_box = np.array([min(input_box[0], input_box[2]), min(input_box[1], input_box[3]),max(input_box[0], input_box[2]), max(input_box[1], input_box[3])])# 生成分割掩碼predictor.set_image(image_rgb)masks, _, _ = predictor.predict(box=input_box, multimask_output=False)# 顯示分割結果highlighted_image = image_rgb.copy()mask = masks[0]blue_overlay = np.zeros_like(image_rgb, dtype=np.uint8)blue_overlay[mask == 1] = [0, 0, 255] # 藍色掩碼highlighted_image = cv2.addWeighted(highlighted_image, 1, blue_overlay, 0.7, 0)cv2.imshow("Segmented Image", cv2.cvtColor(highlighted_image, cv2.COLOR_RGB2BGR))cv2.waitKey(0)cv2.destroyAllWindows()breakcv2.destroyAllWindows()
實現效果：
### 2. 代碼說明
-
全局變量：
start_point和end_point用于記錄鼠標點擊和釋放時的坐標。drawing用于標記是否正在繪制矩形。
-
鼠標回調函數：
draw_rectangle函數用于處理鼠標事件，包括按下、移動和釋放。- 當鼠標按下時，記錄起始點。
- 當鼠標移動時，更新終點。
- 當鼠標釋放時，繪制矩形并更新圖像。
-
加載 SAM 模型：
- 使用
sam_model_registry加載預訓練的 SAM 模型。 - 將模型移動到 GPU（如果可用）。
- 使用
-
加載圖像：
- 使用 OpenCV 讀取圖像，并將其轉換為 RGB 格式以供 SAM 使用。
-
創建窗口并綁定鼠標回調：
- 使用
cv2.namedWindow創建窗口。 - 使用
cv2.setMouseCallback綁定鼠標回調函數。
- 使用
-
主循環：
- 顯示圖像并等待用戶操作。
- 如果用戶繪制了矩形，將矩形的坐標轉換為 SAM 需要的格式，并調用 SAM 進行分割。
- 顯示分割結果。
3. 運行效果
-
繪制矩形：
- 運行程序后，使用鼠標在圖像上繪制一個矩形框，表示需要分割的目標區域。
-
顯示分割結果：
- 按下
q鍵后，程序會根據繪制的矩形框調用 SAM 進行分割，并顯示分割結果。
- 按下
4. 注意事項
- 圖像路徑：確保
image_path指向正確的圖像文件。 - 模型路徑：確保
sam_checkpoint指向正確的 SAM 模型文件。 - 環境依賴：確保安裝了必要的依賴庫，如
opencv-python、torch和segment_anything。
