【小貪】項目實戰——Zero-shot根據文字提示分割出圖片目標掩碼

目標描述

給定RGB視頻或圖片，目標是分割出圖像中的指定目標掩碼。我們需要復現兩個Zero-shot的開源項目，分別為IDEA研究院的GroundingDINO和Facebook的SAM。首先使用目標檢測方法GroundingDINO，輸入想檢測目標的文字提示，可以獲得目標的anchor box。將上一步獲得的box信息作為SAM的提示，分割出目標mask。具體效果如下（測試數據來自VolumeDeform數據集）：

在這里插入圖片描述

其中GroundingDINO根據white shirt的文字輸入計算的box信息為："shirt_000500": "[194.23726, 2.378189, 524.09503, 441.5135]"。項目實測下來單張圖片的預測速度GroundingDINO要慢于SAM。GroundingDINO和SAM均會給出多個預測結果，當選擇置信度最高的結果時兩個模型也會存在預測不準確的情況。

論文簡介

GroundingDINO

GroundingDINO extends a closedset detector DINO by performing vision-language modality fusion at multiple phases, including a feature enhancer, a language-guided query selection module, and a cross-modality decoder. Such a deep fusion strategy effectively improves open-set object detection.

在這里插入圖片描述

SAM

簡介：使用三個組件建立圖像分割的foundation model，解決一系列下游分割問題，可zero-shot生成
關鍵技術：
1. promptable分割任務：使用prompt engineering，prompt不確定時輸出多目標mask
2. 分割模型：image encoder + prompt encoder -> mask decoder
3. 數據驅動：SA-1B（1B masks from 11M imgs）手工標注->半自動->全自動
Limitation：存在不連貫不精細的mask結果；交互式實時mask生成但是img encoder耗時；text-to-mask任務效果不魯棒

在這里插入圖片描述

項目實戰

兩個項目的復現很簡單，按照github的readme配置相關環境并運行程序。當然也可以直接使用一站式項目Grounded Segment Anything等。當需要分割的圖片較多時，可以修改GroundingDINO的demo.sh和demo/inference_on_a_image.py文件將檢測結果保存至json文件。

demo/inference_on_a_image.py文件

# 修改plot_boxes_to_image函數輸出box信息
image_with_box, mask, box_coor = plot_boxes_to_image(image_pil, pred_dict)
# obj為目標名稱，i為當前圖片的索引
obj = 'shirt'
data = {f'{obj}_{str(i).zfill(6)}': str(list(box_coor.cpu().detach().numpy()))}
with open("box.json", "r", encoding="utf-8") as f:old_data = json.load(f)old_data.update(data)
with open("box.json", "w", encoding="utf-8") as f:json.dump(old_data, f, indent=4)# f.write(json.dumps(old_data, indent=4, ensure_ascii=False))
f.close()

然后SAM再讀取json文件獲取box信息，將SAM的輸入提示改為box。

測試代碼

import os
import numpy as np
import matplotlib.pyplot as plt
import cv2
import glob
import jsoncoords = []def show_mask(mask, ax, random_color=False):if random_color:color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)else:color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])h, w = mask.shape[-2:]mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)ax.imshow(mask_image)def show_points(coords, labels, ax, marker_size=375):pos_points = coords[labels == 1]neg_points = coords[labels == 0]ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',linewidth=1.25)ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',linewidth=1.25)def show_box(box, ax):x0, y0 = box[0], box[1]w, h = box[2] - box[0], box[3] - box[1]ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))def on_click(event):global coordsif event.button == 1:x, y = event.xdata, event.ydataprint(f"鼠標左鍵點擊：x={x:.2f}, y={y:.2f}")coords.append([x, y])# if len(coords) == 2:#     fig.canvas.mpl_disconnect(cid)elif event.button == 3:print("鼠標右鍵點擊")def get_mask(image, mask_id=1, click_coords=False, choose_mask=False, box=None):image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# plt.figure(figsize=(10, 10))# plt.imshow(image)# plt.axis('on')if click_coords:global coordsfig, ax = plt.subplots()  # 創建畫布和子圖對象fig.set_size_inches(30, 20)  # 設置寬度和高度，單位為英寸（inch）ax.imshow(image)cid = fig.canvas.mpl_connect('button_press_event', on_click)plt.show()else:  # 如果使用 必須全局coords = []from segment_anything import SamPredictor, sam_model_registrysam_checkpoint = "sam_vit_h_4b8939.pth"model_type = "vit_h"device = "cuda"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)sam.to(device=device)predictor = SamPredictor(sam)predictor.set_image(image)input_point = np.array(coords)input_label = np.array([1] * len(coords))# plt.figure(figsize=(10, 10))# plt.imshow(image)# show_points(input_point, input_label, plt.gca())# plt.axis('on')# plt.show()input_box = boxif len(coords) == 0:input_point = Noneinput_label = Nonemasks, scores, logits = predictor.predict(point_coords=input_point,point_labels=input_label,box=input_box[None, :],multimask_output=True)if choose_mask:plt.figure(figsize=(60, 20))plt.subplot(1, 3, 1)plt.imshow(image)show_mask(masks[0], plt.gca())# show_points(input_point, input_label, plt.gca())plt.title(f"Mask 0, Score: {scores[0]:.3f}", fontsize=18)plt.subplot(1, 3, 2)plt.imshow(image)show_mask(masks[1], plt.gca())# show_points(input_point, input_label, plt.gca())plt.title(f"Mask 1, Score: {scores[1]:.3f}", fontsize=18)plt.subplot(1, 3, 3)plt.imshow(image)show_mask(masks[2], plt.gca())# show_points(input_point, input_label, plt.gca())plt.title(f"Mask 2, Score: {scores[1]:.3f}", fontsize=18)plt.show()mask_id = int(input())  # 通過輸入idx或者設置特定的idx輸出mask = masks[mask_id]mask = np.tile(np.expand_dims(mask, axis=-1), 3)mask_data = np.where(mask, 255, 0)# mask_image = np.where(mask, image/255, 0.)# plt.figure(figsize=(10, 10))# plt.imshow(mask_image)# plt.show()if click_coords: coords.clear()return mask_dataif __name__ == '__main__':obj = 'shirt'color_path = f'/Data/VolumeDeformData/{obj}/data/'mask_path = f'/Data/VolumeDeformData/{obj}/mask/'if not os.path.exists(mask_path):os.makedirs(mask_path)img_paths = []for extension in ["jpg", "png", "jpeg"]:img_paths += glob.glob(os.path.join(color_path, "*.{}".format(extension)))json_path = 'GroundingDINO-main/box.json'with open(json_path, "r", encoding="utf-8") as f:data = json.load(f)for i in range(len(img_paths) // 2):img_name = f'frame-{str(i).zfill(6)}.color.png'img = cv2.imread(color_path + img_name)id = f'{obj}_{str(i).zfill(6)}'box = np.array(list(map(float, data[id][1:-1].split(','))))mask = get_mask(img, mask_id=2, click_coords=False, choose_mask=False, box=box)cv2.imwrite(mask_path + str(i).zfill(6) + '.png', mask)print(img_name)f.close()