Augment Segmentation Maps

Segmentation maps are 2D arrays in which every spatial position is assigned to exactly one class. They are represented in imgaug using imgaug.augmentables.segmaps.SegmentationMapsOnImage. The class is instantiated as SegmentationMapsOnImage(arr, shape). arr contains the 2D segmentation map and shape is the shape of the corresponding image (not of the segmentation map array!). The class is called "...MapsOnImage" and not "...MapOnImage", because you can actually provide many segmentation maps for the same image (provided they all have the same height and width). That means that arr may have the shape (H,W) or (H,W,1), but it can also have shape (H,W,C) with C>1. That is useful for e.g. stacked boolean masks or instance segmentation maps (one map of instances per class).

Analogous to the constructor arguments, SegmentationMapsOnImage has the attributes .shape (shape of the corresponding image) and .arr (internal segmentation map representation).

Noteworthy methods of SegmentationMapsOnImage are:

  • get_arr(): Converts the internal representation of the segmentation map to the same shape and dtype as was originally provided to the constructor. (The internal representation is of shape (H,W,C) (with possibly C=1) and dtype int32.)
  • draw([size], [colors]): Converts the segmentation map to an RGB image.
  • draw_on_image(image, [alpha], [resize], [colors], [draw_background]): Converts the segmentation map to an RGB image and blends it with a provided image.
  • pad([top], [right], [bottom], [left], [mode], [cval]): Pad the segmentation map on its sides.
  • pad_to_aspect_ratio(aspect_ration, [mode], [cval], [return_pad_amounts]): Pad the segmentation map to an aspect ratio (ratio=width/height).
  • resize(sizes, [interpolation]): Resize the segmentation map to the provided size. Uses by default nearest neighbour interpolation.

To augment segmentation maps, use augment_segmentation_maps(), which is offered by all augmenters. It expects a single SegmentationMapsOnImage instance or a list of SegmentationMapsOnImage instances. You may also call augment(image=..., segmentation_map=...) or its alias __call__(images=..., segmentation_maps=...) (e.g. Affine(scale=1.5)(images=..., segmentation_maps=...)), which both allow to provide the segmentation map as an int-like array.

For more details, see the API: imgaug.augmentables.segmaps.SegmentationMapsOnImage, imgaug.augmenters.meta.Augmenter.augment_segmentation_maps(), imgaug.augmenters.meta.Augmenter.augment().

For drawing routines SegmentationMapsOnImage uses a predefined set of colors. These are currently saved in the constant SegmentationMapsOnImage.DEFAULT_SEGMENT_COLORS. They will likely be replaced at some point in the future by a matplotlib colormap.

Important: imgaug's segmentation map augmentation is geared towards ground truth outputs. As such, only augmentation techniques that change the image geometry will be applied to segmentation maps, even when other augmentation techniques are part of a pipeline. Examples for that are e.g. horizontal flips or affine transformations. To also apply non-geometric augmentation techniques, feed the segmentation map array through augmenter.augment_images() instead.

Creating an Example Segmentation Map from Polygons Given as Points

The goal of our first example is to load an image, create a segmentation map and augment both of them. Let's first load and visualize our example image:

In [1]:
import imageio
import imgaug as ia
%matplotlib inline

image = imageio.imread("https://upload.wikimedia.org/wikipedia/commons/f/f4/Tamias_striatus_CT.jpg")
image = ia.imresize_single_image(image, 0.15)
print(image.shape)
ia.imshow(image)
(319, 479, 3)

Now we need a segmentation map for that image. We will create two classes, one for the tree (bottom) and one for the chipmunk (center). Everything else will be background. Both classes will be created as polygons and then drawn on a segmentation map array. First, we define the four corner points of the tree polygon:

In [2]:
import numpy as np
from imgaug.augmentables.kps import KeypointsOnImage

tree_kps_xy = np.float32([
    [0, 300],  # left top of the tree
    [image.shape[1]-1, 230],  # right top
    [image.shape[1]-1, image.shape[0]-1],  # right bottom
    [0, image.shape[0]-1]  # left bottom
])

# visualize
kpsoi_tree = KeypointsOnImage.from_xy_array(tree_kps_xy, shape=image.shape)
ia.imshow(kpsoi_tree.draw_on_image(image, size=13))

Now we have to create the chipmunk polygon. That one requires significantly more corner points, but the underlying method is the same:

In [3]:
chipmunk_kps_xy = np.float32([
    [200, 50],  # left ear, top (from camera perspective)
    [220, 70],
    [260, 70],
    [280, 50],  # right ear, top
    [290, 80],
    [285, 110],
    [310, 140],
    [330, 175], # right of cheek
    [310, 260], # right of right paw
    [175, 275], # left of left paw
    [170, 220],
    [150, 200],
    [150, 170], # left of cheek
    [160, 150],
    [186, 120], # left of eye
    [185, 70]
])

# visualize
kpsoi_chipmunk = KeypointsOnImage.from_xy_array(chipmunk_kps_xy, shape=image.shape)
ia.imshow(kpsoi_chipmunk.draw_on_image(image, size=7))

In the next step, we convert both sets of corner points to instances of imgaug.augmentables.polys.Polygon:

In [4]:
from imgaug.augmentables.polys import Polygon

# create polygons
poly_tree = Polygon(kpsoi_tree.keypoints)
poly_chipmunk = Polygon(kpsoi_chipmunk.keypoints)

# visualize polygons
ia.imshow(np.hstack([
    poly_tree.draw_on_image(image),
    poly_chipmunk.draw_on_image(image)
]))