使用 CLIPSeg 進行零樣本影像分割

本指南展示瞭如何使用 CLIPSeg（一個零樣本影像分割模型）和 🤗 transformers。CLIPSeg 建立的粗略分割掩碼可用於機器人感知、影像修復和許多其他任務。如果您需要更精確的分割掩碼，我們將展示如何在 Segments.ai 上最佳化 CLIPSeg 的結果。

影像分割是計算機視覺領域的一個眾所周知的任務。它使計算機不僅能知道影像中有什麼（分類），物體在影像中的位置（檢測），還能知道這些物體的輪廓。瞭解物體的輪廓在機器人和自動駕駛等領域至關重要。例如，機器人必須知道物體的形狀才能正確抓取它。分割還可以與 影像修復 結合，允許使用者描述他們想要替換的影像部分。

大多數影像分割模型的一個限制是它們只適用於固定的類別列表。例如，您不能簡單地使用在橙子上訓練的分割模型來分割蘋果。要教授分割模型一個額外的類別，您必須標記新類別的資料並訓練一個新模型，這可能成本高昂且耗時。但是，如果有一個模型已經可以分割幾乎任何種類的物體，而無需進一步訓練，那會怎麼樣？這正是 CLIPSeg（一個零樣本分割模型）所實現的。

目前，CLIPSeg 仍有其侷限性。例如，該模型使用 352 x 352 畫素的影像，因此輸出解析度相當低。這意味著當我們使用現代相機拍攝的影像時，我們無法期望獲得畫素級完美的結果。如果我們需要更精確的分割，我們可以對最先進的分割模型進行微調，如 我們之前的部落格文章 所示。在這種情況下，我們仍然可以使用 CLIPSeg 生成一些粗略的標籤，然後在使用 Segments.ai 等標註工具中對其進行最佳化。在我們描述如何做到這一點之前，我們首先看看 CLIPSeg 的工作原理。

CLIP：CLIPSeg 背後的神奇模型

CLIP，全稱是 Contrastive Language–Image Pre-training（對比語言-影像預訓練），是 OpenAI 於 2021 年開發的一個模型。您可以給 CLIP 輸入一張影像或一段文字，CLIP 將輸出您輸入的抽象表示。這種抽象表示，也稱為嵌入，實際上只是一個向量（一串數字）。您可以將此向量視為高維空間中的一個點。CLIP 經過訓練，使得相似影像和文字的表示也相似。這意味著如果我們輸入一張影像和一段與該影像匹配的文字描述，影像和文字的表示將相似（即，高維點將彼此靠近）。

起初，這可能看起來不是很有用，但它實際上非常強大。舉個例子，讓我們快速看看 CLIP 如何用於影像分類，而從未在該任務上進行過訓練。為了對影像進行分類，我們將影像和我們想要選擇的不同類別輸入到 CLIP（例如，我們輸入一張影像和“蘋果”、“橙子”等詞）。然後 CLIP 返回影像和每個類別的嵌入。現在，我們只需檢查哪個類別嵌入與影像的嵌入最接近，瞧！感覺像魔法，不是嗎？

更重要的是，CLIP 不僅對分類有用，它還可以用於影像搜尋（您能看出這與分類有何相似之處嗎？）、文字到影像模型（DALL-E 2 由 CLIP 提供支援）、物件檢測（OWL-ViT），以及對我們來說最重要的是：影像分割。現在您明白為什麼 CLIP 確實是機器學習領域的一個突破了。

CLIP 之所以效果如此好，是因為該模型在包含文字標題的大型影像資料集上進行了訓練。該資料集包含多達 4 億張從網際網路上獲取的影像-文字對。這些影像包含各種各樣的物件和概念，CLIP 在為它們中的每一個建立表示方面表現出色。

CLIPSeg：使用 CLIP 進行影像分割

CLIPSeg 是一個使用 CLIP 表示來建立影像分割掩碼的模型。它由 Timo Lüddecke 和 Alexander Ecker 釋出。他們透過在 CLIP 模型之上訓練一個基於 Transformer 的解碼器來實現零樣本影像分割，該解碼器保持凍結狀態。解碼器接收影像的 CLIP 表示以及您要分割的事物的 CLIP 表示。使用這兩個輸入，CLIPSeg 解碼器建立一個二進位制分割掩碼。更精確地說，解碼器不僅使用我們要分割的影像的最終 CLIP 表示，它還使用 CLIP 的某些層的輸出。

該解碼器在 PhraseCut 資料集上進行訓練，該資料集包含超過 340,000 個短語及其對應的影像分割掩碼。作者還嘗試了各種增強方法來擴大資料集的大小。這裡的目標不僅是能夠分割資料集中存在的類別，還要能夠分割未見過的類別。實驗確實表明解碼器可以泛化到未見過的類別。

