AI分割一切——用OpenVINO?加速M(fèi)eta SAM大模型

ChatGPT的火爆讓大家看到了通用AI大模型的威力，也帶動(dòng)了近期一批自然語(yǔ)言處理（NLP）領(lǐng)域大模型的不斷被推出。你方唱罷我登場(chǎng)，最近，計(jì)算機(jī)視覺(jué)領(lǐng)域也迎來(lái)了自己的物體分割大模型，由Meta開(kāi)源的 “萬(wàn)物可分割 （Segment Anything Model, SAM）”物體分割模型。

物體分割是計(jì)算機(jī)視覺(jué)中的核心任務(wù)之一，旨在識(shí)別圖像中屬于特定對(duì)象的像素。通常實(shí)現(xiàn)圖像分割的方法有兩種，即交互式分割和自動(dòng)分割。交互式分割可以對(duì)任何類別的對(duì)象進(jìn)行分割，但需要人工引導(dǎo)，并通過(guò)反復(fù)精細(xì)化掩碼來(lái)完成。而自動(dòng)分割可以對(duì)預(yù)定義的特定對(duì)象類別進(jìn)行分割，但需要大量手動(dòng)標(biāo)注的對(duì)象進(jìn)行訓(xùn)練，同時(shí)需要大量的計(jì)算資源和具有技術(shù)專業(yè)知識(shí)的人員來(lái)訓(xùn)練分割模型。然而，這兩種方法都沒(méi)有提供一種通用的、完全自動(dòng)的分割方法。

SAM是這兩種方法的泛化，它是一個(gè)單一的模型，可以輕松地執(zhí)行交互式分割和自動(dòng)分割。SAM可以從輸入提示（例如點(diǎn)或框）生成高質(zhì)量的對(duì)象掩碼，并且可以用于生成圖像中所有對(duì)象的掩碼。它已經(jīng)在一個(gè)包含1.1億個(gè)掩碼的1100萬(wàn)個(gè)圖像數(shù)據(jù)集上進(jìn)行了訓(xùn)練，并且在各種分割任務(wù)上具有較強(qiáng)的零樣本性能。它創(chuàng)建了一個(gè)通用的物體分割模型，可以在從分析科學(xué)圖像到編輯照片各種應(yīng)用程序中使用。

圖1. SAM推理結(jié)果示例

這個(gè)強(qiáng)大的通用分割模型，我們的OpenVINO 當(dāng)然也是可以對(duì)它進(jìn)行優(yōu)化以及推理的加速，使其可以方便快速地在英特爾的CPU上部署運(yùn)行起來(lái)。為了方便各位開(kāi)發(fā)者的使用，我們同樣提供了Jupyter Notebook形式的源代碼，大家只需要跟隨我們代碼里的步驟，就可以在自己的機(jī)器上運(yùn)行SAM，對(duì)圖像進(jìn)行任意分割了。

SAM模型由三個(gè)部分組成。

圖像編碼器（Image Encoder）：這是一個(gè)Vision Transformer模型（VIT），使用Masked Auto Encoders方法（MAE）對(duì)圖像進(jìn)行編碼，將圖像轉(zhuǎn)換為嵌入空間。圖像編碼器對(duì)每個(gè)圖像運(yùn)行一次，可以在向模型輸入提示之前應(yīng)用它。

提示編碼器（Prompt Encoder ）：這是一個(gè)用于分割條件的編碼器?？梢允褂靡韵聴l件進(jìn)行分割提示：

1. 點(diǎn)（points）- 與應(yīng)分割的對(duì)象相關(guān)的一組點(diǎn)。Prompt編碼器使用位置編碼將點(diǎn)轉(zhuǎn)換為嵌入值。

2. 框（boxes）- 應(yīng)分割的對(duì)象所在的邊界框。類似于points，邊界框的坐標(biāo)通過(guò)位置編碼來(lái)進(jìn)行編碼。

3. 分割掩碼-由用戶提供的分割掩碼使用卷積進(jìn)行嵌入，并與圖像嵌入進(jìn)行element-wise求和。

4. 文本（text）- 由CLIP模型編碼的文本表示。

掩碼解碼器（Mask Decoder） ：掩碼解碼器有效地將圖像嵌入、提示嵌入和輸出標(biāo)記映射到掩碼。

下圖描述了SAM生成掩碼的流程圖。

向右滑動(dòng)查看完整圖片

