使用 Mask2Former 和 OneFormer 進行通用影像分割

釋出日期：2023 年 1 月 19 日

在 GitHub 上更新

贊

本指南介紹了 Mask2Former 和 OneFormer，這是兩種最先進的影像分割神經網路模型。這些模型現已在 🤗 transformers 中提供，這是一個開源庫，提供易於使用的最先進模型實現。在此過程中，您將瞭解各種影像分割形式之間的區別。

影像分割

影像分割是將影像中不同的“區域”進行識別的任務，例如人或汽車。更專業地說，影像分割是將具有不同語義的畫素分組的任務。有關簡要介紹，請參閱 Hugging Face 任務頁面。

影像分割大致可以分為 3 個子任務——例項分割、語義分割和全景分割——每個子任務都有許多方法和模型架構。

例項分割 是識別影像中不同“例項”（如個體人物）的任務。例項分割與目標檢測非常相似，只是我們希望輸出一組二值分割掩碼，而不是邊界框，並帶有相應的類標籤。例項通常也稱為“物件”或“事物”。請注意，單個例項可能會重疊。
語義分割 是識別影像中每個畫素的不同“語義類別”（如“人”或“天空”）的任務。與例項分割相反，語義分割不區分給定語義類別的單個例項；例如，只希望為“人”類別而不是為單個個體建立掩碼。沒有單個例項的語義類別，如“天空”或“草地”，通常被稱為“背景物”，以區別於“事物”（很棒的名字，是吧？）。請注意，語義類別之間不可能有重疊，因為每個畫素都屬於一個類別。
全景分割，由 Kirillov 等人於 2018 年引入，旨在透過讓模型簡單地識別一組“段”，每個段都有相應的二值掩碼和類標籤，從而統一例項分割和語義分割。段可以是“事物”或“背景物”。與例項分割不同，不同段之間不可能重疊。

下圖說明了 3 個子任務之間的區別（取自此部落格文章）。

drawing

在過去幾年中，研究人員提出了幾種架構，這些架構通常非常適合例項分割、語義分割或全景分割。例項分割和全景分割通常透過為每個物件例項輸出一組二值掩碼 + 相應的標籤來解決（與目標檢測非常相似，只是為每個例項輸出二值掩碼而不是邊界框）。這通常被稱為“二值掩碼分類”。另一方面，語義分割通常透過讓模型輸出一個帶有每個畫素一個標籤的“分割圖”來解決。因此，語義分割被視為“逐畫素分類”問題。採用此正規化的流行語義分割模型包括 SegFormer（我們為此撰寫了一篇詳細的部落格文章）和 UPerNet。

通用影像分割

幸運的是，自 2020 年左右以來，人們開始提出能夠使用統一架構和相同正規化解決所有 3 個任務（例項、語義和全景分割）的模型。這始於 DETR，它是第一個使用“二值掩碼分類”正規化解決全景分割的模型，透過統一處理“事物”和“背景物”類別。關鍵創新是讓 Transformer 解碼器以並行方式生成一組二值掩碼 + 類別。這在 MaskFormer 論文中得到了改進，該論文表明“二值掩碼分類”正規化也適用於語義分割。

Mask2Former 透過進一步改進神經網路架構將其擴充套件到例項分割。因此，我們已經從獨立的架構發展到研究人員現在所稱的“通用影像分割”架構，能夠解決任何影像分割任務。有趣的是，這些通用模型都採用了“掩碼分類”正規化，完全放棄了“逐畫素分類”正規化。下圖描繪了 Mask2Former 的架構（取自原始論文）。

drawing

簡而言之，影像首先透過骨幹網路（在論文中可以是 ResNet 或 Swin Transformer）以獲得低解析度特徵圖。接下來，使用畫素解碼器模組增強這些特徵圖以獲得高解析度特徵。最後，Transformer 解碼器接收一組查詢，並根據畫素解碼器的特徵將其轉換為一組二值掩碼和類別預測。

