diff --git a/cli/modules/palette.py b/cli/modules/palette.py index b26ced7f6..cef0d4e0a 100755 --- a/cli/modules/palette.py +++ b/cli/modules/palette.py @@ -2,17 +2,21 @@ # based on import os -import sys +import io +import pathlib +import argparse import pandas as pd import numpy as np import extcolors +import filetype import matplotlib.pyplot as plt import matplotlib.patches as patches import matplotlib.image as mpimg from matplotlib.offsetbox import OffsetImage, AnnotationBbox from colormap import rgb2hex from PIL import Image - +from util import log +from grid import grid def color_to_df(input): colors_pre_list = str(input).replace('([(','').split(', (')[0:-1] @@ -26,16 +30,12 @@ def color_to_df(input): return df -def color_wheel(input_image, resize, tolerance, zoom): - #resize - img = Image.open(input_image) - if img.size[0] >= resize: - wpercent = (resize / float(img.size[0])) - hsize = int((float(img.size[1]) * float(wpercent))) - img = img.resize((resize, hsize)) +def palette(img, args, output): + size = 1024 + img.thumbnail((size, size), Image.HAMMING) #crate dataframe - colors_x = extcolors.extract_from_image(img, tolerance = tolerance, limit = 13) + colors_x = extcolors.extract_from_image(img, tolerance = args.color, limit = 13) df_color = color_to_df(colors_x) #annotate text @@ -43,50 +43,87 @@ def color_wheel(input_image, resize, tolerance, zoom): list_precent = [int(i) for i in list(df_color['occurence'])] text_c = [c + ' ' + str(round(p * 100 / sum(list_precent), 1)) +'%' for c, p in zip(list_color, list_precent)] fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(120,60), dpi=10) - + fig.set_facecolor('black') + #donut plot - wedges, _text = ax1.pie(list_precent, labels= text_c, labeldistance= 1.05, colors = list_color, textprops={'fontsize': 140, 'color':'black'}) + wedges, _text = ax1.pie(list_precent, labels= text_c, labeldistance= 1.05, colors = list_color, textprops={'fontsize': 100, 'color':'white'}) plt.setp(wedges, width=0.3) #add image in the center of donut plot data = np.asarray(img) - imagebox = OffsetImage(data, zoom=zoom) + imagebox = OffsetImage(data, zoom=2.5) ab = AnnotationBbox(imagebox, (0, 0)) ax1.add_artist(ab) #color palette - x_posi, y_posi, y_posi2 = 160, -200, -200 + x_posi, y_posi, y_posi2 = 160, -260, -260 for c in list_color: if list_color.index(c) <= 5: - y_posi += 220 - rect = patches.Rectangle((x_posi, y_posi), 480, 200, facecolor = c) + y_posi += 240 + rect = patches.Rectangle((x_posi, y_posi), 540, 230, facecolor = c) ax2.add_patch(rect) - ax2.text(x = x_posi + 40, y = y_posi + 120, s = c, fontdict={'fontsize': 140}) + ax2.text(x = x_posi + 100, y = y_posi + 140, s = c, fontdict={'fontsize': 140}, color = 'white') else: - y_posi2 += 220 - rect = patches.Rectangle((x_posi + 600, y_posi2), 480, 200, facecolor = c) + y_posi2 += 240 + rect = patches.Rectangle((x_posi + 600, y_posi2), 540, 230, facecolor = c) ax2.add_artist(rect) - ax2.text(x = x_posi + 640, y = y_posi2 + 120, s = c, fontdict={'fontsize': 140}) + ax2.text(x = x_posi + 700, y = y_posi2 + 140, s = c, fontdict={'fontsize': 140}, color = 'white') - #background - tmp_file = 'tmp.png' - fig, _ax = plt.subplots(figsize=(200,140),dpi=10) - fig.set_facecolor('white') - plt.savefig(tmp_file) - plt.close(fig) - - fig.set_facecolor('white') + # add background to force layout + fig.set_facecolor('black') ax2.axis('off') - tmp = plt.imread(tmp_file) + tmp = Image.new('RGB', (2000, 1400), (0, 0, 0)) plt.imshow(tmp) - plt.tight_layout() - plt.savefig('palette.jpg') + plt.tight_layout(rect = (-0.08, -0.2, 1.18, 1.05)) + + # save image + if output is not None: + buf = io.BytesIO() + plt.savefig(buf, format='png') + pltimg = Image.open(buf) + pltimg = pltimg.convert('RGB') + pltimg.save(output) + buf.close() + log.info({ 'palette created': output }) + plt.close() - os.remove(tmp_file) return if __name__ == '__main__': - sys.argv.pop(0) - for arg in sys.argv: - color_wheel(arg, 512, 10, 2) + parser = argparse.ArgumentParser(description = 'extract image color palette') + parser.add_argument('--color', type=int, default=20, help="color tolerance threshdold") + parser.add_argument('--output', type=str, required=False, default='', help='folder to store images') + parser.add_argument('--suffix', type=str, required=False, default='pallete', help='add suffix to image name') + parser.add_argument('--grid', default=False, action='store_true', help = "create grid of images before processing") + parser.add_argument('input', type=str, nargs='*') + args = parser.parse_args() + log.info({ 'palette args': vars(args) }) + if args.output != '': + pathlib.Path(args.output).mkdir(parents = True, exist_ok = True) + if not args.grid: + for arg in args.input: + if os.path.isfile(arg) and filetype.is_image(arg): + img = Image.open(arg) + output = os.path.join(args.output, pathlib.Path(arg).stem + '-' + args.suffix + '.jpg') + palette(img, args, output) + elif os.path.isdir(arg): + for root, _dirs, files in os.walk(arg): + for f in files: + if filetype.is_image(os.path.join(root, f)): + img = Image.open(os.path.join(root, f)) + output = os.path.join(args.output, pathlib.Path(f).stem + '-' + args.suffix + '.jpg') + palette(img, args, output) + else: + images = [] + for arg in args.input: + if os.path.isfile(arg) and filetype.is_image(arg): + images.append(Image.open(arg)) + elif os.path.isdir(arg): + for root, _dirs, files in os.walk(arg): + for f in files: + if filetype.is_image(os.path.join(root, f)): + images.append(Image.open(os.path.join(root, f))) + img = grid(images) + output = os.path.join(args.output, args.suffix + '.jpg') + palette(img, args, output) diff --git a/cli/watermark.py b/cli/modules/watermark.py similarity index 100% rename from cli/watermark.py rename to cli/modules/watermark.py diff --git a/extensions-builtin/stable-diffusion-webui-images-browser b/extensions-builtin/stable-diffusion-webui-images-browser index 818d59a66..8bc7d3a11 160000 --- a/extensions-builtin/stable-diffusion-webui-images-browser +++ b/extensions-builtin/stable-diffusion-webui-images-browser @@ -1 +1 @@ -Subproject commit 818d59a6629ed70d13890ccf5d0299939a550596 +Subproject commit 8bc7d3a11b4187d53905464327ea729f5c99f243