图像分割

图像分割模型将图像中对应于不同感兴趣区域的区域分隔开来。这些模型通过为每个像素分配标签来工作。分割有几种类型：语义分割、实例分割和全景分割。

在本指南中，我们将

了解不同类型的分割.
提供一个用于语义分割的端到端微调示例.

在开始之前，请确保您已安装所有必要的库

# uncomment to install the necessary libraries
!pip install -q datasets transformers evaluate accelerate

我们鼓励您登录您的 Hugging Face 帐户，以便您可以上传您的模型并与社区分享。出现提示时，输入您的令牌以登录

>>> from huggingface_hub import notebook_login

>>> notebook_login()

分割的类型

语义分割为图像中的每个像素分配一个标签或类别。让我们看一下语义分割模型的输出。它将为图像中遇到的每个对象实例分配相同的类别，例如，所有猫都将被标记为“猫”而不是“猫-1”、“猫-2”。我们可以使用 transformers 的图像分割 pipeline 快速推断语义分割模型。让我们看一下示例图像。

from transformers import pipeline
from PIL import Image
import requests

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/segmentation_input.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image

我们将使用 nvidia/segformer-b1-finetuned-cityscapes-1024-1024。

semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
results = semantic_segmentation(image)
results

分割 pipeline 的输出包括每个预测类别的掩码。