CLIPSeg 的一個有趣特性是，查詢（我們要分割的影像）和提示（我們要分割的影像中的事物）都作為 CLIP 嵌入輸入。提示的 CLIP 嵌入可以來自一段文字（類別名稱），**或者來自另一張影像**。這意味著您可以透過向 CLIPSeg 提供一個橙子的示例影像來分割照片中的橙子。

這種技術被稱為“視覺提示”，當您想要分割的事物難以描述時，它非常有用。例如，如果您想分割 T 恤圖片中的徽標，描述徽標的形狀並不容易，但 CLIPSeg 允許您只需使用徽標的影像作為提示。

CLIPSeg 論文包含了一些關於提高視覺提示有效性的技巧。他們發現裁剪查詢影像（使其僅包含您要分割的物件）非常有幫助。模糊和調暗查詢影像的背景也有一些幫助。在下一節中，我們將展示如何使用 🤗 transformers 親自嘗試視覺提示。

將 CLIPSeg 與 Hugging Face Transformers 結合使用

使用 Hugging Face Transformers，您可以輕鬆下載並在影像上執行預訓練的 CLIPSeg 模型。讓我們從安裝 transformers 開始。

!pip install -q transformers

要下載模型，只需例項化它即可。

from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation

processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")

現在我們可以載入一張影像來嘗試分割。我們將選擇一張由 Calum Lewis 拍攝的美味早餐圖片。

from PIL import Image
import requests

url = "https://unsplash.com/photos/8Nc_oQsc2qQ/download?ixid=MnwxMjA3fDB8MXxhbGx8fHx8fHx8fHwxNjcxMjAwNzI0&force=true&w=640"
image = Image.open(requests.get(url, stream=True).raw)
image

文字提示

讓我們從定義一些我們想要分割的文字類別開始。

prompts = ["cutlery", "pancakes", "blueberries", "orange juice"]

現在我們有了輸入，我們可以處理它們並將它們輸入到模型中。

import torch

inputs = processor(text=prompts, images=[image] * len(prompts), padding="max_length", return_tensors="pt")
# predict
with torch.no_grad():
  outputs = model(**inputs)
preds = outputs.logits.unsqueeze(1)

最後，讓我們視覺化輸出。

import matplotlib.pyplot as plt

_, ax = plt.subplots(1, len(prompts) + 1, figsize=(3*(len(prompts) + 1), 4))
[a.axis('off') for a in ax.flatten()]
ax[0].imshow(image)
[ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(len(prompts))];
[ax[i+1].text(0, -15, prompt) for i, prompt in enumerate(prompts)];

視覺提示

如前所述，我們也可以使用影像作為輸入提示（即代替類別名稱）。如果難以描述您要分割的事物，這會特別有用。在此示例中，我們將使用由 Daniel Hooper 拍攝的咖啡杯圖片。

url = "https://unsplash.com/photos/Ki7sAc8gOGE/download?ixid=MnwxMjA3fDB8MXxzZWFyY2h8MTJ8fGNvZmZlJTIwdG8lMjBnb3xlbnwwfHx8fDE2NzExOTgzNDQ&force=true&w=640"
prompt = Image.open(requests.get(url, stream=True).raw)
prompt

現在我們可以處理輸入影像和提示影像，並將它們輸入到模型中。

encoded_image = processor(images=[image], return_tensors="pt")
encoded_prompt = processor(images=[prompt], return_tensors="pt")
# predict
with torch.no_grad():
  outputs = model(**encoded_image, conditional_pixel_values=encoded_prompt.pixel_values)
preds = outputs.logits.unsqueeze(1)
preds = torch.transpose(preds, 0, 1)

然後，我們可以像以前一樣視覺化結果。

_, ax = plt.subplots(1, 2, figsize=(6, 4))
[a.axis('off') for a in ax.flatten()]
ax[0].imshow(image)
ax[1].imshow(torch.sigmoid(preds[0]))

讓我們最後一次嘗試使用論文中描述的視覺提示技巧，即裁剪影像和調暗背景。

url = "https://i.imgur.com/mRSORqz.jpg"
alternative_prompt = Image.open(requests.get(url, stream=True).raw)
alternative_prompt

encoded_alternative_prompt = processor(images=[alternative_prompt], return_tensors="pt")
# predict
with torch.no_grad():
  outputs = model(**encoded_image, conditional_pixel_values=encoded_alternative_prompt.pixel_values)
preds = outputs.logits.unsqueeze(1)
preds = torch.transpose(preds, 0, 1)

_, ax = plt.subplots(1, 2, figsize=(6, 4))
[a.axis('off') for a in ax.flatten()]
ax[0].imshow(image)
ax[1].imshow(torch.sigmoid(preds[0]))

在這種情況下，結果幾乎相同。這可能是因為咖啡杯在原始影像中已經與背景很好地分離了。

使用 CLIPSeg 在 Segments.ai 上預標註影像

