Guia de segmentação de imagens para Python

A tarefa MediaPipe Image Segmenter permite dividir imagens em regiões com base em categorias predefinidas para aplicar efeitos visuais, como o desfoque do plano de fundo. Estas instruções mostram como usar o Segmentador de imagens com a linguagem Python. Para mais informações sobre os recursos, modelos e opções de configuração dessa tarefa, consulte a Visão geral.

Exemplo de código

O código de exemplo do Segmenter de imagem oferece uma implementação completa dessa tarefa em Python para referência. Este código ajuda você a testar essa tarefa e começar a criar seu próprio aplicativo de segmento de imagem. É possível visualizar, executar e editar o código de exemplo do Segmentador de imagens usando apenas seu navegador da Web.

Configuração

Esta seção descreve as principais etapas para configurar seu ambiente de desenvolvimento e projetos de código especificamente para usar o Segmentador de imagem. Para informações gerais sobre a configuração do ambiente de desenvolvimento para o uso de tarefas do MediaPipe, incluindo requisitos de versão da plataforma, consulte o Guia de configuração para Python. Revise o código-fonte deste exemplo no GitHub (em inglês).

Pacotes

A tarefa MediaPipe Image Segmenter requer o pacote mediapipe. Instale as dependências necessárias com o seguinte comando:

$ python -m pip install mediapipe

Importações

Importe as seguintes classes para acessar as funções de tarefa do Segmentador de imagens:

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

Modelo

A tarefa do Segmenter de imagem do MediaPipe requer um modelo treinado compatível com ela. Para mais informações sobre os modelos treinados disponíveis para o Image Segmenter, consulte a seção Modelos de visão geral da tarefa.

Selecione e faça o download do modelo e armazene-o no diretório do projeto:

model_path = '/absolute/path/to/model.tflite'

Especifique o caminho do modelo no parâmetro Model Name, conforme mostrado abaixo:

base_options = BaseOptions(model_asset_path=model_path)

Criar a tarefa

A tarefa do Segmenter de imagem do MediaPipe usa a função create_from_options para configurar a tarefa. A função create_from_options aceita valores para as opções de configuração processar. Para mais informações sobre a configuração da tarefa, consulte Opções de configuração.

Esses exemplos também mostram as variações de construção de tarefas para imagens, arquivos de vídeo e streams de vídeo ao vivo.

Imagem

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the image mode:
options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.IMAGE,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
  

Video

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the video mode:
options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.VIDEO,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
  

Transmissão ao vivo

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the live stream mode:
def print_result(result: List[Image], output_image: Image, timestamp_ms: int):
    print('segmented masks size: {}'.format(len(result)))

options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.VIDEO,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
    

Opções de configuração

Esta tarefa tem as seguintes opções de configuração para aplicativos Python:

Nome da opção Descrição Intervalo de valor Valor padrão
running_mode Define o modo de execução da tarefa. Há três modos:

IMAGE: o modo para entradas de imagem única.

VIDEO: o modo para frames decodificados de um vídeo.

LIVE_STREAM: é o modo para uma transmissão ao vivo de dados de entrada, como de uma câmera. Nesse modo, o resultListener precisa ser chamado para configurar um listener para receber resultados de forma assíncrona.
{IMAGE, VIDEO, LIVE_STREAM} IMAGE
output_category_mask Se definido como True, a saída incluirá uma máscara de segmentação como uma imagem uint8, em que cada valor de pixel indica o valor da categoria vencedora. {True, False} False
output_confidence_masks Se definido como True, a saída incluirá uma máscara de segmentação como uma imagem de valor flutuante, em que cada valor flutuante representa o mapa de pontuação de confiança da categoria. {True, False} True
display_names_locale Define o idioma dos rótulos a ser usado para nomes de exibição fornecidos nos metadados do modelo da tarefa, se disponível. O padrão é en para inglês. É possível adicionar rótulos localizados aos metadados de um modelo personalizado usando a API TensorFlow Lite Metadata Writer Código da localidade en
result_callback Define o listener de resultados para receber os resultados da segmentação de forma assíncrona quando o segmento de imagem está no modo de transmissão ao vivo. Só pode ser usado quando o modo de corrida está definido como LIVE_STREAM N/A N/A