[{'score': None,
  'label': 'road',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'sidewalk',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'building',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'wall',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'pole',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'traffic sign',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'vegetation',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'terrain',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'sky',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': None,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>}]

查看汽车类别的掩码，我们可以看到每辆汽车都使用相同的掩码进行分类。

results[-1]["mask"]

在实例分割中，目标不是对每个像素进行分类，而是为给定图像中对象的每个实例预测一个掩码。它的工作原理与目标检测非常相似，目标检测为每个实例都有一个边界框，而实例分割则使用分割掩码代替。我们将使用 facebook/mask2former-swin-large-cityscapes-instance 来实现这一点。

instance_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-instance")
results = instance_segmentation(image)
results

如下所示，有多辆汽车被分类，并且除了属于汽车和人实例的像素之外，没有对其他像素进行分类。

[{'score': 0.999944,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999945,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999652,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.903529,
  'label': 'person',
  'mask': <PIL.Image.Image image mode=L size=612x415>}]

查看下面其中一个汽车掩码。

results[2]["mask"]

全景分割结合了语义分割和实例分割，其中每个像素都被分类为一个类别和该类别的实例，并且每个类别的实例都有多个掩码。我们可以使用 facebook/mask2former-swin-large-cityscapes-panoptic 来实现这一点。

panoptic_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-panoptic")
results = panoptic_segmentation(image)
results

如下所示，我们有更多的类别。稍后我们将说明每个像素都被分类为其中一个类别。

[{'score': 0.999981,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999958,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.99997,
  'label': 'vegetation',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999575,
  'label': 'pole',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999958,
  'label': 'building',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999634,
  'label': 'road',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.996092,
  'label': 'sidewalk',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.999221,
  'label': 'car',
  'mask': <PIL.Image.Image image mode=L size=612x415>},
 {'score': 0.99987,
  'label': 'sky',
  'mask': <PIL.Image.Image image mode=L size=612x415>}]

让我们对所有类型的分割进行并排比较。

了解了所有类型的分割之后，让我们深入探讨如何微调语义分割模型。

语义分割常见的现实世界应用包括训练自动驾驶汽车以识别行人和重要的交通信息、识别医学图像中的细胞和异常情况，以及通过卫星图像监测环境变化。

微调分割模型

现在我们将

在 SceneParse150 数据集上微调 SegFormer。
使用您微调的模型进行推理。

要查看与此任务兼容的所有架构和 checkpoint，我们建议查看 task-page

加载 SceneParse150 数据集

首先从 🤗 Datasets 库加载 SceneParse150 数据集的一个较小子集。这将让您有机会进行实验，并确保一切正常，然后再花费更多时间在完整数据集上进行训练。

>>> from datasets import load_dataset

>>> ds = load_dataset("scene_parse_150", split="train[:50]")

使用 train_test_split 方法将数据集的 train 拆分拆分为训练集和测试集

>>> ds = ds.train_test_split(test_size=0.2)
>>> train_ds = ds["train"]
>>> test_ds = ds["test"]

然后查看一个示例

>>> train_ds[0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x683 at 0x7F9B0C201F90>,
 'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=512x683 at 0x7F9B0C201DD0>,
 'scene_category': 368}

# view the image
>>> train_ds[0]["image"]

image：场景的 PIL 图像。
annotation：分割图的 PIL 图像，这也是模型的target。
scene_category：描述图像场景的类别 ID，例如“厨房”或“办公室”。在本指南中，您只需要 image 和 annotation，它们都是 PIL 图像。

您还需要创建一个字典，将标签 ID 映射到标签类别，这在您稍后设置模型时会很有用。从 Hub 下载映射并创建 id2label 和 label2id 字典

>>> import json
>>> from pathlib import Path
>>> from huggingface_hub import hf_hub_download

>>> repo_id = "huggingface/label-files"
>>> filename = "ade20k-id2label.json"
>>> id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text())
>>> id2label = {int(k): v for k, v in id2label.items()}
>>> label2id = {v: k for k, v in id2label.items()}
>>> num_labels = len(id2label)

自定义数据集

如果您更喜欢使用 run_semantic_segmentation.py 脚本而不是 notebook 实例进行训练，您也可以创建并使用您自己的数据集。该脚本需要

一个 DatasetDict，其中包含两个 Image 列，“image”和“label”

from datasets import Dataset, DatasetDict, Image

image_paths_train = ["path/to/image_1.jpg/jpg", "path/to/image_2.jpg/jpg", ..., "path/to/image_n.jpg/jpg"]
label_paths_train = ["path/to/annotation_1.png", "path/to/annotation_2.png", ..., "path/to/annotation_n.png"]

image_paths_validation = [...]
label_paths_validation = [...]

def create_dataset(image_paths, label_paths):
    dataset = Dataset.from_dict({"image": sorted(image_paths),
                                "label": sorted(label_paths)})
    dataset = dataset.cast_column("image", Image())
    dataset = dataset.cast_column("label", Image())
    return dataset

# step 1: create Dataset objects
train_dataset = create_dataset(image_paths_train, label_paths_train)
validation_dataset = create_dataset(image_paths_validation, label_paths_validation)

# step 2: create DatasetDict
dataset = DatasetDict({
     "train": train_dataset,
     "validation": validation_dataset,
     }
)

# step 3: push to Hub (assumes you have ran the huggingface-cli login command in a terminal/notebook)
dataset.push_to_hub("your-name/dataset-repo")

# optionally, you can push to a private repo on the Hub
# dataset.push_to_hub("name of repo on the hub", private=True)

一个 id2label 字典，将类整数映射到它们的类名

import json
# simple example
id2label = {0: 'cat', 1: 'dog'}
with open('id2label.json', 'w') as fp:
json.dump(id2label, fp)

例如，查看这个示例数据集，它是使用上面显示的步骤创建的。

预处理

下一步是加载 SegFormer 图像处理器，以准备模型所需的图像和 annotation。某些数据集（例如此数据集）使用零索引作为背景类。但是，背景类实际上不包含在 150 个类别中，因此您需要设置 do_reduce_labels=True 以从所有标签中减去 1。零索引被替换为 255，因此 SegFormer 的损失函数会忽略它

>>> from transformers import AutoImageProcessor

>>> checkpoint = "nvidia/mit-b0"
>>> image_processor = AutoImageProcessor.from_pretrained(checkpoint, do_reduce_labels=True)

Pytorch

隐藏 Pytorch 内容

通常会对图像数据集应用一些数据增强，以使模型更健壮，防止过拟合。在本指南中，您将使用 torchvision 中的 ColorJitter 函数来随机更改图像的颜色属性，但您也可以使用任何您喜欢的图像库。

>>> from torchvision.transforms import ColorJitter

>>> jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)

现在创建两个预处理函数，以准备模型所需的图像和 annotation。这些函数将图像转换为 pixel_values，并将 annotation 转换为 labels。对于训练集，在将图像提供给图像处理器之前应用 jitter。对于测试集，图像处理器裁剪并标准化 images，并且仅裁剪 labels，因为在测试期间不应用数据增强。

>>> def train_transforms(example_batch):
...     images = [jitter(x) for x in example_batch["image"]]
...     labels = [x for x in example_batch["annotation"]]
...     inputs = image_processor(images, labels)
...     return inputs


>>> def val_transforms(example_batch):
...     images = [x for x in example_batch["image"]]
...     labels = [x for x in example_batch["annotation"]]
...     inputs = image_processor(images, labels)
...     return inputs

要将 jitter 应用于整个数据集，请使用 🤗 Datasets set_transform 函数。该转换是动态应用的，速度更快，并且占用更少的磁盘空间

>>> train_ds.set_transform(train_transforms)
>>> test_ds.set_transform(val_transforms)

TensorFlow

隐藏 TensorFlow 内容

通常会对图像数据集应用一些数据增强，以使模型更健壮，防止过拟合。在本指南中，您将使用 tf.image 来随机更改图像的颜色属性，但您也可以使用任何您喜欢的图像库。定义两个单独的转换函数

包括图像增强的训练数据转换
仅转置图像的验证数据转换，因为 🤗 Transformers 中的计算机视觉模型期望通道优先布局

>>> import tensorflow as tf


>>> def aug_transforms(image):
...     image = tf.keras.utils.img_to_array(image)
...     image = tf.image.random_brightness(image, 0.25)
...     image = tf.image.random_contrast(image, 0.5, 2.0)
...     image = tf.image.random_saturation(image, 0.75, 1.25)
...     image = tf.image.random_hue(image, 0.1)
...     image = tf.transpose(image, (2, 0, 1))
...     return image


>>> def transforms(image):
...     image = tf.keras.utils.img_to_array(image)
...     image = tf.transpose(image, (2, 0, 1))
...     return image

接下来，创建两个预处理函数，以准备模型所需的批量图像和 annotation。这些函数应用图像转换，并使用之前加载的 image_processor 将图像转换为 pixel_values，并将 annotation 转换为 labels。ImageProcessor 还负责调整大小和标准化图像。

>>> def train_transforms(example_batch):
...     images = [aug_transforms(x.convert("RGB")) for x in example_batch["image"]]
...     labels = [x for x in example_batch["annotation"]]
...     inputs = image_processor(images, labels)
...     return inputs


>>> def val_transforms(example_batch):
...     images = [transforms(x.convert("RGB")) for x in example_batch["image"]]
...     labels = [x for x in example_batch["annotation"]]
...     inputs = image_processor(images, labels)
...     return inputs

要将预处理转换应用于整个数据集，请使用 🤗 Datasets set_transform 函数。该转换是动态应用的，速度更快，并且占用更少的磁盘空间

>>> train_ds.set_transform(train_transforms)
>>> test_ds.set_transform(val_transforms)

评估

在训练期间包含指标通常有助于评估模型的性能。您可以使用 🤗 Evaluate 库快速加载评估方法。对于此任务，加载平均交并比 (IoU) 指标（请参阅 🤗 Evaluate 快速入门，以了解有关如何加载和计算指标的更多信息）

>>> import evaluate

>>> metric = evaluate.load("mean_iou")

然后创建一个函数来 compute 指标。您的预测需要首先转换为 logits，然后重塑以匹配 labels 的大小，然后才能调用 compute

Pytorch

隐藏 Pytorch 内容

>>> import numpy as np
>>> import torch
>>> from torch import nn

>>> def compute_metrics(eval_pred):
...     with torch.no_grad():
...         logits, labels = eval_pred
...         logits_tensor = torch.from_numpy(logits)
...         logits_tensor = nn.functional.interpolate(
...             logits_tensor,
...             size=labels.shape[-2:],
...             mode="bilinear",
...             align_corners=False,
...         ).argmax(dim=1)

...         pred_labels = logits_tensor.detach().cpu().numpy()
...         metrics = metric.compute(
...             predictions=pred_labels,
...             references=labels,
...             num_labels=num_labels,
...             ignore_index=255,
...             reduce_labels=False,
...         )
...         for key, value in metrics.items():
...             if isinstance(value, np.ndarray):
...                 metrics[key] = value.tolist()
...         return metrics

TensorFlow

隐藏 TensorFlow 内容

>>> def compute_metrics(eval_pred):
...     logits, labels = eval_pred
...     logits = tf.transpose(logits, perm=[0, 2, 3, 1])
...     logits_resized = tf.image.resize(
...         logits,
...         size=tf.shape(labels)[1:],
...         method="bilinear",
...     )

