diff --git a/CHANGELOG.md b/CHANGELOG.md index 4369714d8..1e5f153fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,8 @@ * Added `augmenters.color.KMeansColorQuantization` and corresponding `augmenters.color.quantize_colors_kmeans()`. Both deal with quantizing similar colors using k-Means clustering. #347 + * Added a check script for `KMeansColorQuantization` under + `checks/check_kmeans_color_quantization.py`. #347 ## Fixes diff --git a/checks/check_kmeans_color_quantization.py b/checks/check_kmeans_color_quantization.py new file mode 100644 index 000000000..8007db072 --- /dev/null +++ b/checks/check_kmeans_color_quantization.py @@ -0,0 +1,31 @@ +from __future__ import print_function, division +import numpy as np +import imgaug as ia +import imgaug.augmenters as iaa + + +def main(): + image = ia.quokka_square((256, 256)) + image_q2 = iaa.quantize_colors_kmeans(image, 2) + image_q16 = iaa.quantize_colors_kmeans(image, 16) + ia.imshow(np.hstack([image_q2, image_q16])) + + from_cs = "RGB" + to_cs = ["RGB", "Lab"] + kwargs = {"from_colorspace": from_cs, "to_colorspace": to_cs} + augs = [ + iaa.KMeansColorQuantization(2, **kwargs), + iaa.KMeansColorQuantization(4, **kwargs), + iaa.KMeansColorQuantization(8, **kwargs), + iaa.KMeansColorQuantization((2, 16), **kwargs), + ] + + images_aug = [] + for aug in augs: + images_aug.extend(aug(images=[image]*8)) + + ia.imshow(ia.draw_grid(images_aug, cols=8)) + + +if __name__ == "__main__": + main()