Segmentação de Imagens

Modelos de segmentação de imagem separam áreas correspondentes a diferentes regiões de interesse em uma imagem. Esses modelos funcionam atribuindo um rótulo a cada pixel. Existem vários tipos de segmentação: segmentação semântica, segmentação de instâncias e segmentação panótica.

Neste guia, vamos:

Dê uma olhada nos diferentes tipos de segmentação.
Disponibilizar um exemplo de ajuste fino de ponta a ponta para segmentação semântica.

Antes de começar, verifique se você tem todas as bibliotecas necessárias instaladas:

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

Recomendamos que você faça login na sua conta do Hugging Face para que possa enviar e compartilhar seu modelo com a comunidade. Quando solicitado, insira seu token para fazer login:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Tipos de segmentação

A segmentação semântica atribui um rótulo ou classe a cada pixel de uma imagem. Vamos examinar a saída de um modelo de segmentação semântica. Ele atribuirá a mesma classe a cada instância de um objeto que encontrar em uma imagem; por exemplo, todos os gatos serão rotulados como “gato”, em vez de “gato-1”, “gato-2”. Podemos usar o pipeline de segmentação de imagem do transformers para inferir rapidamente um modelo de segmentação semântica. Vamos examinar a imagem de exemplo.

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

Usaremos nvidia/segformer-b1-finetuned-cityscapes-1024-1024.

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

A saída do pipeline de segmentação inclui uma máscara para cada classe prevista.