接下來(lái)，我們一起來(lái)看看運(yùn)行利用OpenVINO 來(lái)優(yōu)化加速SAM的推理有哪些重點(diǎn)步驟吧。 注意：以下步驟中的所有代碼來(lái)自O(shè)penVINO Notebooks開(kāi)源倉(cāng)庫(kù)中的237-segment-anything notebook 代碼示例，您可以點(diǎn)擊以下鏈接直達(dá)源代碼。 https://github.com/openvinotoolkit/openvino_notebooks/tree/main/notebooks/237-segment-anything

第一步：安裝相應(yīng)工具包、加載模型并轉(zhuǎn)換為OpenVINO IR格式

本次代碼示例需要首先安裝SAM相應(yīng)工具包。

 !pip install -q "segment_anything" "gradio>=3.25"

向右滑動(dòng)查看完整代碼

然后下載及加載相應(yīng)的PyTorch模型。

有幾個(gè)SAM checkpoint可供下載。在本次代碼示例中，我們將使用基于vit_b的模型，但模型加載的方法是通用的，也適用于其他SAM模型。將下面的模型URL、保存checkpoint的路徑和模型類型設(shè)置為對(duì)應(yīng)的SAM模型checkpoint，然后使用SAM_model_registry加載模型。

 import sys
 
 sys.path.append("../utils")
 from notebook_utils import download_file
 
 checkpoint = "sam_vit_b_01ec64.pth"
 model_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
 model_type = "vit_b"
 
 download_file(model_url)

向右滑動(dòng)查看完整代碼

加載模型

 from segment_anything import sam_model_registry
 
 sam = sam_model_registry[model_type](checkpoint=checkpoint)

向右滑動(dòng)查看完整代碼

正如我們已經(jīng)討論過(guò)的，每個(gè)圖像可以使用一次圖像編碼器，然后可以多次運(yùn)行更改提示、提示編碼器和掩碼解碼器來(lái)從同一圖像中檢索不同的對(duì)象?？紤]到這一事實(shí)，我們將模型分為兩個(gè)獨(dú)立的部分：image_encoder和mask_pr預(yù)測(cè)器（提示編碼器和掩碼解碼器的組合）。

第二步：定義圖像編碼器和掩碼預(yù)測(cè)器

圖像編碼器輸入是NCHW格式的形狀為1x3x1024x1024的張量，包含用于分割的圖像。圖像編碼器輸出為圖像嵌入，張量形狀為1x256x64x64。代碼如下

 import warnings
 from pathlib import Path
 import torch
 from openvino.tools import mo
 from openvino.runtime import serialize, Core
 
 core = Core()
 
 ov_encoder_path = Path("sam_image_encoder.xml")
 
 if not ov_encoder_path.exists():
   onnx_encoder_path = ov_encoder_path.with_suffix(".onnx")
   if not onnx_encoder_path.exists():
     with warnings.catch_warnings():
       warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
       warnings.filterwarnings("ignore", category=UserWarning)
       
       torch.onnx.export(sam.image_encoder, torch.zeros(1,3,1024,1024), onnx_encoder_path)
   
   ov_encoder_model = mo.convert_model(onnx_encoder_path, compress_to_fp16=True)
   serialize(ov_encoder_model, str(ov_encoder_path))
 else:
   ov_encoder_model = core.read_model(ov_encoder_path)
 ov_encoder = core.compile_model(ov_encoder_model)

掩碼預(yù)測(cè)器

本次代碼示例需要導(dǎo)出的模型帶有參數(shù)return_single_mask=True。這意味著模型將只返回最佳掩碼，而不是返回多個(gè)掩碼。對(duì)于高分辨率圖像，這可以在放大掩碼開(kāi)銷大的情況下提升運(yùn)行時(shí)速度。

組合提示編碼器和掩碼解碼器模型具有以下輸入列表：

image_embeddings：從image_encoder中嵌入的圖像。具有長(zhǎng)度為1的批索引。

point_coords：稀疏輸入提示的坐標(biāo)，對(duì)應(yīng)于點(diǎn)輸入和框輸入。方框使用兩個(gè)點(diǎn)進(jìn)行編碼，一個(gè)用于左上角，另一個(gè)用于右下角。坐標(biāo)必須已轉(zhuǎn)換為長(zhǎng)邊1024。具有長(zhǎng)度為1的批索引。

point_labels：稀疏輸入提示的標(biāo)簽。0是負(fù)輸入點(diǎn)，1是正輸入點(diǎn)，2是左上角，3是右下角，-1是填充點(diǎn)。*如果沒(méi)有框輸入，則應(yīng)連接標(biāo)簽為-1且坐標(biāo)為（0.0，0.0）的單個(gè)填充點(diǎn)。

模型輸出：

