通用影像分割任務- 使用 Mask2Former 和 OneFormer

HuggingFace發表於2023-03-10
ORM

本文介紹兩個領先的影像分割神經網路模型: Mask2Former 和 OneFormer。相關模型已經在 ? Transformers 提供。? Transformers 是一個開源庫，提供了很多便捷的先進模型。在本文中，你也會學到各種影像分割任務的不同之處。

影像分割

影像分割任務旨在鑑別區分出一張圖片的不同部分，比如人物、汽車等等。從技術角度講，影像分割任務需要根據不同的語義資訊區分並聚集起對應相同語義的畫素點。讀者可以參考 Hugging Face 的 任務頁面 來簡要了解。

大體上，影像分割可以分為三個子任務: 例項分割 (instance segmentation) 、語義分割 (semantic segmentation) 、全景分割 (panoptic segmentation)。這三個子任務都有著大量的演算法與模型。

  • 例項分割 任務旨在區分不同的“例項”，例如影像中不同的人物個體。例項分割從某種角度看和物體檢測很像，不同的是在這裡我們需要的是一個對應類別的二元的分割掩膜，而不是一個檢測框。例項也可以稱為“物體 (objects)”或“實物 (things)”。需要注意的是，不同的個體可能在影像中是相互重疊的。
  • 語義分割 區分的是不同的“語義類別”，比如屬於人物、天空等類別的各個畫素點。與例項分割不同的是，這裡我們不需要區分開同一類別下的不同個體，例如這裡我們只需要得到“人物”類別的畫素級掩膜即可，不需要區分開不同的人。有些類別根本不存在個體的區分，比如天空、草地，這種類別我們稱之為“東西 (stuff)”，以此區分開其它類別，稱之為“實物 (things)”。請注意這裡不存在不同語義類別間的重疊，因為一個畫素點只能屬於一個類別。
  • 全景分割 在 2018 年由 Kirillov et al. 提出，目的是為了統一例項分割和語義分割。模型單純地鑑別出一系列的影像部分，每個部分既有對應的二元掩膜，也有對應的類別標籤。這些區分出來的部分，既可以是“東西”也可以是“實物”。與例項分割不同的是，不同部分間不存在重疊。

下圖展示了三個子任務的不同: (圖片來自 這篇部落格文章)

近年來，研究者們已經推出了很多針對例項、語義、全景分割精心設計的模型架構。例項分割和全景分割基本上是透過輸出一系列例項的二元掩膜和對應類別標籤來處理的 (和物體檢測很像，只不過這裡不是輸出每個例項的檢測框)。這一操作也常常被稱為“二元掩膜分類”。語義分割則不同，通常是讓模型輸出一個“分割圖”，令每一個畫素點都有一個標籤。所以語義分割也常被視為一個“畫素級分類”的任務。採用這一正規化的語義分割模組包括 SegFormer 和 UPerNet。針對 SegFormer 我們還寫了一篇 詳細的部落格。

通用影像分割

幸運的是，從大約 2020 年開始，人們開始研究能同時解決三個任務 (例項、語義和全景分割) 的統一模型。DETR 是開山之作，它透過“二元掩膜分類”的正規化去解決全景分割問題，把“實物”和“東西”的類別用統一的方法對待。其核心點是使用一個 Transformer 的解碼器 (decoder) 來並行地生成一系列的二元掩膜和類別。隨後 MaskFormer 又在此基礎上進行了改進，表明了“二元掩膜分類”的正規化也可以用在語義分割上。

Mask2Former 又將此方法擴充套件到了例項分割上，進一步改進了神經網路的結構。因此，各自分離的子任務框架現在已經進化到了“通用影像分割”的框架，可以解決任何影像分割任務。有趣的是，這些通用模型全都採取了“掩膜分類”的正規化，徹底拋棄了“畫素級分類”這一方法。下圖就展示了 Mask2Former 的網路結構 (影像取自 這篇論文)。