[{'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>}]

Ao examinar a máscara da classe de carro, podemos ver que todos os carros são classificados com a mesma máscara.

results[-1]["mask"]

Na segmentação de instâncias, o objetivo não é classificar cada pixel, mas prever uma máscara para cada instância de um objeto em uma imagem dada. Funciona de forma muito semelhante à detecção de objetos, onde há uma caixa delimitadora para cada instância; aqui, há uma máscara de segmentação no lugar. Usaremos facebook/mask2former-swin-large-cityscapes-instance para isso.

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

Como você pode ver abaixo, há vários carros classificados, e não há classificação para pixels que não pertencem a instâncias de carro e pessoa.

[{'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>}]

Verificando uma das máscaras de carro abaixo.

results[2]["mask"]

A segmentação panóptica combina a segmentação semântica e a segmentação de instâncias, na qual cada pixel é classificado em uma classe e em uma instância dessa classe, e há múltiplas máscaras para cada instância de uma classe. Podemos usar facebook/mask2former-swin-large-cityscapes-panoptic para isso.

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

Como você pode ver abaixo, temos mais classes. Mais adiante, ilustraremos que cada pixel é classificado em uma das classes.

[{'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>}]

Vamos fazer uma comparação lado a lado para todos os tipos de segmentação.

Vendo todos os tipos de segmentação, vamos fazer uma análise aprofundada sobre o ajuste fino de um modelo para segmentação semântica.

Aplicações comuns do mundo real para segmentação semântica incluem treinar carros autônomos para identificar pedestres e informações importantes de tráfego, identificar células e anomalias em imagens médicas e monitorar mudanças ambientais a partir de imagens de satélite.

Ajuste fino de um modelo para segmentação

Agora, vamos:

Ajuste fino de SegFormer no conjunto de dados SceneParse150.
Use seu modelo ajustado para inferência.

Para ver todas as arquiteturas e checkpoints compatíveis com esta tarefa, recomendamos consultar a página da tarefa

Carregar o conjunto de dados SceneParse150

Comece carregando um subconjunto menor do conjunto de dados SceneParse150 a partir da biblioteca 🤗 Datasets. Isso lhe dará a oportunidade de experimentar e garantir que tudo funcione antes de gastar mais tempo treinando no conjunto de dados completo.

>>> from datasets import load_dataset

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

Divida o split train do conjunto de dados em um conjunto de treinamento e um conjunto de teste com o método train_test_split:

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

Em seguida, veja um exemplo:

>>> 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: uma imagem PIL da cena.
annotation: uma imagem PIL do mapa de segmentação, que também é o alvo do modelo.
scene_category: um id de categoria que descreve a cena da imagem, como “kitchen” ou “office”. Neste guia, você precisará apenas de image e annotation, ambos sendo imagens PIL.

Você também deve criar um dicionário que mapeia um id de rótulo para uma classe de rótulo, o que será útil ao configurar o modelo mais tarde. Baixe os mapeamentos do Hub e crie os dicionários id2label e 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)

Pré-processar

O próximo passo é carregar um processador de imagens do SegFormer para preparar as imagens e as anotações para o modelo. Alguns conjuntos de dados, como este, usam o índice zero como a classe de fundo. No entanto, a classe de fundo não está realmente incluída nas 150 classes, então você precisará definir do_reduce_labels=True para subtrair um de todos os rótulos. O índice zero é substituído por 255, para que seja ignorado pela função de perda do SegFormer:

>>> from transformers import AutoImageProcessor

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

É comum aplicar algumas ampliações de dados a um conjunto de dados de imagens para tornar um modelo mais robusto contra o overfitting. Neste guia, você usará a função ColorJitter do torchvision para alterar aleatoriamente as propriedades de cor de uma imagem, mas você também pode usar qualquer biblioteca de imagens que preferir.

>>> from torchvision.transforms import ColorJitter

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

Agora, crie duas funções de pré-processamento para preparar as imagens e as anotações para o modelo. Essas funções convertem as imagens em pixel_values e as anotações em labels. Para o conjunto de treinamento, jitter é aplicado antes de fornecer as imagens ao processador de imagens. Para o conjunto de teste, o processador de imagens recorta e normaliza as images e apenas recorta os labels, pois nenhum aumento de dados é aplicado durante o teste.

>>> 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

Para aplicar o jitter em todo o conjunto de dados, use a função set_transform do 🤗 Datasets. A transformação é aplicada em tempo de execução, o que é mais rápido e consome menos espaço em disco:

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

Avaliar

Incluir uma métrica durante o treinamento é frequentemente útil para avaliar o desempenho do seu modelo. Você pode carregar rapidamente um método de avaliação com a biblioteca 🤗 Evaluate. Para esta tarefa, carregue a métrica mean Intersection over Union (IoU) (consulte o quick tour do 🤗 Evaluate para saber mais sobre como carregar e calcular uma métrica):

>>> import evaluate

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

Em seguida, crie uma função para calcular as métricas. Suas previsões precisam primeiro ser convertidas em logits e, em seguida, redimensionadas para corresponder ao tamanho dos rótulos antes de você poder chamar compute:

>>> 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

Sua função compute_metrics está pronta e você voltará a ela ao configurar o treinamento.

Treinar

Se você não está familiarizado com o ajuste fino de um modelo usando o Trainer, confira o tutorial básico aqui!

Agora você está pronto para começar o treinamento do seu modelo! Carregue o SegFormer com AutoModelForSemanticSegmentation e passe ao modelo o mapeamento entre os ids de rótulo e as classes de rótulo:

>>> from transformers import AutoModelForSemanticSegmentation, TrainingArguments, Trainer

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

Neste ponto, restam apenas três etapas:

Defina os hiperparâmetros de treinamento em TrainingArguments. É importante não remover colunas não utilizadas, pois isso descartará a coluna image. Sem a coluna image, não é possível criar pixel_values. Defina remove_unused_columns=False para evitar esse comportamento! O único outro parâmetro obrigatório é output_dir, que especifica onde salvar seu modelo. Você enviará este modelo para o Hub definindo push_to_hub=True (é necessário estar conectado ao Hugging Face para enviar seu modelo). No final de cada epoch, o Trainer avaliará a métrica IoU e salvará o checkpoint de treinamento.
Passe os argumentos de treinamento para o Trainer junto com o modelo, o conjunto de dados, o tokenizer, o data collator e a função compute_metrics.
Chame train() para fazer o ajuste fino do seu modelo.

>>> 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,
...     report_to="trackio",
...     run_name="scene-parse-150",
...     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()

Quando o treinamento for concluído, compartilhe seu modelo no Hub com o método push_to_hub() para que todos possam usá-lo:

>>> trainer.push_to_hub()

Inferência

Ótimo, agora que você ajustou um modelo, pode usá-lo para inferência!

Recarregue o conjunto de dados e carregue uma imagem para inferência.

>>> from datasets import load_dataset

>>> ds = load_dataset("merve/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

Agora veremos como fazer inferência sem um pipeline. Processe a imagem com um processador de imagens e coloque o pixel_values em uma GPU:

>>> from accelerate import Accelerator
>>> device = Accelerator().device
>>> encoding = image_processor(image, return_tensors="pt")
>>> pixel_values = encoding.pixel_values.to(device)

Passe a sua entrada para o modelo e retorne os logits:

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

Em seguida, reescale os logits para o tamanho original da imagem:

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

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

Para visualizar os resultados, carregue a paleta de cores do conjunto de dados como ade_palette(), que mapeia cada classe para seus valores 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],
  ])

Em seguida, você pode combinar e plotar sua imagem e o mapa de segmentação previsto:

>>> 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

Update on GitHub

Transformers

Segmentação de Imagens

Tipos de segmentação

Ajuste fino de um modelo para segmentação

Carregar o conjunto de dados SceneParse150

Pré-processar

Avaliar

Treinar

Inferência