Source code for imgaug.augmenters.artistic

"""
Augmenters that apply artistic image filters.

List of augmenters:

    * :class:`Cartoon`

Added in 0.4.0.

"""

from __future__ import print_function, division, absolute_import

import numpy as np
import cv2

from imgaug.imgaug import _normalize_cv2_input_arr_
from . import meta
from . import color as colorlib
from .. import dtypes as iadt
from .. import parameters as iap


[docs]def stylize_cartoon(image, blur_ksize=3, segmentation_size=1.0, saturation=2.0, edge_prevalence=1.0, suppress_edges=True, from_colorspace=colorlib.CSPACE_RGB): """Convert the style of an image to a more cartoonish one. This function was primarily designed for images with a size of ``200`` to ``800`` pixels. Smaller or larger images may cause issues. Note that the quality of the results can currently not compete with learned style transfer, let alone human-made images. A lack of detected edges or also too many detected edges are probably the most significant drawbacks. This method is loosely based on the one proposed in https://stackoverflow.com/a/11614479/3760780 Added in 0.4.0. **Supported dtypes**: * ``uint8``: yes; fully tested * ``uint16``: no * ``uint32``: no * ``uint64``: no * ``int8``: no * ``int16``: no * ``int32``: no * ``int64``: no * ``float16``: no * ``float32``: no * ``float64``: no * ``float128``: no * ``bool``: no Parameters ---------- image : ndarray A ``(H,W,3) uint8`` image array. blur_ksize : int, optional Kernel size of the median blur filter applied initially to the input image. Expected to be an odd value and ``>=0``. If an even value, thn automatically increased to an odd one. If ``<=1``, no blur will be applied. segmentation_size : float, optional Size multiplier to decrease/increase the base size of the initial mean-shift segmentation of the image. Expected to be ``>=0``. Note that the base size is increased by roughly a factor of two for images with height and/or width ``>=400``. edge_prevalence : float, optional Multiplier for the prevalance of edges. Higher values lead to more edges. Note that the default value of ``1.0`` is already fairly conservative, so there is limit effect from lowerin it further. saturation : float, optional Multiplier for the saturation. Set to ``1.0`` to not change the image's saturation. suppress_edges : bool, optional Whether to run edge suppression to remove blobs containing too many or too few edge pixels. from_colorspace : str, optional The source colorspace. Use one of ``imgaug.augmenters.color.CSPACE_*``. Defaults to ``RGB``. Returns ------- ndarray Image in cartoonish style. """ iadt.gate_dtypes( image, allowed=["uint8"], disallowed=["bool", "uint16", "uint32", "uint64", "uint128", "uint256", "int8", "int16", "int32", "int64", "int128", "int256", "float16", "float32", "float64", "float96", "float128", "float256"], augmenter=None) assert image.ndim == 3 and image.shape[2] == 3, ( "Expected to get a (H,W,C) image, got shape %s." % (image.shape,)) blur_ksize = max(int(np.round(blur_ksize)), 1) segmentation_size = max(segmentation_size, 0.0) saturation = max(saturation, 0.0) is_small_image = max(image.shape[0:2]) < 400 image = _blur_median(image, blur_ksize) image_seg = np.zeros_like(image) if is_small_image: spatial_window_radius = int(10 * segmentation_size) color_window_radius = int(20 * segmentation_size) else: spatial_window_radius = int(15 * segmentation_size) color_window_radius = int(40 * segmentation_size) if segmentation_size <= 0: image_seg = image else: cv2.pyrMeanShiftFiltering(_normalize_cv2_input_arr_(image), sp=spatial_window_radius, sr=color_window_radius, dst=image_seg) if is_small_image: edges_raw = _find_edges_canny(image_seg, edge_prevalence, from_colorspace) else: edges_raw = _find_edges_laplacian(image_seg, edge_prevalence, from_colorspace) edges = edges_raw edges = ((edges > 100) * 255).astype(np.uint8) if suppress_edges: # Suppress dense 3x3 blobs full of detected edges. They are visually # ugly. edges = _suppress_edge_blobs(edges, 3, 8, inverse=False) # Suppress spurious few-pixel edges (5x5 size with <=3 edge pixels). edges = _suppress_edge_blobs(edges, 5, 3, inverse=True) return _saturate(_blend_edges(image_seg, edges), saturation, from_colorspace)
# Added in 0.4.0. def _find_edges_canny(image, edge_multiplier, from_colorspace): image_gray = colorlib.change_colorspace_(np.copy(image), to_colorspace=colorlib.CSPACE_GRAY, from_colorspace=from_colorspace) image_gray = image_gray[..., 0] thresh = min(int(200 * (1/edge_multiplier)), 254) edges = cv2.Canny(_normalize_cv2_input_arr_(image_gray), thresh, thresh) return edges # Added in 0.4.0. def _find_edges_laplacian(image, edge_multiplier, from_colorspace): image_gray = colorlib.change_colorspace_(np.copy(image), to_colorspace=colorlib.CSPACE_GRAY, from_colorspace=from_colorspace) image_gray = image_gray[..., 0] edges_f = cv2.Laplacian(_normalize_cv2_input_arr_(image_gray / 255.0), cv2.CV_64F) edges_f = np.abs(edges_f) edges_f = edges_f ** 2 vmax = np.percentile(edges_f, min(int(90 * (1/edge_multiplier)), 99)) edges_f = np.clip(edges_f, 0.0, vmax) / vmax edges_uint8 = np.clip(np.round(edges_f * 255), 0, 255.0).astype(np.uint8) edges_uint8 = _blur_median(edges_uint8, 3) edges_uint8 = _threshold(edges_uint8, 50) return edges_uint8 # Added in 0.4.0. def _blur_median(image, ksize): if ksize % 2 == 0: ksize += 1 if ksize <= 1: return image return cv2.medianBlur(_normalize_cv2_input_arr_(image), ksize) # Added in 0.4.0. def _threshold(image, thresh): mask = (image < thresh) result = np.copy(image) result[mask] = 0 return result # Added in 0.4.0. def _suppress_edge_blobs(edges, size, thresh, inverse): kernel = np.ones((size, size), dtype=np.float32) counts = cv2.filter2D(_normalize_cv2_input_arr_(edges / 255.0), -1, kernel) if inverse: mask = (counts < thresh) else: mask = (counts >= thresh) edges = np.copy(edges) edges[mask] = 0 return edges # Added in 0.4.0. def _saturate(image, factor, from_colorspace): image = np.copy(image) if np.isclose(factor, 1.0, atol=1e-2): return image hsv = colorlib.change_colorspace_(image, to_colorspace=colorlib.CSPACE_HSV, from_colorspace=from_colorspace) sat = hsv[:, :, 1] sat = np.clip(sat.astype(np.int32) * factor, 0, 255).astype(np.uint8) hsv[:, :, 1] = sat image_sat = colorlib.change_colorspace_(hsv, to_colorspace=from_colorspace, from_colorspace=colorlib.CSPACE_HSV) return image_sat # Added in 0.4.0. def _blend_edges(image, image_edges): image_edges = 1.0 - (image_edges / 255.0) image_edges = np.tile(image_edges[..., np.newaxis], (1, 1, 3)) return np.clip( np.round(image * image_edges), 0.0, 255.0 ).astype(np.uint8)
[docs]class Cartoon(meta.Augmenter): """Convert the style of images to a more cartoonish one. This augmenter was primarily designed for images with a size of ``200`` to ``800`` pixels. Smaller or larger images may cause issues. Note that the quality of the results can currently not compete with learned style transfer, let alone human-made images. A lack of detected edges or also too many detected edges are probably the most significant drawbacks. Added in 0.4.0. **Supported dtypes**: See :func:`~imgaug.augmenters.artistic.stylize_cartoon`. Parameters ---------- blur_ksize : number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional Median filter kernel size. See :func:`~imgaug.augmenters.artistic.stylize_cartoon` for details. * If ``number``: That value will be used for all images. * If ``tuple (a, b) of number``: A random value will be uniformly sampled per image from the interval ``[a, b)``. * If ``list``: A random value will be picked per image from the ``list``. * If ``StochasticParameter``: The parameter will be queried once per batch for ``(N,)`` values, where ``N`` is the number of images. segmentation_size : number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional Mean-Shift segmentation size multiplier. See :func:`~imgaug.augmenters.artistic.stylize_cartoon` for details. * If ``number``: That value will be used for all images. * If ``tuple (a, b) of number``: A random value will be uniformly sampled per image from the interval ``[a, b)``. * If ``list``: A random value will be picked per image from the ``list``. * If ``StochasticParameter``: The parameter will be queried once per batch for ``(N,)`` values, where ``N`` is the number of images. saturation : number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional Saturation multiplier. See :func:`~imgaug.augmenters.artistic.stylize_cartoon` for details. * If ``number``: That value will be used for all images. * If ``tuple (a, b) of number``: A random value will be uniformly sampled per image from the interval ``[a, b)``. * If ``list``: A random value will be picked per image from the ``list``. * If ``StochasticParameter``: The parameter will be queried once per batch for ``(N,)`` values, where ``N`` is the number of images. edge_prevalence : number or tuple of number or list of number or imgaug.parameters.StochasticParameter, optional Multiplier for the prevalence of edges. See :func:`~imgaug.augmenters.artistic.stylize_cartoon` for details. * If ``number``: That value will be used for all images. * If ``tuple (a, b) of number``: A random value will be uniformly sampled per image from the interval ``[a, b)``. * If ``list``: A random value will be picked per image from the ``list``. * If ``StochasticParameter``: The parameter will be queried once per batch for ``(N,)`` values, where ``N`` is the number of images. from_colorspace : str, optional The source colorspace. Use one of ``imgaug.augmenters.color.CSPACE_*``. Defaults to ``RGB``. seed : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional See :func:`~imgaug.augmenters.meta.Augmenter.__init__`. name : None or str, optional See :func:`~imgaug.augmenters.meta.Augmenter.__init__`. random_state : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional Old name for parameter `seed`. Its usage will not yet cause a deprecation warning, but it is still recommended to use `seed` now. Outdated since 0.4.0. deterministic : bool, optional Deprecated since 0.4.0. See method ``to_deterministic()`` for an alternative and for details about what the "deterministic mode" actually does. Examples -------- >>> import imgaug.augmenters as iaa >>> aug = iaa.Cartoon() Create an example image, then apply a cartoon filter to it. >>> aug = iaa.Cartoon(blur_ksize=3, segmentation_size=1.0, >>> saturation=2.0, edge_prevalence=1.0) Create a non-stochastic cartoon augmenter that produces decent-looking images. """ # Added in 0.4.0. def __init__(self, blur_ksize=(1, 5), segmentation_size=(0.8, 1.2), saturation=(1.5, 2.5), edge_prevalence=(0.9, 1.1), from_colorspace=colorlib.CSPACE_RGB, seed=None, name=None, random_state="deprecated", deterministic="deprecated"): super(Cartoon, self).__init__( seed=seed, name=name, random_state=random_state, deterministic=deterministic) self.blur_ksize = iap.handle_continuous_param( blur_ksize, "blur_ksize", value_range=(0, None), tuple_to_uniform=True, list_to_choice=True) self.segmentation_size = iap.handle_continuous_param( segmentation_size, "segmentation_size", value_range=(0.0, None), tuple_to_uniform=True, list_to_choice=True) self.saturation = iap.handle_continuous_param( saturation, "saturation", value_range=(0.0, None), tuple_to_uniform=True, list_to_choice=True) self.edge_prevalence = iap.handle_continuous_param( edge_prevalence, "edge_prevalence", value_range=(0.0, None), tuple_to_uniform=True, list_to_choice=True) self.from_colorspace = from_colorspace # Added in 0.4.0. def _augment_batch_(self, batch, random_state, parents, hooks): if batch.images is not None: samples = self._draw_samples(batch, random_state) for i, image in enumerate(batch.images): image[...] = stylize_cartoon( image, blur_ksize=samples[0][i], segmentation_size=samples[1][i], saturation=samples[2][i], edge_prevalence=samples[3][i], from_colorspace=self.from_colorspace ) return batch # Added in 0.4.0. def _draw_samples(self, batch, random_state): nb_rows = batch.nb_rows return ( self.blur_ksize.draw_samples((nb_rows,), random_state=random_state), self.segmentation_size.draw_samples((nb_rows,), random_state=random_state), self.saturation.draw_samples((nb_rows,), random_state=random_state), self.edge_prevalence.draw_samples((nb_rows,), random_state=random_state) ) # Added in 0.4.0.
[docs] def get_parameters(self): """See :func:`~imgaug.augmenters.meta.Augmenter.get_parameters`.""" return [self.blur_ksize, self.segmentation_size, self.saturation, self.edge_prevalence, self.from_colorspace]