請注意，Mask2Former 仍然需要針對每個任務單獨訓練才能獲得最先進的結果。這已透過 OneFormer 模型得到改進，該模型僅透過訓練資料集的全景版本（！）即可在所有 3 個任務上獲得最先進的效能，透過新增文字編碼器來根據“例項”、“語義”或“全景”輸入對模型進行條件化。該模型目前也在 🤗 transformers 中提供。它比 Mask2Former 更準確，但由於額外的文字編碼器，延遲更高。有關 OneFormer 的概述，請參閱下圖。它利用 Swin Transformer 或新的 DiNAT 模型作為骨幹網路。

drawing

使用 Transformers 中的 Mask2Former 和 OneFormer 進行推理

Mask2Former 和 OneFormer 的用法相當簡單，與它們的先行者 MaskFormer 非常相似。讓我們從 Hub 例項化一個在 COCO 全景資料集上訓練的 Mask2Former 模型，以及它的處理器。請注意，作者釋出了不下 30 個檢查點，這些檢查點在各種資料集上進行了訓練。

from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation

processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")

接下來，讓我們載入 COCO 資料集中熟悉的貓影像，我們將對其執行推理。

from PIL import Image

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image

我們使用影像處理器為模型準備影像，然後將其透過模型。

inputs = processor(image, return_tensors="pt")

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

模型輸出一組二值掩碼和相應的類 logits。Mask2Former 的原始輸出可以很容易地使用影像處理器進行後處理，以獲得最終的例項、語義或全景分割預測。

prediction = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
print(prediction.keys())

Output:
----------------------------------------------------------------------------------------------------
dict_keys(['segmentation', 'segments_info'])

在全景分割中，最終的 prediction 包含兩部分：一個形狀為 (高, 寬) 的 segmentation 圖，其中每個值編碼給定畫素的例項 ID，以及相應的 segments_info。segments_info 包含有關圖中單個段的更多資訊（例如它們的類別/類別 ID）。請注意，Mask2Former 為了效率輸出形狀為 (96, 96) 的二值掩碼提案，並且 target_sizes 引數用於將最終掩碼調整到原始影像大小。

讓我們將結果視覺化

from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm

def draw_panoptic_segmentation(segmentation, segments_info):
    # get the used color map
    viridis = cm.get_cmap('viridis', torch.max(segmentation))
    fig, ax = plt.subplots()
    ax.imshow(segmentation)
    instances_counter = defaultdict(int)
    handles = []
    # for each segment, draw its legend
    for segment in segments_info:
        segment_id = segment['id']
        segment_label_id = segment['label_id']
        segment_label = model.config.id2label[segment_label_id]
        label = f"{segment_label}-{instances_counter[segment_label_id]}"
        instances_counter[segment_label_id] += 1
        color = viridis(segment_id)
        handles.append(mpatches.Patch(color=color, label=label))
        
    ax.legend(handles=handles)

draw_panoptic_segmentation(**panoptic_segmentation)

在這裡，我們可以看到模型能夠檢測影像中的貓和遙控器個體。另一方面，語義分割只會為“貓”類別建立一個單一的掩碼。

要使用 OneFormer 執行推理，其 API 除了還接受額外的文字提示作為輸入外，其餘與 Mask2Former 相同，我們參考演示筆記本。

在 Transformers 中微調 Mask2Former 和 OneFormer

要將 Mask2Former/OneFormer 微調到自定義資料集上進行例項、語義和全景分割，請檢視我們的演示筆記本。MaskFormer、Mask2Former 和 OneFormer 共享相似的 API，因此從 MaskFormer 升級很容易，只需要最少的更改。

演示筆記本使用 MaskFormerForInstanceSegmentation 載入模型，而您需要切換到使用 Mask2FormerForUniversalSegmentation 或 OneFormerForUniversalSegmentation。在 Mask2Former 的影像處理情況下，您還需要切換到使用 Mask2FormerImageProcessor。您也可以使用 AutoImageProcessor 類載入影像處理器，該類會自動載入與您的模型對應的正確處理器。另一方面，OneFormer 需要 OneFormerProcessor，它準備影像以及文字輸入，以便模型使用。