...     pred_labels = tf.argmax(logits_resized, axis=-1)
...     metrics = metric.compute(
...         predictions=pred_labels,
...         references=labels,
...         num_labels=num_labels,
...         ignore_index=-1,
...         reduce_labels=image_processor.do_reduce_labels,
...     )

...     per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
...     per_category_iou = metrics.pop("per_category_iou").tolist()

...     metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
...     metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})
...     return {"val_" + k: v for k, v in metrics.items()}

您的 compute_metrics 函数现在已准备就绪，当您设置训练时，您将返回到它。

训练

Pytorch

隐藏 Pytorch 内容

如果您不熟悉使用 Trainer 微调模型，请查看此处的入门教程 here！

您现在可以开始训练您的模型了！使用 AutoModelForSemanticSegmentation 加载 SegFormer，并将标签 ID 和标签类别之间的映射传递给模型

>>> from transformers import AutoModelForSemanticSegmentation, TrainingArguments, Trainer

>>> model = AutoModelForSemanticSegmentation.from_pretrained(checkpoint, id2label=id2label, label2id=label2id)

此时，仅剩三个步骤

在 TrainingArguments 中定义您的训练超参数。重要的是您不要删除未使用的列，因为这会删除 image 列。如果没有 image 列，您将无法创建 pixel_values。设置 remove_unused_columns=False 以防止此行为！唯一其他必需的参数是 output_dir，它指定保存模型的位置。您将通过设置 push_to_hub=True 将此模型推送到 Hub（您需要登录 Hugging Face 才能上传您的模型）。在每个 epoch 结束时，Trainer 将评估 IoU 指标并保存训练 checkpoint。
将训练参数以及模型、数据集、tokenizer、data collator 和 compute_metrics 函数传递给 Trainer。
调用 train() 以微调您的模型。

