add palette util

2026-09-19 01:04:32 +02:00 · 2023-02-06 08:37:08 -05:00
parent 9f1aa22b14
commit 2fdf23321c
3 changed files with 73 additions and 36 deletions
@@ -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)
@@ -0,0 +1,128 @@
+#!/bin/env python
+import os
+import io
+import pathlib
+import argparse
+import filetype
+import numpy as np
+from imwatermark import WatermarkEncoder, WatermarkDecoder
+from PIL import Image
+from PIL.ExifTags import TAGS
+from PIL.TiffImagePlugin import ImageFileDirectory_v2
+from modules.util import log, Map
+import piexif
+import piexif.helper
+
+
+options = Map({ 'method': 'dwtDctSvd', 'type': 'bytes' })
+
+
+def get_exif(image):
+    # using piexif
+    res1 = {}
+    try:
+        exif = piexif.load(image.info["exif"])
+        exif = exif.get("Exif", {})
+        for k, v in exif.items():
+            key = list(vars(piexif.ExifIFD).keys())[list(vars(piexif.ExifIFD).values()).index(k)]
+            res1[key] = piexif.helper.UserComment.load(v)
+    except:
+        pass
+    # using pillow
+    res2 = {}
+    try:
+        res2 = { TAGS[k]: v for k, v in image.getexif().items() if k in TAGS }
+    except:
+        pass
+    return {**res1, **res2}
+
+
+def set_exif(d: dict):
+    ifd = ImageFileDirectory_v2()
+    _TAGS = dict(((v, k) for k, v in TAGS.items())) # enumerate possible exif tags
+    for k, v in d.items():
+        ifd[_TAGS[k]] = v
+    exif_stream = io.BytesIO()
+    ifd.save(exif_stream)
+    bytes = b'Exif\x00\x00' + exif_stream.getvalue()
+    return bytes
+
+
+def get_watermark(image, args):
+    data = np.asarray(image)
+    decoder = WatermarkDecoder(options.type, args.length)
+    bytes = decoder.decode(data, options.method)
+    try:
+        watermark = str(bytes, 'UTF-8').replace('\x00', '')
+    except:
+        watermark = ''
+    return watermark
+
+
+def set_watermark(image, args):
+    data = np.asarray(image)
+    encoder = WatermarkEncoder()
+    encoder.set_watermark(options.type, args.wm.encode('utf-8'))
+    encoded = encoder.encode(data, options.method)
+    image = Image.fromarray(encoded)
+    return image
+
+
+def watermark(args, file):
+    if not os.path.exists(file):
+        log.error({ 'watermark': 'file not found' })
+        return
+    if not filetype.is_image(file):
+        log.error({ 'watermark': 'file is not an image' })
+        return
+    image = Image.open(file)
+    if image.width * image.height < 256 * 256:
+        log.error({ 'watermark': 'image too small' })
+        return
+
+    exif = get_exif(image)
+
+    if args.command == 'read':
+        watermark = get_watermark(image, args)
+        log.info({ 'file': file, 'watermark': watermark, 'exif': exif, 'resolution': f'{image.width}x{image.height}' })
+
+    elif args.command == 'write':
+        metadata = b'' if args.strip else set_exif(exif)
+        if args.output != '':
+            pathlib.Path(args.output).mkdir(parents = True, exist_ok = True)
+        image=set_watermark(image, args)
+        fn = os.path.join(args.output, file)
+        image.save(fn, exif=metadata)
+
+        if args.verify:
+            data = np.asarray(image)
+            decoder = WatermarkDecoder(options.type, args.length)
+            bytes = decoder.decode(data, options.method)
+            if bytes.startswith(b'\xff'):
+                watermark = ''
+            else:
+                watermark = str(bytes, 'UTF-8').replace('\x00', '')
+        else:
+            watermark = args.wm
+
+        log.info({ 'file': fn, 'watermark': watermark, 'exif': None if args.strip else exif, 'resolution': f'{image.width}x{image.height}' })
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'image watermarking')
+    parser.add_argument('command', choices = ['read', 'write'])
+    parser.add_argument('--wm', type=str, required=False, default='mm', help='watermark string')
+    parser.add_argument('--strip', default=False, action='store_true', help = "strip existing exif data")
+    parser.add_argument('--verify', default=False, action='store_true', help = "verify watermark during write")
+    parser.add_argument('--length', type=int, default=16, help="watermark length in bits")
+    parser.add_argument('--output', type=str, required=False, default='', help='folder to store images, default is overwrite in-place')
+    parser.add_argument('input', type=str, nargs='*')
+    args = parser.parse_args()
+    log.info({ 'watermark args': vars(args), 'options': options })
+    for arg in args.input:
+        if os.path.isfile(arg):
+            watermark(args, arg)
+        elif os.path.isdir(arg):
+            for root, _dirs, files in os.walk(arg):
+                for f in files:
+                    watermark(args, os.path.join(root, f))