AI分割一切——用OpenVINO?加速Meta SAM大模型

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

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

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

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

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

SAM模型由三個部分組成。

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

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

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

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

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

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

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

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

向右滑動查看完整圖片

接下來，我們一起來看看運行利用OpenVINO 來優(yōu)化加速SAM的推理有哪些重點步驟吧。 注意：以下步驟中的所有代碼來自O(shè)penVINO Notebooks開源倉庫中的237-segment-anything notebook 代碼示例，您可以點擊以下鏈接直達(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"

向右滑動查看完整代碼

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

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

向右滑動查看完整代碼

加載模型

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

向右滑動查看完整代碼

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

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

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

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

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

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

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

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

模型輸出：

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

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

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

第三步：在交互式分割模式下運行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()

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

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

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

調(diào)整圖像保存縱橫比，其中最長尺寸等于圖像編碼器輸入尺寸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于縱橫比）向輸入張量添加零填充。

視頻編碼

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

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

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

點輸入舉例

在本例中，我們選擇一個點作為輸入（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()

添加一個批索引，連接一個填充點，并將其轉(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ù)測器中運行。

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

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

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

多點輸入舉例

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

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

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

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

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

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

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

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

添加批次索引，連接方框和點輸入，為方框角添加適當(dāng)?shù)臉?biāo)簽，然后進(jìn)行變換。本次沒有填充點，因為輸入包括一個框輸入。

 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ù)測

 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é)果如圖所示

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

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

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

定義自動分割函數(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ù)測

 prediction = automatic_mask_generation(image)

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

分割：掩碼

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

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

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

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

穩(wěn)定性核心：衡量掩碼質(zhì)量的一個附加指標(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")

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

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

小結(jié)：

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

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

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

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

審核編輯 ：李倩