add palette util

cli/modules/palette.py

@@ -2,17 +2,21 @@
 # based on <https://towardsdatascience.com/image-color-extraction-with-python-in-4-steps-8d9370d9216e>
 
 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)