簡短來說，一張圖片首先被送入骨幹網路 (backbone) 裡面來獲取一系列，在論文中，骨幹網路既可以是 ResNet 也可以是 Swin Transformer。接下來，這些特徵圖會被一個叫做 Pixel Decoder 的模組增強成為高解析度特徵圖。最終，一個 transformer 的解碼器會接收一系列的 query，基於上一步得到的特徵，把它們轉換成一些列二元掩膜和分類預測。

需要注意的是，MasksFormer 仍然需要在每個單獨的任務上訓練來獲取領先的結果。這一點被 OneFormer 進行了改進，並透過在全景資料集上訓練，達到了領先水平。OneFormer 增加了一個文字編碼器 (text encoder)，使得模型有了一個基於文字條件 (例項、語義或全景) 的輸入。該模型已經收錄入 ? Transformers 之中，比 Mask2Former 更準確，但由於文字編碼器的引入，所以速度略慢。下圖展示了 OneFormer 的基本結構，它使用 Swin Transformer 或 DiNAT 作為骨幹網路。

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

使用 Mask2Former 和 OneFormer 方法相當直接，而且和它們的前身 MaskFormer 非常相似。我們這裡從 Hub 中使用一個在 COCO 全景資料集上訓練的一個模型來例項化一個 Mask2Former 以及對應的 processor。需要注意的是，在不同資料集上訓練出來的 checkpoints 已經公開，數量不下 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

我們使用 processor 處理原始圖片，然後送入模型進行前向推理。

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

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

模型輸出了一系列二元掩膜以及對應類別的 logit。Mask2Former 的原始輸出還可以使用 processor 進行處理，來得到最終的例項、語義或全景分割結果:

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

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

在全景分割中，最終的 prediction 包含兩樣東西: 一個是形狀為 (height, width) 的 segmentation 圖，裡面針對每一個畫素都給出了編碼例項 ID 的值; 另一個是與之對應的 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，它和 Mask2Former 的 API 幾乎一樣，只不過多了一個文字提示的輸入; 可以參考這裡的 demo notebook。

使用 transformers 微調 Mask2Former 和 OneFormer

讀者可以參考這裡的 demo notebooks 來在自定義的例項、語義或全景分割資料集上微調 Mask2Former 或 OneFormer 模型。MaskFormer、Mask2Former 和 OneFormer 都有著相似的 API，所以基於 MaskFormer 進行改進十分方便、需要的修改很少。

在上述 notebooks 中，都是使用 MaskFormerForInstanceSegmentation 來載入模型，而你需要換成使用 Mask2FormerForUniversalSegmentation 或 OneFormerForUniversalSegmentation。對於 Mask2Former 中的影像處理，你也需要使用 Mask2FormerImageProcessor。你也可以使用 AutoImageProcessor 來自動地載入適合你的模型的 processor。OneFormer 則需要使用 OneFormerProcessor，因為它不僅預處理圖片，還需要處理文字。

總結

總的來說就這些內容！你現在知道例項分割、語義分割以及全景分割都有什麼不同了，你也知道如何使用 ? Transformers 中的 Mask2Former 和 OneFormer 之類的“通用架構”了。

我們希望你喜歡本文並學有所學。如果你微調了 Mask2Former 或 OneFormer，也請讓我們知道你是否對結果足夠滿意。

如果想深入學習，我們推薦以下資源:

  • 我們針對 MaskFormer, Mask2Former and OneFormer, 推出的 demo notebooks，將會給出更多關於推理 (包括視覺化) 和微調的知識。
  • 在 Hugging Face Hub 上， Mask2Former 和 OneFormer 的 live demo spaces，可以讓你快速用自己的輸入資料嘗試不同模型。

原文連結: https://hf.co/blog/mask2former

作者: Niels Rogge、Shivalika Singh、Alara Dirik

譯者: Hoi2022

審校、排版: zhongdongy (阿東)

相關文章