pipeline_tag: 图像分割

实例分割示例

内容：

• 使用Accelerate的PyTorch版本
• 重新加载并执行推理

使用Accelerate的PyTorch版本

该模型基于脚本 run_instance_segmentation_no_trainer.py。 该脚本使用 🤗 Accelerate 在PyTorch中编写自定义训练循环，并支持在多种环境中运行，包括CPU、多CPU、GPU、多GPU和TPU，同时支持混合精度。

首先，配置环境：

accelerate config

回答有关训练环境的问题。然后运行：

accelerate test

此命令确保训练准备就绪。最后，启动训练：

accelerate launch run_instance_segmentation_no_trainer.py \
    --model_name_or_path facebook/mask2former-swin-tiny-coco-instance \
    --output_dir finetune-instance-segmentation-ade20k-mini-mask2former-no-trainer \
    --dataset_name qubvel-hf/ade20k-mini \
    --do_reduce_labels \
    --image_height 256 \
    --image_width 256 \
    --num_train_epochs 40 \
    --learning_rate 1e-5 \
    --lr_scheduler_type constant \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 2 \
    --dataloader_num_workers 8 \
    --push_to_hub

重新加载并执行推理

您可以轻松加载此训练好的模型并执行推理，如下所示：

import torch
import requests
import matplotlib.pyplot as plt

from PIL import Image
from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor

# 加载图像
image = Image.open(requests.get("http://farm4.staticflickr.com/3017/3071497290_31f0393363_z.jpg", stream=True).raw)

# 加载模型和图像处理器
device = "cuda"
checkpoint = "qubvel-hf/finetune-instance-segmentation-ade20k-mini-mask2former-no-trainer"

model = Mask2FormerForUniversalSegmentation.from_pretrained(checkpoint, device_map=device)
image_processor = Mask2FormerImageProcessor.from_pretrained(checkpoint)

# 在图像上运行推理
inputs = image_processor(images=[image], return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)

# 后处理输出
outputs = image_processor.post_process_instance_segmentation(outputs, target_sizes=[image.size[::-1]])

print("掩码形状: ", outputs[0]["segmentation"].shape)
print("掩码值: ", outputs[0]["segmentation"].unique())
for segment in outputs[0]["segments_info"]:
    print("分割段: ", segment)

掩码形状:  torch.Size([427, 640])
掩码值:  tensor([-1.,  0.,  1.,  2.,  3.,  4.,  5.,  6.])
分割段:  {'id': 0, 'label_id': 0, 'was_fused': False, 'score': 0.946127}
分割段:  {'id': 1, 'label_id': 1, 'was_fused': False, 'score': 0.961582}
分割段:  {'id': 2, 'label_id': 1, 'was_fused': False, 'score': 0.968367}
分割段:  {'id': 3, 'label_id': 1, 'was_fused': False, 'score': 0.819527}
分割段:  {'id': 4, 'label_id': 1, 'was_fused': False, 'score': 0.655761}
分割段:  {'id': 5, 'label_id': 1, 'was_fused': False, 'score': 0.531299}
分割段:  {'id': 6, 'label_id': 1, 'was_fused': False, 'score': 0.929477}

使用以下代码可视化结果：

import numpy as np
import matplotlib.pyplot as plt

segmentation = outputs[0]["segmentation"].numpy()

plt.figure(figsize=(10, 10))
plt.subplot(1, 2, 1)
plt.imshow(np.array(image))
plt.axis("off")
plt.subplot(1, 2, 2)
plt.imshow(segmentation)
plt.axis("off")
plt.show()