ChatGPT的火爆讓大家看到了通用AI大模型的威力，也帶動了近期一批自然語言處理（NLP）領域大模型的不斷被推出。你方唱罷我登場，最近，計算機視覺領域也迎來了自己的物體分割大模型，由Meta開源的 “萬物可分割 （Segment Anything Model, SAM）”物體分割模型。

物體分割是計算機視覺中的核心任務之一，旨在識別圖像中屬于特定對象的像素。通常實現圖像分割的方法有兩種，即交互式分割和自動分割。交互式分割可以對任何類別的對象進行分割，但需要人工引導，并通過反復精細化掩碼來完成。而自動分割可以對預定義的特定對象類別進行分割，但需要大量手動標注的對象進行訓練，同時需要大量的計算資源和具有技術專業知識的人員來訓練分割模型。然而，這兩種方法都沒有提供一種通用的、完全自動的分割方法。

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

圖1. SAM推理結果示例

這個強大的通用分割模型，我們的OpenVINO 當然也是可以對它進行優化以及推理的加速，使其可以方便快速地在英特爾的CPU上部署運行起來。為了方便各位開發者的使用，我們同樣提供了Jupyter Notebook形式的源代碼，大家只需要跟隨我們代碼里的步驟，就可以在自己的機器上運行SAM，對圖像進行任意分割了。

SAM模型由三個部分組成。

圖像編碼器（Image Encoder）：這是一個Vision Transformer模型（VIT），使用Masked Auto Encoders方法（MAE）對圖像進行編碼，將圖像轉換為嵌入空間。圖像編碼器對每個圖像運行一次，可以在向模型輸入提示之前應用它。

提示編碼器（Prompt Encoder ）：這是一個用于分割條件的編碼器。可以使用以下條件進行分割提示：

1. 點（points）- 與應分割的對象相關的一組點。Prompt編碼器使用位置編碼將點轉換為嵌入值。

2. 框（boxes）- 應分割的對象所在的邊界框。類似于points，邊界框的坐標通過位置編碼來進行編碼。

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

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

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

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

向右滑動查看完整圖片

接下來，我們一起來看看運行利用OpenVINO 來優化加速SAM的推理有哪些重點步驟吧。 注意：以下步驟中的所有代碼來自OpenVINO Notebooks開源倉庫中的237-segment-anything notebook 代碼示例，您可以點擊以下鏈接直達源代碼。 https://github.com/openvinotoolkit/openvino_notebooks/tree/main/notebooks/237-segment-anything

第一步：安裝相應工具包、加載模型并轉換為OpenVINO IR格式

本次代碼示例需要首先安裝SAM相應工具包。

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

向右滑動查看完整代碼

然后下載及加載相應的PyTorch模型。

有幾個SAM checkpoint可供下載。在本次代碼示例中，我們將使用基于vit_b的模型，但模型加載的方法是通用的，也適用于其他SAM模型。將下面的模型URL、保存checkpoint的路徑和模型類型設置為對應的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)

向右滑動查看完整代碼

加載模型

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

向右滑動查看完整代碼

正如我們已經討論過的，每個圖像可以使用一次圖像編碼器，然后可以多次運行更改提示、提示編碼器和掩碼解碼器來從同一圖像中檢索不同的對象。考慮到這一事實，我們將模型分為兩個獨立的部分：image_encoder和mask_pr預測器（提示編碼器和掩碼解碼器的組合）。

第二步：定義圖像編碼器和掩碼預測器

圖像編碼器輸入是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)

掩碼預測器

本次代碼示例需要導出的模型帶有參數return_single_mask=True。這意味著模型將只返回最佳掩碼，而不是返回多個掩碼。對于高分辨率圖像，這可以在放大掩碼開銷大的情況下提升運行時速度。

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

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

point_coords：稀疏輸入提示的坐標，對應于點輸入和框輸入。方框使用兩個點進行編碼，一個用于左上角，另一個用于右下角。坐標必須已轉換為長邊1024。具有長度為1的批索引。

point_labels：稀疏輸入提示的標簽。0是負輸入點，1是正輸入點，2是左上角，3是右下角，-1是填充點。*如果沒有框輸入，則應連接標簽為-1且坐標為（0.0，0.0）的單個填充點。

模型輸出：

掩碼-預測的掩碼大小調整為原始圖像大小，以獲得二進制掩碼，應與閾值（通常等于0.0）進行比較。

iou_predictions-并集預測上的交集。

low_res_masks-后處理之前的預測掩碼，可以用作模型的掩碼輸入。

