Spaces:
Runtime error
Runtime error
File size: 5,269 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
#!/usr/bin/env python
# based on <https://towardsdatascience.com/image-color-extraction-with-python-in-4-steps-8d9370d9216e>
import os
import io
import pathlib
import argparse
import importlib
import pandas as pd
import numpy as np
import extcolors
import filetype
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from colormap import rgb2hex
from PIL import Image
from util import log
grid = importlib.import_module('image-grid').grid
def color_to_df(param):
colors_pre_list = str(param).replace('([(','').split(', (')[0:-1]
df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list]
df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list]
#convert RGB to HEX code
df_color_up = [rgb2hex(int(i.split(", ")[0].replace("(","")),
int(i.split(", ")[1]),
int(i.split(", ")[2].replace(")",""))) for i in df_rgb]
df = pd.DataFrame(zip(df_color_up, df_percent), columns = ['c_code','occurence'])
return df
def palette(img, params, output):
size = 1024
img.thumbnail((size, size), Image.HAMMING)
#crate dataframe
colors_x = extcolors.extract_from_image(img, tolerance = params.color, limit = 13)
df_color = color_to_df(colors_x)
#annotate text
list_color = list(df_color['c_code'])
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': 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=2.5)
ab = AnnotationBbox(imagebox, (0, 0))
ax1.add_artist(ab)
#color palette
x_posi, y_posi, y_posi2 = 160, -260, -260
for c in list_color:
if list_color.index(c) <= 5:
y_posi += 240
rect = patches.Rectangle((x_posi, y_posi), 540, 230, facecolor = c)
ax2.add_patch(rect)
ax2.text(x = x_posi + 100, y = y_posi + 140, s = c, fontdict={'fontsize': 140}, color = 'white')
else:
y_posi2 += 240
rect = patches.Rectangle((x_posi + 600, y_posi2), 540, 230, facecolor = c)
ax2.add_artist(rect)
ax2.text(x = x_posi + 700, y = y_posi2 + 140, s = c, fontdict={'fontsize': 140}, color = 'white')
# add background to force layout
fig.set_facecolor('black')
ax2.axis('off')
tmp = Image.new('RGB', (2000, 1400), (0, 0, 0))
plt.imshow(tmp)
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()
if __name__ == '__main__':
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):
image = Image.open(arg)
fn = os.path.join(args.output, pathlib.Path(arg).stem + '-' + args.suffix + '.jpg')
palette(image, args, fn)
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)):
image = Image.open(os.path.join(root, f))
fn = os.path.join(args.output, pathlib.Path(f).stem + '-' + args.suffix + '.jpg')
palette(image, args, fn)
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)))
image = grid(images)
fn = os.path.join(args.output, args.suffix + '.jpg')
palette(image, args, fn)
|