Руководство по сегментации изображений для Python

Задача MediaPipe Image Segmenter позволяет разделять изображения на области на основе предопределенных категорий для применения визуальных эффектов, таких как размытие фона. Эти инструкции покажут вам, как использовать сегментатор изображений с языком Python. Дополнительные сведения о возможностях, моделях и параметрах конфигурации этой задачи см. в разделе Обзор .

Пример кода

В примере кода Image Segmenter представлена ​​полная реализация этой задачи на Python. Этот код поможет вам протестировать эту задачу и приступить к созданию собственного приложения для сегментирования изображений. Вы можете просматривать, запускать и редактировать пример кода Image Segmenter, используя только веб-браузер.

Настраивать

В этом разделе описаны ключевые шаги по настройке среды разработки и проектов кода специально для использования Image Segmenter. Общие сведения о настройке среды разработки для использования задач MediaPipe, включая требования к версии платформы, см. в руководстве по настройке Python . Вы можете просмотреть исходный код этого примера на GitHub.

Пакеты

Для задачи сегментации изображений MediaPipe требуется пакет mediapipe . Вы можете установить необходимые зависимости с помощью следующей команды:

$ python -m pip install mediapipe

Импорт

Импортируйте следующие классы для доступа к функциям задачи сегментатора изображений:

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

Модель

Для задачи сегментации изображений MediaPipe требуется обученная модель, совместимая с этой задачей. Дополнительную информацию о доступных обученных моделях для Image Segmenter см. в разделе «Модели » обзора задач.

Выберите и загрузите модель, а затем сохраните ее в каталоге вашего проекта:

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

Укажите путь к модели в параметре «Имя модели», как показано ниже:

base_options = BaseOptions(model_asset_path=model_path)

Создать задачу

Задача сегментации изображений MediaPipe использует функцию create_from_options для настройки задачи. Функция create_from_options принимает значения для обрабатываемых параметров конфигурации. Дополнительные сведения о настройке задачи см. в разделе Параметры конфигурации .

В этих примерах также показаны варианты построения задач для изображений, видеофайлов и прямых видеопотоков.

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:
  

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:
  

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:
    

Варианты конфигурации

Эта задача имеет следующие параметры конфигурации для приложений Python:

Подготовьте данные

Подготовьте входные данные в виде файла изображения или массива numpy, а затем преобразуйте их в объект mediapipe.Image . Если ваши входные данные представляют собой видеофайл или прямую трансляцию с веб-камеры, вы можете использовать внешнюю библиотеку, такую ​​как OpenCV, для загрузки входных кадров в виде числовых массивов.

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

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

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

Пример кода, показывающий подготовку данных для Image Segmenter, см. в примере кода .

Запустить задачу

Сегментатор изображений использует функции segment , segment_for_video и segment_async для инициирования выводов. Для сегментации изображения это включает предварительную обработку входных данных, запуск модели сегментации и постобработку необработанных выходных данных модели в сегментированные маски.

В следующих примерах кода показано, как выполнить обработку с помощью модели задачи.

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

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

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

Обратите внимание на следующее:

• При работе в режиме видео или режиме прямой трансляции вы также должны предоставить задаче «Сегментатор изображений» временную метку входного кадра.
• При работе с изображением или видеомоделью задача «Сегментатор изображений» блокирует текущий поток до тех пор, пока он не завершит обработку входного изображения или кадра.

Более полный пример выполнения выводов Image Segmenter смотрите в примере кода .

Обработка и отображение результатов

Сегментатор изображений выводит список данных Image . Если output_type — CATEGORY_MASK , вывод представляет собой список, содержащий одну сегментированную маску в виде изображения uint8. Пиксель указывает распознанный индекс категории входного изображения. Если output_type равен CONFIDENCE_MASK , вывод представляет собой вектор с размером номера категории. Каждая сегментированная маска представляет собой плавающее изображение в диапазоне [0,1] , представляющее показатель достоверности пикселя, принадлежащего категории.

В следующих разделах показаны примеры выходных данных этой задачи:

Категория доверия

На следующих изображениях показана визуализация результатов задачи для маски достоверности категории. Выходные данные маски доверия содержат значения с плавающей запятой между [0, 1] .

Вывод исходного изображения и маски достоверности категории. Исходное изображение из набора данных Pascal VOC 2012 .

Значение категории

На следующих изображениях показана визуализация выходных данных задачи для маски значения категории. Диапазон маски категории — [0, 255] , и каждое значение пикселя представляет индекс выигрышной категории выходных данных модели. Индекс выигрышной категории имеет наивысший балл среди категорий, которые может распознать модель.

Вывод исходного изображения и маски категории. Исходное изображение из набора данных Pascal VOC 2012 .