如何用 Hugging Face Transformers 使用 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 自行尝试视觉提示。

使用 Hugging Face Transformers 的 CLIPSeg

使用 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]))

在这种情况下，结果几乎相同。这可能是因为在原始图像中咖啡杯已经与背景很好地分开了。

在 Segments.ai 上使用 CLIPSeg 预标注图像

如你所见，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。