>>> training_args = TrainingArguments(
...     output_dir="segformer-b0-scene-parse-150",
...     learning_rate=6e-5,
...     num_train_epochs=50,
...     per_device_train_batch_size=2,
...     per_device_eval_batch_size=2,
...     save_total_limit=3,
...     eval_strategy="steps",
...     save_strategy="steps",
...     save_steps=20,
...     eval_steps=20,
...     logging_steps=1,
...     eval_accumulation_steps=5,
...     remove_unused_columns=False,
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=train_ds,
...     eval_dataset=test_ds,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

训练完成后，使用 push_to_hub() 方法将您的模型分享到 Hub，以便所有人都可以使用您的模型

>>> trainer.push_to_hub()

TensorFlow

隐藏 TensorFlow 内容

如果您不熟悉使用 Keras 微调模型，请先查看入门教程！

要在 TensorFlow 中微调模型，请按照以下步骤操作

定义训练超参数，并设置优化器和学习率 schedule。
实例化一个预训练模型。
将 🤗 Dataset 转换为 tf.data.Dataset。
编译您的模型。
添加 callbacks 以计算指标并将您的模型上传到 🤗 Hub
使用 fit() 方法运行训练。

首先定义超参数、优化器和学习率 schedule

>>> from transformers import create_optimizer

>>> batch_size = 2
>>> num_epochs = 50
>>> num_train_steps = len(train_ds) * num_epochs
>>> learning_rate = 6e-5
>>> weight_decay_rate = 0.01

>>> optimizer, lr_schedule = create_optimizer(
...     init_lr=learning_rate,
...     num_train_steps=num_train_steps,
...     weight_decay_rate=weight_decay_rate,
...     num_warmup_steps=0,
... )

然后，使用 TFAutoModelForSemanticSegmentation 加载 SegFormer 以及标签映射，并使用优化器对其进行编译。请注意，Transformers 模型都具有默认的任务相关损失函数，因此除非您想指定一个损失函数，否则无需指定

>>> from transformers import TFAutoModelForSemanticSegmentation

>>> model = TFAutoModelForSemanticSegmentation.from_pretrained(
...     checkpoint,
...     id2label=id2label,
...     label2id=label2id,
... )
>>> model.compile(optimizer=optimizer)  # No loss argument!

使用 to_tf_dataset 和 DefaultDataCollator 将您的数据集转换为 tf.data.Dataset 格式

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator(return_tensors="tf")

>>> tf_train_dataset = train_ds.to_tf_dataset(
...     columns=["pixel_values", "label"],
...     shuffle=True,
...     batch_size=batch_size,
...     collate_fn=data_collator,
... )

>>> tf_eval_dataset = test_ds.to_tf_dataset(
...     columns=["pixel_values", "label"],
...     shuffle=True,
...     batch_size=batch_size,
...     collate_fn=data_collator,
... )

要从预测中计算准确率并将您的模型推送到 🤗 Hub，请使用 Keras callbacks。将您的 compute_metrics 函数传递给 KerasMetricCallback，并使用 PushToHubCallback 上传模型

>>> from transformers.keras_callbacks import KerasMetricCallback, PushToHubCallback

>>> metric_callback = KerasMetricCallback(
...     metric_fn=compute_metrics, eval_dataset=tf_eval_dataset, batch_size=batch_size, label_cols=["labels"]
... )

>>> push_to_hub_callback = PushToHubCallback(output_dir="scene_segmentation", tokenizer=image_processor)

>>> callbacks = [metric_callback, push_to_hub_callback]

最后，您已准备好训练您的模型！使用您的训练和验证数据集、epoch 数量和您的 callbacks 调用 fit() 以微调模型

>>> model.fit(
...     tf_train_dataset,
...     validation_data=tf_eval_dataset,
...     callbacks=callbacks,
...     epochs=num_epochs,
... )

恭喜！您已经微调了您的模型并将其分享到 🤗 Hub 上。您现在可以使用它进行推理了！

推理

太棒了，现在您已经微调了一个模型，您可以使用它进行推理了！

重新加载数据集并加载用于推理的图像。

>>> from datasets import load_dataset

>>> ds = load_dataset("scene_parse_150", split="train[:50]")
>>> ds = ds.train_test_split(test_size=0.2)
>>> test_ds = ds["test"]
>>> image = ds["test"][0]["image"]
>>> image

Pytorch

隐藏 Pytorch 内容

我们现在将看到如何在没有 pipeline 的情况下进行推理。使用图像处理器处理图像并将 pixel_values 放在 GPU 上

>>> from accelerate.test_utils.testing import get_backend
# automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
>>> device, _, _ = get_backend()
>>> encoding = image_processor(image, return_tensors="pt")
>>> pixel_values = encoding.pixel_values.to(device)

将您的输入传递给模型并返回 logits

>>> outputs = model(pixel_values=pixel_values)
>>> logits = outputs.logits.cpu()

接下来，将 logits 重新缩放到原始图像大小

>>> upsampled_logits = nn.functional.interpolate(
...     logits,
...     size=image.size[::-1],
...     mode="bilinear",
...     align_corners=False,
... )

>>> pred_seg = upsampled_logits.argmax(dim=1)[0]

TensorFlow

隐藏 TensorFlow 内容

加载图像处理器以预处理图像并将输入作为 TensorFlow tensors 返回

>>> from transformers import AutoImageProcessor

>>> image_processor = AutoImageProcessor.from_pretrained("MariaK/scene_segmentation")
>>> inputs = image_processor(image, return_tensors="tf")

将您的输入传递给模型并返回 logits

>>> from transformers import TFAutoModelForSemanticSegmentation

>>> model = TFAutoModelForSemanticSegmentation.from_pretrained("MariaK/scene_segmentation")
>>> logits = model(**inputs).logits

接下来，将 logits 重新缩放到原始图像大小，并在类别维度上应用 argmax

>>> logits = tf.transpose(logits, [0, 2, 3, 1])

>>> upsampled_logits = tf.image.resize(
...     logits,
...     # We reverse the shape of `image` because `image.size` returns width and height.
...     image.size[::-1],
... )

>>> pred_seg = tf.math.argmax(upsampled_logits, axis=-1)[0]

要可视化结果，请加载数据集调色板作为 ade_palette()，它将每个类别映射到其 RGB 值。

def ade_palette():
  return np.asarray([
      [0, 0, 0],
      [120, 120, 120],
      [180, 120, 120],
      [6, 230, 230],
      [80, 50, 50],
      [4, 200, 3],
      [120, 120, 80],
      [140, 140, 140],
      [204, 5, 255],
      [230, 230, 230],
      [4, 250, 7],
      [224, 5, 255],
      [235, 255, 7],
      [150, 5, 61],
      [120, 120, 70],
      [8, 255, 51],
      [255, 6, 82],
      [143, 255, 140],
      [204, 255, 4],
      [255, 51, 7],
      [204, 70, 3],
      [0, 102, 200],
      [61, 230, 250],
      [255, 6, 51],
      [11, 102, 255],
      [255, 7, 71],
      [255, 9, 224],
      [9, 7, 230],
      [220, 220, 220],
      [255, 9, 92],
      [112, 9, 255],
      [8, 255, 214],
      [7, 255, 224],
      [255, 184, 6],
      [10, 255, 71],
      [255, 41, 10],
      [7, 255, 255],
      [224, 255, 8],
      [102, 8, 255],
      [255, 61, 6],
      [255, 194, 7],
      [255, 122, 8],
      [0, 255, 20],
      [255, 8, 41],
      [255, 5, 153],
      [6, 51, 255],
      [235, 12, 255],
      [160, 150, 20],
      [0, 163, 255],
      [140, 140, 140],
      [250, 10, 15],
      [20, 255, 0],
      [31, 255, 0],
      [255, 31, 0],
      [255, 224, 0],
      [153, 255, 0],
      [0, 0, 255],
      [255, 71, 0],
      [0, 235, 255],
      [0, 173, 255],
      [31, 0, 255],
      [11, 200, 200],
      [255, 82, 0],
      [0, 255, 245],
      [0, 61, 255],
      [0, 255, 112],
      [0, 255, 133],
      [255, 0, 0],
      [255, 163, 0],
      [255, 102, 0],
      [194, 255, 0],
      [0, 143, 255],
      [51, 255, 0],
      [0, 82, 255],
      [0, 255, 41],
      [0, 255, 173],
      [10, 0, 255],
      [173, 255, 0],
      [0, 255, 153],
      [255, 92, 0],
      [255, 0, 255],
      [255, 0, 245],
      [255, 0, 102],
      [255, 173, 0],
      [255, 0, 20],
      [255, 184, 184],
      [0, 31, 255],
      [0, 255, 61],
      [0, 71, 255],
      [255, 0, 204],
      [0, 255, 194],
      [0, 255, 82],
      [0, 10, 255],
      [0, 112, 255],
      [51, 0, 255],
      [0, 194, 255],
      [0, 122, 255],
      [0, 255, 163],
      [255, 153, 0],
      [0, 255, 10],
      [255, 112, 0],
      [143, 255, 0],
      [82, 0, 255],
      [163, 255, 0],
      [255, 235, 0],
      [8, 184, 170],
      [133, 0, 255],
      [0, 255, 92],
      [184, 0, 255],
      [255, 0, 31],
      [0, 184, 255],
      [0, 214, 255],
      [255, 0, 112],
      [92, 255, 0],
      [0, 224, 255],
      [112, 224, 255],
      [70, 184, 160],
      [163, 0, 255],
      [153, 0, 255],
      [71, 255, 0],
      [255, 0, 163],
      [255, 204, 0],
      [255, 0, 143],
      [0, 255, 235],
      [133, 255, 0],
      [255, 0, 235],
      [245, 0, 255],
      [255, 0, 122],
      [255, 245, 0],
      [10, 190, 212],
      [214, 255, 0],
      [0, 204, 255],
      [20, 0, 255],
      [255, 255, 0],
      [0, 153, 255],
      [0, 41, 255],
      [0, 255, 204],
      [41, 0, 255],
      [41, 255, 0],
      [173, 0, 255],
      [0, 245, 255],
      [71, 0, 255],
      [122, 0, 255],
      [0, 255, 184],
      [0, 92, 255],
      [184, 255, 0],
      [0, 133, 255],
      [255, 214, 0],
      [25, 194, 194],
      [102, 255, 0],
      [92, 0, 255],
  ])

然后，您可以组合并绘制您的图像和预测的分割图

>>> import matplotlib.pyplot as plt
>>> import numpy as np

>>> color_seg = np.zeros((pred_seg.shape[0], pred_seg.shape[1], 3), dtype=np.uint8)
>>> palette = np.array(ade_palette())
>>> for label, color in enumerate(palette):
...     color_seg[pred_seg == label, :] = color
>>> color_seg = color_seg[..., ::-1]  # convert to BGR

>>> img = np.array(image) * 0.5 + color_seg * 0.5  # plot the image with the segmentation map
>>> img = img.astype(np.uint8)

>>> plt.figure(figsize=(15, 10))
>>> plt.imshow(img)
>>> plt.show()

Image of bedroom overlaid with segmentation map

< > 在 GitHub 上更新