Introducing Google AI Edge Portal: Benchmark Edge AI at scale. Sign-up to request access during private preview.

Interactive image segmentation guide for Python

The MediaPipe Interactive Image Segmenter task takes strokes indicating a location in an image, estimates the boundaries of an object at that location, and returns the segmentation for the object as image data. These instructions show you how to use the Interactive Image Segmenter with the Python language. For more information about the capabilities, models, and configuration options of this task, see the Overview.

Code example

The example code for Interactive Image Segmenter provides a complete implementation of this task in Python for your reference. This code helps you test this task and get started on building your own interactive image segmentation application. You can view, run, edit, and download the Interactive Image Segmenter example code using just your web browser.

Setup

This section describes key steps for setting up your development environment and code projects specifically to use Interactive Image Segmenter. For general information on setting up your development environment for using MediaPipe tasks, including platform version requirements, see the Setup guide for Python. You can review the source code for this example on GitHub

Packages

The MediaPipe Interactive Image Segmenter task requires the mediapipe package. You can install the required dependencies with the following command:

$ python -m pip install mediapipe

Imports

Import the following classes to access the Interactive Image Segmenter task functions:

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

Model

The MediaPipe Interactive Image Segmenter task requires a trained model that is compatible with this task. For more information on available trained models for Interactive Image Segmenter, see the task overview Models section.

Select and download the model, and then store it within your project directory:

model_path = 'interactive_segmentation.task'

Specify the path of the model within the model_asset_path parameter, as shown in the following example:

base_options = BaseOptions(model_asset_path=model_path)

Create the task

The MediaPipe Interactive Image Segmenter task uses the create_from_options function to set up the task. The create_from_options function accepts values for configuration options to handle. For more information on configuration options, see Configuration options. The following code demonstrates how to build and configure this task.

BaseOptions = mp.tasks.BaseOptions
InteractiveSegmenter = mp.tasks.vision.InteractiveSegmenter
InteractiveSegmenterOptions = mp.tasks.vision.InteractiveSegmenterOptions

# Create a image segmenter instance:
options = InteractiveSegmenterOptions(
    base_options=BaseOptions(model_asset_path='interactive_segmentation.task'))
with InteractiveSegmenter.create_from_options(options) as segmenter:
# segmenter is initialized and ready to use

Configuration options

This task has no task-specific configuration options. It only supports the BaseOptions configuration common to all tasks, which lets you specify:

Model asset path: The location of the pre-trained task model within your project.
Delegate: The processing hardware for running the task (such as CPU or GPU).

Prepare data

Prepare your input as an image file or a numpy array, then convert it to a mediapipe.Image object.

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

For a code example showing preparation of data for Interactive Image Segmenter, see the code example.

Run the task

The Interactive Image Segmenter uses the set_image and segment functions to trigger inferences. The set_image function handles preprocessing input data and running the segmentation model. The segment function accepts strokes to postprocess raw model outputs to the segmented mask.

The following code example shows how to execute processing with the task model.

BrushMode = mp.tasks.vision.BrushMode
StrokePoint = mp.tasks.vision.StrokePoint
Stroke = mp.tasks.vision.Stroke

# Set the image to be segmented.
segmenter.set_image(mp_image)

# Perform image segmentation on the provided image using a stroke.
strokes = [Stroke(brush_mode=BrushMode.POSITIVE, points=[StrokePoint(0.5, 0.5)], is_completed=True)]
segmented_mask = segmenter.segment(strokes)

For a more complete example of running Interactive Image Segmenter inferences, see the code example.

Handle and display results

The output result for Interactive Image Segmenter is an Image data representing the confidence score of the pixel belonging to the object at the area of interest.

The following sections further explain the output data from this task:

Confidence mask

The output for the interactive segmenter is a single-channel confidence mask containing float values between [0.0, 1.0] for each pixel of the image, with dimensions matching the input image. Higher values indicate a higher confidence that the image pixel is part of the object located at the area of interest.