图像分割¶
图像分割模型用于将图像中对应不同兴趣区域的部分进行分离。这些模型通过为每个像素分配一个标签来实现。分割类型包括语义分割、实例分割和全景分割。
在本指南中，我们将：
- 查看不同类型的分割。
- 提供一个端到端的语义分割微调示例。
在开始之前，请确保您已安装所有必要的库：
In [ ]:
pip install -q datasets transformers evaluate accelerate
我们鼓励您登录您的 Hugging Face 账户，这样您就可以上传并与社区分享您的模型。当提示时，输入您的令牌以登录：
In [ ]:
from huggingface_hub import notebook_login
notebook_login()
分割类型¶
语义分割为图像中的每个像素分配一个标签或类别。让我们看一下语义分割模型的输出。它将为图像中遇到的对象的每个实例分配相同的类别，例如，所有猫都将被标记为“猫”，而不是“猫-1”、“猫-2”。我们可以使用 transformers 的图像分割管道快速推断语义分割模型。让我们看一下示例图像。
In [ ]:
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
In [ ]:
semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
results = semantic_segmentation(image)
results
分割管道输出包括每个预测类别的掩码。
查看汽车类别的掩码，我们可以看到每辆汽车都被分类为相同的掩码。
In [ ]:
results[-1]["mask"]
在实例分割中，目标不是对每个像素进行分类，而是预测给定图像中对象的每个实例的掩码。它的工作原理与目标检测非常相似，目标检测中每个实例都有一个边界框，而实例分割中则是一个分割掩码。我们将使用 facebook/mask2former-swin-large-cityscapes-instance 进行此操作。
In [ ]:
instance_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-instance")
results = instance_segmentation(image)
results
如您所见，这里有多辆汽车被分类，并且除了属于汽车和人实例的像素外，没有对其他像素进行分类。
查看其中一个汽车掩码如下。
In [ ]:
results[2]["mask"]
全景分割结合了语义分割和实例分割，每个像素都被分类为一个类别和该类别的实例，并且每个类别的每个实例都有多个掩码。我们可以使用 facebook/mask2former-swin-large-cityscapes-panoptic 进行此操作。
In [ ]:
panoptic_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-panoptic")
results = panoptic_segmentation(image)
results
如您所见，我们有更多的类别。稍后我们将说明每个像素都被分类为其中一个类别。
让我们对所有类型的分割进行并排比较。
看到所有类型的分割后，让我们深入了解为语义分割微调模型。
语义分割的常见现实应用包括训练自动驾驶汽车以识别行人和重要交通信息、在医学图像中识别细胞和异常情况，以及从卫星图像中监测环境变化。
为分割微调模型¶
我们现在将：
- 在 SceneParse150 数据集上微调 SegFormer。
- 使用您微调的模型进行推理。
要查看与此任务兼容的所有架构和检查点，我们建议查看 任务页面。
加载 SceneParse150 数据集¶
首先从 🤗 Datasets 库中加载 SceneParse150 数据集的一个较小的子集。这将让您有机会进行实验并确保一切正常工作，然后再在完整数据集上花费更多时间进行训练。
In [ ]:
from datasets import load_dataset
ds = load_dataset("scene_parse_150", split="train[:50]")
使用 train_test_split 方法将数据集的 train 分割成训练集和测试集：
In [ ]:
ds = ds.train_test_split(test_size=0.2)
train_ds = ds["train"]
test_ds = ds["test"]
然后查看一个示例：
In [ ]:
train_ds[0]
# 查看图像
train_ds[0]["image"]
image：场景的 PIL 图像。annotation：分割图的 PIL 图像，这也是模型的目标。scene_category：描述图像场景的类别 ID，如“厨房”或“办公室”。在本指南中，您只需要image和annotation，它们都是 PIL 图像。
您还需要创建一个字典，将标签 ID 映射到标签类，这在稍后设置模型时将非常有用。从 Hub 下载映射并创建 id2label 和 label2id 字典：
In [ ]:
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 脚本而不是笔记本实例进行训练，可以按照以下步骤操作。该脚本需要：
- 一个包含两个 Image 列的 DatasetDict，即“image”和“label”。
In [ ]:
from datasets import Dataset, DatasetDict, Image
image_paths_train = ["path/to/image_1.jpg", "path/to/image_2.jpg", ..., "path/to/image_n.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
# 步骤 1: 创建 Dataset 对象
train_dataset = create_dataset(image_paths_train, label_paths_train)
validation_dataset = create_dataset(image_paths_validation, label_paths_validation)
# 步骤 2: 创建 DatasetDict
dataset = DatasetDict({
"train": train_dataset,
"validation": validation_dataset,
}
)
# 步骤 3: 推送到 Hub（假设您已在终端/笔记本中运行了 huggingface-cli login 命令）
dataset.push_to_hub("your-name/dataset-repo")
# 可选地，您可以将模型推送到 Hub 上的私有仓库
# dataset.push_to_hub("name of repo on the hub", private=True)
- 一个将类整数映射到类名的 id2label 字典。
In [ ]:
import json
# 简单示例
id2label = {0: 'cat', 1: 'dog'}
with open('id2label.json', 'w') as fp:
json.dump(id2label, fp)
In [ ]:
from transformers import AutoImageProcessor
checkpoint = "nvidia/mit-b0"
image_processor = AutoImageProcessor.from_pretrained(checkpoint, do_reduce_labels=True)
PyTorch
通常对图像数据集应用一些数据增强，以使模型更能抵抗过拟合。在本指南中，您将使用 torchvision 的 ColorJitter 函数随机更改图像的颜色属性，但您也可以使用您喜欢的任何图像库。
In [ ]:
from torchvision.transforms import ColorJitter
jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
现在创建两个预处理函数来准备模型所需的图像和注释。这些函数将图像转换为 pixel_values，并将注释转换为 labels。对于训练集，jitter 在提供图像给图像处理器之前应用。对于测试集，图像处理器只裁剪和规范化 images，并且只裁剪 labels，因为在测试期间不应用数据增强。
In [ ]:
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 函数。转换是即时应用的，这更快，并且消耗更少的磁盘空间：
In [ ]:
train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)
In [ ]:
import evaluate
metric = evaluate.load("mean_iou")
In [ ]:
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
训练¶
PyTorch
如果您不熟悉使用 Trainer 微调模型，请查看 这里 的基本教程！
您现在可以开始训练您的模型了！使用 AutoModelForSemanticSegmentation 加载 SegFormer，并将模型传递给标签 ID 和标签类之间的映射：
In [ ]:
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 指标并保存训练检查点。 - 将训练参数传递给 Trainer，以及模型、数据集、tokenizer、数据整理器和
compute_metrics函数。 - 调用 train() 来微调您的模型。
In [ ]:
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，以便每个人都可以使用您的模型：
In [ ]:
trainer.push_to_hub()
In [ ]:
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
我们现在将看到如何在没有管道的情况下进行推理。使用图像处理器处理图像，并将 pixel_values 放在 GPU 上：
In [ ]:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 如果可用，使用 GPU，否则使用 CPU
encoding = image_processor(image, return_tensors="pt")
pixel_values = encoding.pixel_values.to(device)
将您的输入传递给模型并返回 logits：
In [ ]:
outputs = model(pixel_values=pixel_values)
logits = outputs.logits.cpu()
接下来，将 logits 缩放到原始图像大小：
In [ ]:
upsampled_logits = nn.functional.interpolate(
logits,
size=image.size[::-1],
mode="bilinear",
align_corners=False,
)
pred_seg = upsampled_logits.argmax(dim=1)[0]
为了可视化结果，加载 数据集颜色调色板 作为 ade_palette()，它将每个类映射到它们的 RGB 值。
In [ ]:
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],
])
然后，您可以组合并绘制您的图像和预测的分割图：
In [ ]:
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] # 转换为 BGR
img = np.array(image) * 0.5 + color_seg * 0.5 # 绘制带有分割图的图像
img = img.astype(np.uint8)
plt.figure(figsize=(15, 10))
plt.imshow(img)
plt.show()
恭喜！您已经成功地微调了一个图像分割模型，并使用它进行了推理。您现在可以将这个模型用于各种图像分割任务，如自动驾驶、医学图像分析和环境监测。