如您所見，CLIPSeg 的結果有點模糊且解析度很低。如果我們需要獲得更好的結果，您可以微調最先進的分割模型，如 我們之前的部落格文章 中所述。為了微調模型，我們需要標註資料。在本節中，我們將向您展示如何使用 CLIPSeg 建立一些粗略的分割掩碼，然後在 Segments.ai（一個帶有智慧影像分割標註工具的標註平臺）上對其進行最佳化。

首先，在 https://segments.ai/join 建立一個帳戶並安裝 Segments Python SDK。然後，您可以使用 API 金鑰初始化 Segments.ai Python 客戶端。此金鑰可以在 帳戶頁面 上找到。

!pip install -q segments-ai

from segments import SegmentsClient
from getpass import getpass

api_key = getpass('Enter your API key: ')
segments_client = SegmentsClient(api_key)

接下來，讓我們使用 Segments 客戶端從資料集中載入一張影像。我們將使用 a2d2 自動駕駛資料集。您也可以按照 這些說明 建立自己的資料集。

samples = segments_client.get_samples("admin-tobias/clipseg")

# Use the last image as an example
sample = samples[1]
image = Image.open(requests.get(sample.attributes.image.url, stream=True).raw)
image

我們還需要從資料集屬性中獲取類別名稱。

dataset = segments_client.get_dataset("admin-tobias/clipseg")
category_names = [category.name for category in dataset.task_attributes.categories]

現在我們可以像以前一樣在影像上使用 CLIPSeg。這次，我們還會將輸出放大，使其與輸入影像的大小匹配。

from torch import nn

inputs = processor(text=category_names, images=[image] * len(category_names), padding="max_length", return_tensors="pt")

# predict
with torch.no_grad():
  outputs = model(**inputs)

# resize the outputs
preds = nn.functional.interpolate(
    outputs.logits.unsqueeze(1),
    size=(image.size[1], image.size[0]),
    mode="bilinear"
)

然後我們可以再次視覺化結果。

len_cats = len(category_names)
_, ax = plt.subplots(1, len_cats + 1, figsize=(3*(len_cats + 1), 4))
[a.axis('off') for a in ax.flatten()]
ax[0].imshow(image)
[ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(len_cats)];
[ax[i+1].text(0, -15, category_name) for i, category_name in enumerate(category_names)];

現在我們必須將預測結果組合成一個單一的分割影像。我們只需對每個補丁選擇 Sigmoid 值最大的類別即可。我們還會確保低於某個閾值的所有值都不計入在內。

threshold = 0.1

flat_preds = torch.sigmoid(preds.squeeze()).reshape((preds.shape[0], -1))

# Initialize a dummy "unlabeled" mask with the threshold
flat_preds_with_treshold = torch.full((preds.shape[0] + 1, flat_preds.shape[-1]), threshold)
flat_preds_with_treshold[1:preds.shape[0]+1,:] = flat_preds

# Get the top mask index for each pixel
inds = torch.topk(flat_preds_with_treshold, 1, dim=0).indices.reshape((preds.shape[-2], preds.shape[-1]))

讓我們快速視覺化結果。

plt.imshow(inds)

最後，我們可以將預測上傳到 Segments.ai。為此，我們首先將點陣圖轉換為 png 檔案，然後將此檔案上傳到 Segments，最後將標籤新增到樣本中。

from segments.utils import bitmap2file
import numpy as np

inds_np = inds.numpy().astype(np.uint32)
unique_inds = np.unique(inds_np).tolist()
f = bitmap2file(inds_np, is_segmentation_bitmap=True)

asset = segments_client.upload_asset(f, "clipseg_prediction.png")

attributes = {
      'format_version': '0.1',
      'annotations': [{"id": i, "category_id": i} for i in unique_inds if i != 0],
      'segmentation_bitmap': { 'url': asset.url },
  }

segments_client.add_label(sample.uuid, 'ground-truth', attributes)

如果你檢視 Segments.ai 上上傳的預測，你會發現它並不完美。但是，你可以手動糾正最大的錯誤，然後你可以使用糾正後的資料集來訓練一個比 CLIPSeg 更好的模型。

總結

CLIPSeg 是一種零樣本分割模型，可以使用文字和影像提示。該模型在 CLIP 的基礎上添加了一個解碼器，幾乎可以分割任何東西。然而，目前輸出的分割掩碼解析度仍然很低，因此如果精度很重要，您可能仍然需要微調不同的分割模型。

請注意，目前正在進行更多關於零樣本分割的研究，因此預計在不久的將來會新增更多模型。一個例子是 GroupViT，它已在 🤗 Transformers 中可用。要了解分割研究的最新新聞，您可以關注我們的 Twitter：@TobiasCornille、@NielsRogge 和 @huggingface。

如果您有興趣瞭解如何微調最先進的分割模型，請檢視我們之前的部落格文章：https://huggingface.co/blog/fine-tune-segformer。