preparar dados

Prepare sua entrada como um arquivo de imagem ou uma matriz numpy e a converta em um objeto mediapipe.Image. Se a entrada for um arquivo de vídeo ou uma transmissão ao vivo de uma webcam, use uma biblioteca externa, como a OpenCV (link em inglês), para carregar os frames de entrada como matrizes numpy.

Imagem

# Load the input image from an image file.
mp_image = mp.Image.create_from_file('/path/to/image')

# Load the input image from a numpy array.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_image)
    

Video

# Use OpenCV’s VideoCapture to load the input video.
# Load the frame rate of the video using OpenCV’s CV_CAP_PROP_FPS
# You need the frame rate to calculate the timestamp for each frame.
# Loop through each frame in the video using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

Transmissão ao vivo

# Use OpenCV’s VideoCapture to start capturing from the webcam.
# Create a loop to read the latest frame from the camera using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

Para conferir um exemplo de código que mostra a preparação de dados para o segmento de imagem, consulte o exemplo de código.

Executar a tarefa

O Segmentador de imagens usa as funções segment, segment_for_video e segment_async para acionar inferências. Para segmentação de imagens, isso envolve o pré-processamento de dados de entrada, a execução do modelo de segmentação e o pós-processamento das saídas do modelo bruto para as máscaras segmentadas.

Os exemplos de código a seguir mostram como executar o processamento com o modelo de tarefa.

Imagem

# Perform image segmentation on the provided single image.
# The image segmenter must be created with the image mode.
segmented_masks = segmenter.segment(mp_image)
    

Video

# Perform image segmentation on the provided single image.
# The image segmenter must be created with the video mode.
segmented_masks = segmenter.segment_for_video(mp_image, frame_timestamp_ms)
    

Transmissão ao vivo

# Send live image data to perform image segmentation.
# The results are accessible via the `result_callback` provided in
# the `ImageSegmenterOptions` object.
# The image segmenter must be created with the live stream mode.
segmenter.segment_async(mp_image, frame_timestamp_ms)
    

Observe o seguinte:

  • Ao executar no modo de vídeo ou de transmissão ao vivo, você também precisa fornecer à tarefa Segmentador de imagem o carimbo de data/hora do frame de entrada.
  • Quando executada no modelo de imagem ou de vídeo, a tarefa "Segmentação de imagem" bloqueia a linha de execução atual até que ela termine de processar a imagem ou o frame de entrada.

Para conferir um exemplo mais completo da execução de inferências do Segmentador de imagem, consulte o exemplo de código.

Gerenciar e mostrar resultados

O Segmentador de imagem gera uma lista de dados de Image. Se output_type for CATEGORY_MASK, a saída será uma lista que contém uma única máscara segmentada como uma imagem uint8. O pixel indica o índice de categoria reconhecido da imagem de entrada. Se output_type for CONFIDENCE_MASK, a saída será um vetor com o tamanho do número de categoria. Cada máscara segmentada é uma imagem flutuante dentro do intervalo [0,1], representando a pontuação de confiança do pixel pertencente à categoria.

As seções abaixo mostram exemplos dos dados de saída dessa tarefa:

Confiança da categoria

As imagens abaixo mostram uma visualização da saída da tarefa para uma máscara de confiança de categoria. A saída da máscara de confiança contém valores flutuantes entre [0, 1].

Saída da máscara de confiança da categoria e da imagem original. Imagem de origem do conjunto de dados Pascal VOC 2012.

Valor da categoria

As imagens a seguir mostram uma visualização da saída da tarefa para uma máscara de valor de categoria. O intervalo da máscara de categoria é [0, 255], e cada valor de pixel representa o índice da categoria vencedora da saída do modelo. O índice de categorias vencedor tem a pontuação mais alta entre as categorias que o modelo consegue reconhecer.

Saída original da máscara de imagem e categoria. Imagem de origem do conjunto de dados Pascal VOC 2012.