掩碼-預(yù)測(cè)的掩碼大小調(diào)整為原始圖像大小，以獲得二進(jìn)制掩碼，應(yīng)與閾值（通常等于0.0）進(jìn)行比較。

iou_predictions-并集預(yù)測(cè)上的交集。

low_res_masks-后處理之前的預(yù)測(cè)掩碼，可以用作模型的掩碼輸入。

第三步：在交互式分割模式下運(yùn)行OpenVINO 推理

加載分割用的測(cè)試圖片。

 import numpy as np
 import cv2
 import matplotlib.pyplot as plt
 
 download_file("https://raw.githubusercontent.com/facebookresearch/segment-anything/main/notebooks/images/truck.jpg")
 image = cv2.imread('truck.jpg')
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

原始圖片如下，

 plt.figure(figsize=(10,10))
 plt.imshow(image)
 plt.axis('off')
 plt.show()

預(yù)處理及可視化函數(shù)定義：

為圖像編碼器準(zhǔn)備輸入，包含以下步驟：

將BGR圖像轉(zhuǎn)換為RGB

調(diào)整圖像保存縱橫比，其中最長(zhǎng)尺寸等于圖像編碼器輸入尺寸1024。

歸一化圖像減去平均值（123.675、116.28、103.53）并除以標(biāo)準(zhǔn)差（58.395、57.12、57.375）

將HWC數(shù)據(jù)布局轉(zhuǎn)換為CHW并添加批次維度。

根據(jù)圖像編碼器預(yù)期的輸入形狀，按高度或?qū)挾龋ㄈQ于縱橫比）向輸入張量添加零填充。

視頻編碼

要開(kāi)始處理圖像，我們應(yīng)該對(duì)其進(jìn)行預(yù)處理，并使用ov_encoder獲得圖像嵌入。我們將在所有實(shí)驗(yàn)中使用相同的圖像，因此可以運(yùn)行一次視頻編碼、生成一次圖像嵌入，然后重用它們。

 preprocessed_image = preprocess_image(image)
 encoding_results = ov_encoder(preprocessed_image)
 
 image_embeddings = encoding_results[ov_encoder.output(0)]

現(xiàn)在，我們可以嘗試為掩碼生成提供不同的提示。

點(diǎn)輸入舉例

在本例中，我們選擇一個(gè)點(diǎn)作為輸入（input_point）。綠色星形符號(hào)在下圖中顯示了它的位置。

 input_point = np.array([[500, 375]])
 input_label = np.array([1])
 
 plt.figure(figsize=(10,10))
 plt.imshow(image)
 show_points(input_point, input_label, plt.gca())
 plt.axis('off')
 plt.show()

添加一個(gè)批索引，連接一個(gè)填充點(diǎn)，并將其轉(zhuǎn)換為輸入張量坐標(biāo)系。

 coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
 label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)
 coord = resizer.apply_coords(coord, image.shape[:2]).astype(np.float32)

將輸入打包以在掩碼預(yù)測(cè)器中運(yùn)行。

 inputs = {
   "image_embeddings": image_embeddings,
   "point_coords": coord,
   "point_labels": label,
 }

預(yù)測(cè)一個(gè)掩碼并設(shè)置閾值以獲得二進(jìn)制掩碼（0-無(wú)對(duì)象，1-對(duì)象）。

 results = ov_predictor(inputs)
 
 masks = results[ov_predictor.output(0)]
 masks = postprocess_masks(masks, image.shape[:-1])
 masks = masks > 0.0

繪制結(jié)果

 plt.figure(figsize=(10,10))
 plt.imshow(image)
 show_mask(masks, plt.gca())
 show_points(input_point, input_label, plt.gca())
 plt.axis('off')
 plt.show()

多點(diǎn)輸入舉例

 input_point = np.array([[500, 375], [1125, 625], [575, 750])
 input_label = np.array([1, 1, 1])

模型輸入的提示反應(yīng)在測(cè)試圖片上為

 plt.figure(figsize=(10,10))
 plt.imshow(image)
 show_points(input_point, input_label, plt.gca())
 plt.axis('off')
 plt.show()

像上面單點(diǎn)輸入的例子一樣，講輸入點(diǎn)變換為張量坐標(biāo)系，進(jìn)而將輸入打包成所需格式，最后獲得的分割結(jié)果如下圖所示

帶負(fù)標(biāo)簽的框和點(diǎn)輸入

在這個(gè)例中，我們使用邊界框和邊界框內(nèi)的點(diǎn)來(lái)定義輸入提示。邊界框表示為其左上角和右下角的一組點(diǎn)。點(diǎn)的標(biāo)簽0表示該點(diǎn)應(yīng)從掩碼中排除。

 input_box = np.array([425, 600, 700, 875])
 input_point = np.array([[575, 750]])
 input_label = np.array([0])

反應(yīng)在測(cè)試圖片中為

添加批次索引，連接方框和點(diǎn)輸入，為方框角添加適當(dāng)?shù)臉?biāo)簽，然后進(jìn)行變換。本次沒(méi)有填充點(diǎn)，因?yàn)檩斎氚ㄒ粋€(gè)框輸入。

 box_coords = input_box.reshape(2, 2)
 box_labels = np.array([2,3])
 
 coord = np.concatenate([input_point, box_coords], axis=0)[None, :, :]
 label = np.concatenate([input_label, box_labels], axis=0)[None, :].astype(np.float32)
 
 coord = resizer.apply_coords(coord, image.shape[:2]).astype(np.float32)

打包輸入，并進(jìn)行預(yù)測(cè)

 inputs = {
   "image_embeddings": image_embeddings,
   "point_coords": coord,
   "point_labels": label,
 }
 
 results = ov_predictor(inputs)
 
 masks = results[ov_predictor.output(0)]
 masks = postprocess_masks(masks, image.shape[:-1])
 masks = masks > 0.0

結(jié)果如圖所示

第四步：在自動(dòng)分割模式下運(yùn)行OpenVINO 推理

由于SAM可以有效地處理提示，因此可以通過(guò)在圖像上采樣大量提示來(lái)生成整個(gè)圖像的掩碼。automatic_mask_generation函數(shù)實(shí)現(xiàn)了這一功能。它的工作原理是在圖像上的網(wǎng)格中對(duì)單點(diǎn)輸入提示進(jìn)行采樣，SAM可以從每個(gè)提示中預(yù)測(cè)多個(gè)掩碼。然后，對(duì)掩碼進(jìn)行質(zhì)量過(guò)濾，并使用非最大抑制進(jìn)行去重復(fù)。額外的選項(xiàng)允許進(jìn)一步提高掩模的質(zhì)量和數(shù)量，例如對(duì)圖像的多個(gè)裁剪進(jìn)行預(yù)測(cè)，或?qū)ρ谀＿M(jìn)行后處理以去除小的斷開(kāi)區(qū)域和孔洞。

 from segment_anything.utils.amg import (
   MaskData, 
   generate_crop_boxes, 
   uncrop_boxes_xyxy, 
   uncrop_masks, 
   uncrop_points, 
   calculate_stability_score, 
   rle_to_mask, 
   batched_mask_to_box, 
   mask_to_rle_pytorch, 
   is_box_near_crop_edge,
   batch_iterator,
   remove_small_regions,
   build_all_layer_point_grids,
   box_xyxy_to_xywh,
   area_from_rle
)
 from torchvision.ops.boxes import batched_nms, box_area
 from typing import Tuple, List, Dict, Any

在自動(dòng)掩碼生成中有幾個(gè)可調(diào)參數(shù)，用于控制采樣點(diǎn)的密度以及去除低質(zhì)量或重復(fù)掩碼的閾值。此外，生成可以在圖像的裁剪上自動(dòng)運(yùn)行，以提高對(duì)較小對(duì)象的性能，后處理可以去除雜散像素和孔洞。

定義自動(dòng)分割函數(shù)

 def automatic_mask_generation(
   image: np.ndarray, min_mask_region_area: int = 0, points_per_side: int = 32, crop_n_layers: int = 0, crop_n_points_downscale_factor: int = 1, crop_overlap_ratio: float = 512 / 1500, box_nms_thresh: float = 0.7, crop_nms_thresh: float = 0.7
) -> List[Dict[str, Any]]:
   """
   Generates masks for the given image.
   
   Arguments:
    image (np.ndarray): The image to generate masks for, in HWC uint8 format.
   
   Returns:
     list(dict(str, any)): A list over records for masks. Each record is
      a dict containing the following keys:
       segmentation (dict(str, any) or np.ndarray): The mask. If
        output_mode='binary_mask', is an array of shape HW. Otherwise,
        is a dictionary containing the RLE.
       bbox (list(float)): The box around the mask, in XYWH format.
       area (int): The area in pixels of the mask.
       predicted_iou (float): The model's own prediction of the mask's
        quality. This is filtered by the pred_iou_thresh parameter.
       point_coords (list(list(float))): The point coordinates input
        to the model to generate this mask.
       stability_score (float): A measure of the mask's quality. This
        is filtered on using the stability_score_thresh parameter.
       crop_box (list(float)): The crop of the image used to generate
        the mask, given in XYWH format.
   """
   point_grids = build_all_layer_point_grids(
     points_per_side,
     crop_n_layers,
     crop_n_points_downscale_factor,
   )
   mask_data = generate_masks(
     image, point_grids, crop_n_layers, crop_overlap_ratio, crop_nms_thresh)
   
   # Filter small disconnected regions and holes in masks
   if min_mask_region_area > 0:
     mask_data = postprocess_small_regions(
       mask_data,
       min_mask_region_area,
       max(box_nms_thresh, crop_nms_thresh),
     )
   
   mask_data["segmentations"] = [
     rle_to_mask(rle) for rle in mask_data["rles"]]
   
   # Write mask records
   curr_anns = []
   for idx in range(len(mask_data["segmentations"])):
     ann = {
       "segmentation": mask_data["segmentations"][idx],
       "area": area_from_rle(mask_data["rles"][idx]),
       "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),
       "predicted_iou": mask_data["iou_preds"][idx].item(),
       "point_coords": [mask_data["points"][idx].tolist()],
       "stability_score": mask_data["stability_score"][idx].item(),
       "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),
     }
     curr_anns.append(ann)
   
   return curr_anns

