aleju--imgaug
wehub-resource-sync
570 行
22 KiB
Python
570 行
22 KiB
Python
"""Classes dealing with segmentation maps.
|
|
|
|
E.g. masks, semantic or instance segmentation maps.
|
|
|
|
"""
|
|
from __future__ import print_function, division, absolute_import
|
|
|
|
import numpy as np
|
|
import six.moves as sm
|
|
|
|
from .. import imgaug as ia
|
|
from ..augmenters import blend as blendlib
|
|
from .base import IAugmentable
|
|
|
|
|
|
@ia.deprecated(alt_func="SegmentationMapsOnImage",
|
|
comment="(Note the plural 'Maps' instead of old 'Map'.)")
|
|
def SegmentationMapOnImage(*args, **kwargs):
|
|
"""Object representing a segmentation map associated with an image."""
|
|
# pylint: disable=invalid-name
|
|
return SegmentationMapsOnImage(*args, **kwargs)
|
|
|
|
|
|
class SegmentationMapsOnImage(IAugmentable):
|
|
"""
|
|
Object representing a segmentation map associated with an image.
|
|
|
|
Attributes
|
|
----------
|
|
DEFAULT_SEGMENT_COLORS : list of tuple of int
|
|
Standard RGB colors to use during drawing, ordered by class index.
|
|
|
|
Parameters
|
|
----------
|
|
arr : (H,W) ndarray or (H,W,C) ndarray
|
|
Array representing the segmentation map(s). May have dtypes bool,
|
|
int or uint.
|
|
|
|
shape : tuple of int
|
|
Shape of the image on which the segmentation map(s) is/are placed.
|
|
**Not** the shape of the segmentation map(s) array, unless it is
|
|
identical to the image shape (note the likely difference between the
|
|
arrays in the number of channels).
|
|
This is expected to be ``(H, W)`` or ``(H, W, C)`` with ``C`` usually
|
|
being ``3``.
|
|
If there is no corresponding image, use ``(H_arr, W_arr)`` instead,
|
|
where ``H_arr`` is the height of the segmentation map(s) array
|
|
(analogous ``W_arr``).
|
|
|
|
nb_classes : None or int, optional
|
|
Deprecated.
|
|
|
|
"""
|
|
|
|
# TODO replace this by matplotlib colormap
|
|
DEFAULT_SEGMENT_COLORS = [
|
|
(0, 0, 0), # black
|
|
(230, 25, 75), # red
|
|
(60, 180, 75), # green
|
|
(255, 225, 25), # yellow
|
|
(0, 130, 200), # blue
|
|
(245, 130, 48), # orange
|
|
(145, 30, 180), # purple
|
|
(70, 240, 240), # cyan
|
|
(240, 50, 230), # magenta
|
|
(210, 245, 60), # lime
|
|
(250, 190, 190), # pink
|
|
(0, 128, 128), # teal
|
|
(230, 190, 255), # lavender
|
|
(170, 110, 40), # brown
|
|
(255, 250, 200), # beige
|
|
(128, 0, 0), # maroon
|
|
(170, 255, 195), # mint
|
|
(128, 128, 0), # olive
|
|
(255, 215, 180), # coral
|
|
(0, 0, 128), # navy
|
|
(128, 128, 128), # grey
|
|
(255, 255, 255), # white
|
|
# --
|
|
(115, 12, 37), # dark red
|
|
(30, 90, 37), # dark green
|
|
(127, 112, 12), # dark yellow
|
|
(0, 65, 100), # dark blue
|
|
(122, 65, 24), # dark orange
|
|
(72, 15, 90), # dark purple
|
|
(35, 120, 120), # dark cyan
|
|
(120, 25, 115), # dark magenta
|
|
(105, 122, 30), # dark lime
|
|
(125, 95, 95), # dark pink
|
|
(0, 64, 64), # dark teal
|
|
(115, 95, 127), # dark lavender
|
|
(85, 55, 20), # dark brown
|
|
(127, 125, 100), # dark beige
|
|
(64, 0, 0), # dark maroon
|
|
(85, 127, 97), # dark mint
|
|
(64, 64, 0), # dark olive
|
|
(127, 107, 90), # dark coral
|
|
(0, 0, 64), # dark navy
|
|
(64, 64, 64), # dark grey
|
|
]
|
|
|
|
def __init__(self, arr, shape, nb_classes=None):
|
|
assert ia.is_np_array(arr), (
|
|
"Expected to get numpy array, got %s." % (type(arr),))
|
|
assert arr.ndim in [2, 3], (
|
|
"Expected segmentation map array to be 2- or "
|
|
"3-dimensional, got %d dimensions and shape %s." % (
|
|
arr.ndim, arr.shape))
|
|
assert isinstance(shape, tuple), (
|
|
"Expected 'shape' to be a tuple denoting the shape of the image "
|
|
"on which the segmentation map is placed. Got type %s instead." % (
|
|
type(shape)))
|
|
|
|
if arr.dtype.kind == "f":
|
|
ia.warn_deprecated(
|
|
"Got a float array as the segmentation map in "
|
|
"SegmentationMapsOnImage. That is deprecated. Please provide "
|
|
"instead a (H,W,[C]) array of dtype bool_, int or uint, where "
|
|
"C denotes the segmentation map index."
|
|
)
|
|
|
|
if arr.ndim == 2:
|
|
arr = (arr > 0.5)
|
|
else: # arr.ndim == 3
|
|
arr = np.argmax(arr, axis=2).astype(np.int32)
|
|
|
|
if arr.dtype.name == "bool":
|
|
self._input_was = (arr.dtype, arr.ndim)
|
|
if arr.ndim == 2:
|
|
arr = arr[..., np.newaxis]
|
|
elif arr.dtype.kind in ["i", "u"]:
|
|
assert np.min(arr.flat[0:100]) >= 0, (
|
|
"Expected segmentation map array to only contain values >=0, "
|
|
"got a minimum of %d." % (np.min(arr),))
|
|
if arr.dtype.kind == "u":
|
|
# allow only <=uint16 due to conversion to int32
|
|
assert arr.dtype.itemsize <= 2, (
|
|
"When using uint arrays as segmentation maps, only uint8 "
|
|
"and uint16 are allowed. Got dtype %s." % (arr.dtype.name,)
|
|
)
|
|
elif arr.dtype.kind == "i":
|
|
# allow only <=uint16 due to conversion to int32
|
|
assert arr.dtype.itemsize <= 4, (
|
|
"When using int arrays as segmentation maps, only int8, "
|
|
"int16 and int32 are allowed. Got dtype %s." % (
|
|
arr.dtype.name,)
|
|
)
|
|
|
|
self._input_was = (arr.dtype, arr.ndim)
|
|
if arr.ndim == 2:
|
|
arr = arr[..., np.newaxis]
|
|
else:
|
|
raise Exception((
|
|
"Input was expected to be an array of dtype 'bool', 'int' "
|
|
"or 'uint'. Got dtype '%s'.") % (arr.dtype.name,))
|
|
|
|
if arr.dtype.name != "int32":
|
|
arr = arr.astype(np.int32)
|
|
|
|
self.arr = arr
|
|
self.shape = shape
|
|
|
|
if nb_classes is not None:
|
|
ia.warn_deprecated(
|
|
"Providing nb_classes to SegmentationMapsOnImage is no longer "
|
|
"necessary and hence deprecated. The argument is ignored "
|
|
"and can be safely removed.")
|
|
|
|
def get_arr(self):
|
|
"""Return the seg.map array, with original dtype and shape ndim.
|
|
|
|
Here, "original" denotes the dtype and number of shape dimensions that
|
|
was used when the :class:`SegmentationMapsOnImage` instance was
|
|
created, i.e. upon the call of
|
|
:func:`SegmentationMapsOnImage.__init__`.
|
|
Internally, this class may use a different dtype and shape to simplify
|
|
computations.
|
|
|
|
.. note::
|
|
|
|
The height and width may have changed compared to the original
|
|
input due to e.g. pooling operations.
|
|
|
|
Returns
|
|
-------
|
|
ndarray
|
|
Segmentation map array.
|
|
Same dtype and number of dimensions as was originally used when
|
|
the :class:`SegmentationMapsOnImage` instance was created.
|
|
|
|
"""
|
|
input_dtype, input_ndim = self._input_was
|
|
# The internally used int32 has a wider value range than any other
|
|
# input dtype, hence we can simply convert via astype() here.
|
|
arr_input = self.arr.astype(input_dtype)
|
|
if input_ndim == 2:
|
|
assert arr_input.shape[2] == 1, (
|
|
"Originally got a (H,W) segmentation map. Internal array "
|
|
"should now have shape (H,W,1), but got %s. This might be "
|
|
"an internal error." % (arr_input.shape,))
|
|
return arr_input[:, :, 0]
|
|
return arr_input
|
|
|
|
@ia.deprecated(alt_func="SegmentationMapsOnImage.get_arr()")
|
|
def get_arr_int(self, *args, **kwargs):
|
|
"""Return the seg.map array, with original dtype and shape ndim."""
|
|
# pylint: disable=unused-argument
|
|
return self.get_arr()
|
|
|
|
def draw(self, size=None, colors=None):
|
|
"""
|
|
Render the segmentation map as an RGB image.
|
|
|
|
Parameters
|
|
----------
|
|
size : None or float or iterable of int or iterable of float, optional
|
|
Size of the rendered RGB image as ``(height, width)``.
|
|
See :func:`~imgaug.imgaug.imresize_single_image` for details.
|
|
If set to ``None``, no resizing is performed and the size of the
|
|
segmentation map array is used.
|
|
|
|
colors : None or list of tuple of int, optional
|
|
Colors to use. One for each class to draw.
|
|
If ``None``, then default colors will be used.
|
|
|
|
Returns
|
|
-------
|
|
list of (H,W,3) ndarray
|
|
Rendered segmentation map (dtype is ``uint8``).
|
|
One per ``C`` in the original input array ``(H,W,C)``.
|
|
|
|
"""
|
|
def _handle_sizeval(sizeval, arr_axis_size):
|
|
if sizeval is None:
|
|
return arr_axis_size
|
|
if ia.is_single_float(sizeval):
|
|
return max(int(arr_axis_size * sizeval), 1)
|
|
if ia.is_single_integer(sizeval):
|
|
return sizeval
|
|
raise ValueError("Expected float or int, got %s." % (
|
|
type(sizeval),))
|
|
|
|
if size is None:
|
|
size = [size, size]
|
|
elif not ia.is_iterable(size):
|
|
size = [size, size]
|
|
|
|
height = _handle_sizeval(size[0], self.arr.shape[0])
|
|
width = _handle_sizeval(size[1], self.arr.shape[1])
|
|
image = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
return self.draw_on_image(
|
|
image,
|
|
alpha=1.0,
|
|
resize="segmentation_map",
|
|
colors=colors,
|
|
draw_background=True
|
|
)
|
|
|
|
def draw_on_image(self, image, alpha=0.75, resize="segmentation_map",
|
|
colors=None, draw_background=False,
|
|
background_class_id=0, background_threshold=None):
|
|
"""Draw the segmentation map as an overlay over an image.
|
|
|
|
Parameters
|
|
----------
|
|
image : (H,W,3) ndarray
|
|
Image onto which to draw the segmentation map. Expected dtype
|
|
is ``uint8``.
|
|
|
|
alpha : float, optional
|
|
Alpha/opacity value to use for the mixing of image and
|
|
segmentation map. Larger values mean that the segmentation map
|
|
will be more visible and the image less visible.
|
|
|
|
resize : {'segmentation_map', 'image'}, optional
|
|
In case of size differences between the image and segmentation
|
|
map, either the image or the segmentation map can be resized.
|
|
This parameter controls which of the two will be resized to the
|
|
other's size.
|
|
|
|
colors : None or list of tuple of int, optional
|
|
Colors to use. One for each class to draw.
|
|
If ``None``, then default colors will be used.
|
|
|
|
draw_background : bool, optional
|
|
If ``True``, the background will be drawn like any other class.
|
|
If ``False``, the background will not be drawn, i.e. the respective
|
|
background pixels will be identical with the image's RGB color at
|
|
the corresponding spatial location and no color overlay will be
|
|
applied.
|
|
|
|
background_class_id : int, optional
|
|
Class id to interpret as the background class.
|
|
See `draw_background`.
|
|
|
|
background_threshold : None, optional
|
|
Deprecated.
|
|
This parameter is ignored.
|
|
|
|
Returns
|
|
-------
|
|
list of (H,W,3) ndarray
|
|
Rendered overlays as ``uint8`` arrays.
|
|
Always a **list** containing one RGB image per segmentation map
|
|
array channel.
|
|
|
|
"""
|
|
if background_threshold is not None:
|
|
ia.warn_deprecated(
|
|
"The argument `background_threshold` is deprecated and "
|
|
"ignored. Please don't use it anymore.")
|
|
|
|
assert image.ndim == 3, (
|
|
"Expected to draw on 3-dimensional image, got image with %d "
|
|
"dimensions." % (image.ndim,))
|
|
assert image.shape[2] == 3, (
|
|
"Expected to draw on RGB image, got image with %d channels "
|
|
"instead." % (image.shape[2],))
|
|
assert image.dtype.name == "uint8", (
|
|
"Expected to get image with dtype uint8, got dtype %s." % (
|
|
image.dtype.name,))
|
|
assert 0 - 1e-8 <= alpha <= 1.0 + 1e-8, (
|
|
"Expected 'alpha' to be in interval [0.0, 1.0], got %.4f." % (
|
|
alpha,))
|
|
assert resize in ["segmentation_map", "image"], (
|
|
"Expected 'resize' to be \"segmentation_map\" or \"image\", got "
|
|
"%s." % (resize,))
|
|
|
|
colors = (
|
|
colors
|
|
if colors is not None
|
|
else SegmentationMapsOnImage.DEFAULT_SEGMENT_COLORS
|
|
)
|
|
|
|
if resize == "image":
|
|
image = ia.imresize_single_image(
|
|
image, self.arr.shape[0:2], interpolation="cubic")
|
|
|
|
segmaps_drawn = []
|
|
arr_channelwise = np.dsplit(self.arr, self.arr.shape[2])
|
|
for arr in arr_channelwise:
|
|
arr = arr[:, :, 0]
|
|
|
|
nb_classes = 1 + np.max(arr)
|
|
segmap_drawn = np.zeros((arr.shape[0], arr.shape[1], 3),
|
|
dtype=np.uint8)
|
|
assert nb_classes <= len(colors), (
|
|
"Can't draw all %d classes as it would exceed the maximum "
|
|
"number of %d available colors." % (nb_classes, len(colors),))
|
|
|
|
ids_in_map = np.unique(arr)
|
|
for c, color in zip(sm.xrange(nb_classes), colors):
|
|
if c in ids_in_map:
|
|
class_mask = (arr == c)
|
|
segmap_drawn[class_mask] = color
|
|
|
|
segmap_drawn = ia.imresize_single_image(
|
|
segmap_drawn, image.shape[0:2], interpolation="nearest")
|
|
|
|
segmap_on_image = blendlib.blend_alpha(segmap_drawn, image, alpha)
|
|
|
|
if draw_background:
|
|
mix = segmap_on_image
|
|
else:
|
|
foreground_mask = ia.imresize_single_image(
|
|
(arr != background_class_id),
|
|
image.shape[0:2],
|
|
interpolation="nearest")
|
|
# without this, the merge below does nothing
|
|
foreground_mask = np.atleast_3d(foreground_mask)
|
|
|
|
mix = (
|
|
(~foreground_mask) * image
|
|
+ foreground_mask * segmap_on_image
|
|
)
|
|
segmaps_drawn.append(mix)
|
|
return segmaps_drawn
|
|
|
|
def pad(self, top=0, right=0, bottom=0, left=0, mode="constant", cval=0):
|
|
"""Pad the segmentation maps at their top/right/bottom/left side.
|
|
|
|
Parameters
|
|
----------
|
|
top : int, optional
|
|
Amount of pixels to add at the top side of the segmentation map.
|
|
Must be ``0`` or greater.
|
|
|
|
right : int, optional
|
|
Amount of pixels to add at the right side of the segmentation map.
|
|
Must be ``0`` or greater.
|
|
|
|
bottom : int, optional
|
|
Amount of pixels to add at the bottom side of the segmentation map.
|
|
Must be ``0`` or greater.
|
|
|
|
left : int, optional
|
|
Amount of pixels to add at the left side of the segmentation map.
|
|
Must be ``0`` or greater.
|
|
|
|
mode : str, optional
|
|
Padding mode to use. See :func:`~imgaug.imgaug.pad` for details.
|
|
|
|
cval : number, optional
|
|
Value to use for padding if `mode` is ``constant``.
|
|
See :func:`~imgaug.imgaug.pad` for details.
|
|
|
|
Returns
|
|
-------
|
|
imgaug.augmentables.segmaps.SegmentationMapsOnImage
|
|
Padded segmentation map with height ``H'=H+top+bottom`` and
|
|
width ``W'=W+left+right``.
|
|
|
|
"""
|
|
from ..augmenters import size as iasize
|
|
arr_padded = iasize.pad(self.arr, top=top, right=right, bottom=bottom,
|
|
left=left, mode=mode, cval=cval)
|
|
return self.deepcopy(arr=arr_padded)
|
|
|
|
def pad_to_aspect_ratio(self, aspect_ratio, mode="constant", cval=0,
|
|
return_pad_amounts=False):
|
|
"""Pad the segmentation maps until they match a target aspect ratio.
|
|
|
|
Depending on which dimension is smaller (height or width), only the
|
|
corresponding sides (left/right or top/bottom) will be padded. In
|
|
each case, both of the sides will be padded equally.
|
|
|
|
Parameters
|
|
----------
|
|
aspect_ratio : float
|
|
Target aspect ratio, given as width/height. E.g. ``2.0`` denotes
|
|
the image having twice as much width as height.
|
|
|
|
mode : str, optional
|
|
Padding mode to use.
|
|
See :func:`~imgaug.imgaug.pad` for details.
|
|
|
|
cval : number, optional
|
|
Value to use for padding if `mode` is ``constant``.
|
|
See :func:`~imgaug.imgaug.pad` for details.
|
|
|
|
return_pad_amounts : bool, optional
|
|
If ``False``, then only the padded instance will be returned.
|
|
If ``True``, a tuple with two entries will be returned, where
|
|
the first entry is the padded instance and the second entry are
|
|
the amounts by which each array side was padded. These amounts are
|
|
again a tuple of the form ``(top, right, bottom, left)``, with
|
|
each value being an integer.
|
|
|
|
Returns
|
|
-------
|
|
imgaug.augmentables.segmaps.SegmentationMapsOnImage
|
|
Padded segmentation map as :class:`SegmentationMapsOnImage`
|
|
instance.
|
|
|
|
tuple of int
|
|
Amounts by which the instance's array was padded on each side,
|
|
given as a tuple ``(top, right, bottom, left)``.
|
|
This tuple is only returned if `return_pad_amounts` was set to
|
|
``True``.
|
|
|
|
"""
|
|
from ..augmenters import size as iasize
|
|
arr_padded, pad_amounts = iasize.pad_to_aspect_ratio(
|
|
self.arr,
|
|
aspect_ratio=aspect_ratio,
|
|
mode=mode,
|
|
cval=cval,
|
|
return_pad_amounts=True)
|
|
segmap = self.deepcopy(arr=arr_padded)
|
|
if return_pad_amounts:
|
|
return segmap, pad_amounts
|
|
return segmap
|
|
|
|
@ia.deprecated(alt_func="SegmentationMapsOnImage.resize()",
|
|
comment="resize() has the exactly same interface.")
|
|
def scale(self, *args, **kwargs):
|
|
"""Resize the seg.map(s) array given a target size and interpolation."""
|
|
return self.resize(*args, **kwargs)
|
|
|
|
def resize(self, sizes, interpolation="nearest"):
|
|
"""Resize the seg.map(s) array given a target size and interpolation.
|
|
|
|
Parameters
|
|
----------
|
|
sizes : float or iterable of int or iterable of float
|
|
New size of the array in ``(height, width)``.
|
|
See :func:`~imgaug.imgaug.imresize_single_image` for details.
|
|
|
|
interpolation : None or str or int, optional
|
|
The interpolation to use during resize.
|
|
Nearest neighbour interpolation (``"nearest"``) is almost always
|
|
the best choice.
|
|
See :func:`~imgaug.imgaug.imresize_single_image` for details.
|
|
|
|
Returns
|
|
-------
|
|
imgaug.augmentables.segmaps.SegmentationMapsOnImage
|
|
Resized segmentation map object.
|
|
|
|
"""
|
|
arr_resized = ia.imresize_single_image(self.arr, sizes,
|
|
interpolation=interpolation)
|
|
return self.deepcopy(arr_resized)
|
|
|
|
# TODO how best to handle changes to _input_was due to changed 'arr'?
|
|
def copy(self, arr=None, shape=None):
|
|
"""Create a shallow copy of the segmentation map object.
|
|
|
|
Parameters
|
|
----------
|
|
arr : None or (H,W) ndarray or (H,W,C) ndarray, optional
|
|
Optionally the `arr` attribute to use for the new segmentation map
|
|
instance. Will be copied from the old instance if not provided.
|
|
See
|
|
:func:`~imgaug.augmentables.segmaps.SegmentationMapsOnImage.__init__`
|
|
for details.
|
|
|
|
shape : None or tuple of int, optional
|
|
Optionally the shape attribute to use for the the new segmentation
|
|
map instance. Will be copied from the old instance if not provided.
|
|
See
|
|
:func:`~imgaug.augmentables.segmaps.SegmentationMapsOnImage.__init__`
|
|
for details.
|
|
|
|
Returns
|
|
-------
|
|
imgaug.augmentables.segmaps.SegmentationMapsOnImage
|
|
Shallow copy.
|
|
|
|
"""
|
|
# pylint: disable=protected-access
|
|
segmap = SegmentationMapsOnImage(
|
|
self.arr if arr is None else arr,
|
|
shape=self.shape if shape is None else shape)
|
|
segmap._input_was = self._input_was
|
|
return segmap
|
|
|
|
def deepcopy(self, arr=None, shape=None):
|
|
"""Create a deep copy of the segmentation map object.
|
|
|
|
Parameters
|
|
----------
|
|
arr : None or (H,W) ndarray or (H,W,C) ndarray, optional
|
|
Optionally the `arr` attribute to use for the new segmentation map
|
|
instance. Will be copied from the old instance if not provided.
|
|
See
|
|
:func:`~imgaug.augmentables.segmaps.SegmentationMapsOnImage.__init__`
|
|
for details.
|
|
|
|
shape : None or tuple of int, optional
|
|
Optionally the shape attribute to use for the the new segmentation
|
|
map instance. Will be copied from the old instance if not provided.
|
|
See
|
|
:func:`~imgaug.augmentables.segmaps.SegmentationMapsOnImage.__init__`
|
|
for details.
|
|
|
|
Returns
|
|
-------
|
|
imgaug.augmentables.segmaps.SegmentationMapsOnImage
|
|
Deep copy.
|
|
|
|
"""
|
|
# pylint: disable=protected-access
|
|
segmap = SegmentationMapsOnImage(
|
|
np.copy(self.arr if arr is None else arr),
|
|
shape=self.shape if shape is None else shape)
|
|
segmap._input_was = self._input_was
|
|
return segmap
|