From b389b2304d79366a6644f72e90912020807207fa Mon Sep 17 00:00:00 2001
From: SyfSchydea <syf.p@live.com>
Date: Mon, 20 Apr 2020 08:59:46 +0100
Subject: [PATCH] Fix to handle images with very few unique colours

The `find_palette` function now attempts to find a palette using the
exact colours in the image. If this is not possible due to the image
having 16 or more unique colours, it falls back on the traditional
method.

Issue: #44
---
 palette.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/palette.py b/palette.py
index edcdf61..b7bd8e7 100644
--- a/palette.py
+++ b/palette.py
@@ -201,6 +201,40 @@ def create_indexed_img_dithered(img, palette):
 
 	return indexed
 
+# Attempt to find an exact palette for the image. This will only succeed if the list of colours has `max_size`
+# unique colours or fewer. If there are more unique colours, this operation is impossible, and the function will
+# return None.
+#
+# param colours  - NumPy array of colours
+# param max_size - Maximum number of colours allowed in the palette
+#
+# return         - NumPy array of palette colours if possible, or None if not.
+def find_small_palette(colours, max_size=15):
+	_, depth = colours.shape
+	palette = []
+
+	# Iterate through each colour in the image
+	for col in colours:
+		found_colour = False
+
+		for palette_colour in palette:
+			if np.array_equal(col, palette_colour):
+				found_colour = True
+				break
+
+		if found_colour:
+			continue
+
+		# If this colour is not yet in the palette, add it
+		palette.append(col)
+
+		# If this makes the palette too big, this image has too many
+		# colours to produce an exact palette, so return None here.
+		if len(palette) > max_size:
+			return None
+
+	return np.array(palette)
+
 # Take an image in Lab colour space, and output a colour palette in NH colours
 #
 # param img           - Input Lab image.
@@ -218,6 +252,21 @@ def pick_palette(img, palette_size, seed, weight_map, *, retry_on_dupe=False, lo
 	height, width, depth = img.shape
 	colours = img.reshape((height * width, depth))
 
+	small_palette = find_small_palette(colours, 15)
+	if small_palette is not None:
+		log.info("Found exact palette")
+		log.debug("Palette: \n", small_palette)
+
+		hsv_palette = colour.convert_palette(small_palette, cv2.COLOR_Lab2RGB, cv2.COLOR_RGB2HSV)
+		nh_palette = colour.palette_hsv_to_nh(hsv_palette)
+		log.debug("NH Palette:\n", nh_palette)
+
+		duplicate_colours = colour.get_duplicate_colours(nh_palette)
+		if len(duplicate_colours) > 0:
+			log.warn("Repeated colours in palette:", *duplicate_colours)
+
+		return nh_palette, small_palette
+
 	if weight_map is not None:
 		weight_map = weight_map.reshape(height * width)
 	else: