def get_colors(self,
image_path,
num,
return_color_counts=False,
return_image=False,
use_logs=False):
if num < 2:
print("The image must have at least 2 colors.")
return
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
(h, w) = image.shape[:2]
# Convert color spaces
image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
image = image.reshape((image.shape[0] * image.shape[1], 3))
clusters = MiniBatchKMeans(n_clusters=num)
labels = clusters.fit_predict(image)
res = clusters.cluster_centers_.astype("uint8")[labels]
# To get the colors back to RGB, we must convert to an image and extract the colors
res = res.reshape((h, w, 3))
res = cv2.cvtColor(res, cv2.COLOR_LAB2BGR)
if return_image:
res_image = res
res = res.reshape((h * w, 3))
res_colors = []
color_counts = []
for color in res:
string_color = generate_hex(color, bgr=True)
if string_color not in res_colors:
res_colors.append(string_color)
color_counts.append(0)
else:
color_counts[res_colors.index(string_color)] += 1
palette = Palette()
for i, color in enumerate(res_colors):
palette.add(Color(color, color_counts[i]))
self.scheme_cache = palette
# Oh no
if return_image:
if return_color_counts:
return palette, color_counts, res_image
return palette, res_image
if return_color_counts:
return palette, color_counts
return palette
return res_colors, color_counts, res_image
return res_colors, res_image
if return_color_counts:
return res_colors, color_counts
return res_colors
def get_image_data_from_url(self, url):
req = urllib.urlopen(url)
return np.asarray(bytearray(req.read()), dtype=np.uint8)
if __name__ == "__main__":
print("Running...")
extractor = Extractor()
colors, counts, image = extractor.get_colors(
extractor.get_image_data_from_url(image_url),
colors,
return_color_counts=True,
return_image=True,
use_logs=True)
(h, w) = image.shape[:2]
if colors is not None:
print("Resulting colors:")
palette = Palette([])
for i, color in enumerate(colors):
palette.add(Color(color, counts[i]))
palette.sort(sort_palette_by)
display(palette.get_palette_html())
cv2_imshow(image)