第三步：在交互式分割模式下運行OpenVINO 推理

加載分割用的測試圖片。

 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()

預處理及可視化函數定義：

為圖像編碼器準備輸入，包含以下步驟：

將BGR圖像轉換為RGB

調整圖像保存縱橫比，其中最長尺寸等于圖像編碼器輸入尺寸1024。

歸一化圖像減去平均值（123.675、116.28、103.53）并除以標準差（58.395、57.12、57.375）

將HWC數據布局轉換為CHW并添加批次維度。

根據圖像編碼器預期的輸入形狀，按高度或寬度（取決于縱橫比）向輸入張量添加零填充。

視頻編碼

要開始處理圖像，我們應該對其進行預處理，并使用ov_encoder獲得圖像嵌入。我們將在所有實驗中使用相同的圖像，因此可以運行一次視頻編碼、生成一次圖像嵌入，然后重用它們。

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

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

點輸入舉例

在本例中，我們選擇一個點作為輸入（input_point）。綠色星形符號在下圖中顯示了它的位置。

 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()

添加一個批索引，連接一個填充點，并將其轉換為輸入張量坐標系。

 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)

將輸入打包以在掩碼預測器中運行。

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

預測一個掩碼并設置閾值以獲得二進制掩碼（0-無對象，1-對象）。

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

繪制結果

 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()

多點輸入舉例

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

模型輸入的提示反應在測試圖片上為

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

像上面單點輸入的例子一樣，講輸入點變換為張量坐標系，進而將輸入打包成所需格式，最后獲得的分割結果如下圖所示

帶負標簽的框和點輸入

在這個例中，我們使用邊界框和邊界框內的點來定義輸入提示。邊界框表示為其左上角和右下角的一組點。點的標簽0表示該點應從掩碼中排除。

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

反應在測試圖片中為

添加批次索引，連接方框和點輸入，為方框角添加適當的標簽，然后進行變換。本次沒有填充點，因為輸入包括一個框輸入。

 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)

打包輸入，并進行預測

 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

結果如圖所示

第四步：在自動分割模式下運行OpenVINO 推理

由于SAM可以有效地處理提示，因此可以通過在圖像上采樣大量提示來生成整個圖像的掩碼。automatic_mask_generation函數實現了這一功能。它的工作原理是在圖像上的網格中對單點輸入提示進行采樣，SAM可以從每個提示中預測多個掩碼。然后，對掩碼進行質量過濾，并使用非最大抑制進行去重復。額外的選項允許進一步提高掩模的質量和數量，例如對圖像的多個裁剪進行預測，或對掩模進行后處理以去除小的斷開區域和孔洞。

 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

在自動掩碼生成中有幾個可調參數，用于控制采樣點的密度以及去除低質量或重復掩碼的閾值。此外，生成可以在圖像的裁剪上自動運行，以提高對較小對象的性能，后處理可以去除雜散像素和孔洞。

定義自動分割函數

 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

運行自動分割預測

 prediction = automatic_mask_generation(image)

以上automatic_mask_generation函數返回一個掩碼列表，其中每個掩碼都是一個包含有關掩碼的各種數據的字典：

分割：掩碼

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

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

predicted_out：模型自己對掩模質量的預測

point_coords：生成此掩碼的采樣輸入點

穩定性核心：衡量掩碼質量的一個附加指標

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

查看掩碼的信息

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

獲得如下結果

繪制最后的分割結果

 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")

看看這些分割的效果，是不是非常的驚艷呢。其實除了以上我們介紹的代碼內容，在我們的Jupyter Notebook代碼里，還為大家提供了窗口式鼠標點擊輸入提示的交互式分割體驗，甚至可以在手機端輸入URL地址體驗即時的互動效果，如下圖所示

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

小結：

整個的步驟就是這樣！現在就開始跟著我們提供的代碼和步驟，動手試試用OpenVINO和SAM吧。

關于英特爾OpenVINO 開源工具套件的詳細資料，包括其中我們提供的三百多個經驗證并優化的預訓練模型的詳細資料，請您點擊https://www.intel.com/content/www/us/en/developer/tools/openvino-toolkit/overview.html

除此之外，為了方便大家了解并快速掌握OpenVINO 的使用，我們還提供了一系列開源的Jupyter notebook demo。運行這些notebook，就能快速了解在不同場景下如何利用OpenVINO 實現一系列、包括計算機視覺、語音及自然語言處理任務。OpenVINO notebooks的資源可以在GitHub這里下載安裝：

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

審核編輯 ：李倩