Add augmenter Cartoon

changelogs/master/added/20191020_cartoon.md

@@ -0,0 +1,5 @@
+# New Augmenter `Cartoon` #463
+
+* Added module `imgaug.augmenters.artistic`.
+* Added function `imgaug.augmenters.artistic.stylize_cartoon(image)`.
+* Added augmenter `imgaug.augmenters.artistic.Cartoon`.

imgaug/augmenters/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import absolute_import
 from imgaug.augmenters.arithmetic import *
+from imgaug.augmenters.artistic import *
 from imgaug.augmenters.blend import *
 from imgaug.augmenters.blur import *
 from imgaug.augmenters.color import *

imgaug/augmenters/artistic.py

@@ -0,0 +1,164 @@
+"""
+Augmenters that perform apply artistic image filters.
+
+List of augmenters:
+
+    * Cartoon
+
+"""
+
+from __future__ import print_function, division, absolute_import
+
+import numpy as np
+import cv2
+
+from . import meta
+from .. import dtypes as iadt
+
+
+def stylize_cartoon(image):
+    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)
+
+    is_small_image = max(image.shape[0:2]) < 400
+
+    image = _blur_median(image, 3)
+    image_seg = np.zeros_like(image)
+
+    if is_small_image:
+        spatial_window_radius = 10
+        color_window_radius = 20
+    else:
+        spatial_window_radius = 15
+        color_window_radius = 40
+
+    cv2.pyrMeanShiftFiltering(image,
+                              sp=spatial_window_radius,
+                              sr=color_window_radius,
+                              dst=image_seg)
+
+    if max(image.shape[0:2]) < 400:
+        edges_raw = _find_edges_canny(image_seg)
+    else:
+        edges_raw = _find_edges_laplacian(image_seg)
+
+    edges = edges_raw
+
+    edges = ((edges > 100) * 255).astype(np.uint8)
+    edges = _suppress_edge_blobs(edges, 3, 8, inverse=False)
+    edges = _suppress_edge_blobs(edges, 5, 3, inverse=True)
+
+    return _saturate(_blend_edges(image_seg, edges))
+
+
+def _saturate(image):
+    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
+    sat = hsv[:, :, 1]
+    sat = np.clip(sat.astype(np.int32) * 2, 0, 255).astype(np.uint8)
+    hsv[:, :, 1] = sat
+    image_sat = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+    return image_sat
+
+
+def _find_edges_canny(image):
+    image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(image_gray, 200, 200)
+    return edges
+
+
+def _find_edges_laplacian(image):
+    image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    edges_f = cv2.Laplacian(image_gray / 255.0, cv2.CV_64F)
+    edges_f = np.abs(edges_f)
+    edges_f = edges_f ** 2
+    vmax = np.percentile(edges_f, 90)
+    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
+
+
+def _blur_median(image, ksize):
+    return cv2.medianBlur(image, ksize)
+
+
+def _threshold(image, thresh):
+    assert image.ndim == 2
+    mask = (image < thresh)
+    result = np.copy(image)
+    result[mask] = 0
+    return result
+
+
+def _suppress_edge_blobs(edges, size, thresh, inverse):
+    kernel = np.ones((size, size), dtype=np.float32)
+    counts = cv2.filter2D(edges / 255.0, -1, kernel)
+
+    if inverse:
+        mask = (counts < thresh)
+    else:
+        mask = (counts >= thresh)
+
+    edges = np.copy(edges)
+    edges[mask] = 0
+    return edges
+
+
+def _blend_edges(image, image_edges):
+    assert image_edges.dtype.name == "uint8"
+    assert image_edges.ndim == 2
+    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)
+
+
+class Cartoon(meta.Augmenter):
+    """Convert the style of images to a more cartoonish one.
+
+    Parameters
+    ----------
+    name : None or str, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    deterministic : bool, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    random_state : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.bit_generator.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional
+        See :func:`imgaug.augmenters.meta.Augmenter.__init__`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import imgaug.augmenters as iaa
+    >>> image = np.arange(5*5*3).astype(np.uint8).reshape((5, 5, 3))
+    >>> aug = iaa.Cartoon()
+    >>> image_aug = aug(image=image)
+
+    Create an example image, then apply a cartoon filter to it.
+
+    """
+    def __init__(self, name=None, deterministic=False, random_state=None):
+        super(Cartoon, self).__init__(
+            name=name, deterministic=deterministic, random_state=random_state)
+
+    def _augment_batch(self, batch, random_state, parents, hooks):
+        if batch.images is not None:
+            for image in batch.images:
+                image[...] = stylize_cartoon(image)
+        return batch
+
+    def get_parameters(self):
+        return []