運(yùn)行自動(dòng)分割預(yù)測(cè)

 prediction = automatic_mask_generation(image)

以上automatic_mask_generation函數(shù)返回一個(gè)掩碼列表，其中每個(gè)掩碼都是一個(gè)包含有關(guān)掩碼的各種數(shù)據(jù)的字典：

分割：掩碼

面積：掩碼的面積（以像素為單位）

bbox：XYWH格式的掩碼的邊界框

predicted_out：模型自己對(duì)掩模質(zhì)量的預(yù)測(cè)

point_coords：生成此掩碼的采樣輸入點(diǎn)

穩(wěn)定性核心：衡量掩碼質(zhì)量的一個(gè)附加指標(biāo)

crop_box：用于以XYWH格式生成此掩碼的圖像的裁剪

查看掩碼的信息

 print(f"Number of detected masks: {len(prediction)}")
 print(f"Annotation keys: {prediction[0].keys()}")

獲得如下結(jié)果

繪制最后的分割結(jié)果

 from tqdm.notebook import tqdm
 
 def draw_anns(image, anns):
   if len(anns) == 0:
     return
   segments_image = image.copy()
   sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
   for ann in tqdm(sorted_anns):
     mask = ann["segmentation"]
     mask_color = np.random.randint(0, 255, size=(1, 1, 3)).astype(np.uint8)
     segments_image[mask] = mask_color
   return cv2.addWeighted(image.astype(np.float32), 0.7, segments_image.astype(np.float32), 0.3, 0.0)

 import PIL
 
 out = draw_anns(image, prediction)
 cv2.imwrite("result.png", out[:, :, ::-1])
 
 PIL.Image.open("result.png")

看看這些分割的效果，是不是非常的驚艷呢。其實(shí)除了以上我們介紹的代碼內(nèi)容，在我們的Jupyter Notebook代碼里，還為大家提供了窗口式鼠標(biāo)點(diǎn)擊輸入提示的交互式分割體驗(yàn)，甚至可以在手機(jī)端輸入U(xiǎn)RL地址體驗(yàn)即時(shí)的互動(dòng)效果，如下圖所示

這么多有趣又快速的OpenVINO運(yùn)行物體分割的方式，快在你本地的機(jī)器上克隆我們的代碼示例，自己動(dòng)手試試SAM的效果吧。

小結(jié)：

整個(gè)的步驟就是這樣！現(xiàn)在就開(kāi)始跟著我們提供的代碼和步驟，動(dòng)手試試用OpenVINO和SAM吧。

關(guān)于英特爾OpenVINO 開(kāi)源工具套件的詳細(xì)資料，包括其中我們提供的三百多個(gè)經(jīng)驗(yàn)證并優(yōu)化的預(yù)訓(xùn)練模型的詳細(xì)資料，請(qǐng)您點(diǎn)擊https://www.intel.com/content/www/us/en/developer/tools/openvino-toolkit/overview.html

除此之外，為了方便大家了解并快速掌握OpenVINO 的使用，我們還提供了一系列開(kāi)源的Jupyter notebook demo。運(yùn)行這些notebook，就能快速了解在不同場(chǎng)景下如何利用OpenVINO 實(shí)現(xiàn)一系列、包括計(jì)算機(jī)視覺(jué)、語(yǔ)音及自然語(yǔ)言處理任務(wù)。OpenVINO notebooks的資源可以在GitHub這里下載安裝：

https://github.com/openvinotoolkit/openvino_notebooks 。

審核編輯 ：李倩