Extract dominant colors from image using numpy, opencv, and imagemagick

palette-extraction.py

import sys

import cv2
import numpy as np
from wand.color import Color
from wand.drawing import Drawing
from wand.image import Image

MAX_IMAGE_HEIGHT = 600
COLORSCHEME_COUNT = 12
COLORSCHEME_HEIGHT = 300


def resize_sampling(img_arr: np.ndarray,
                    sample_count: int,
                    interpolation=cv2.INTER_NEAREST) -> list:
    samples = cv2.resize(img_arr,
                         (1, sample_count),
                         interpolation=interpolation)
    samples = [x[0] for x in samples.tolist()]
    return samples


def median_sampling(img_arr: np.ndarray, sample_count: int) -> list:
    height = img_arr.shape[0]
    width = img_arr.shape[1]
    sample_width = round(width / sample_count)

    samples = []
    for i in range(sample_count):
        start_idx = i * sample_width
        last_idx = (i+1) * sample_width
        sample = img_arr[0:height, start_idx:last_idx]
        averaged_sample = np.median(sample, (0, 1))
        samples.append(averaged_sample.astype('uint8').tolist())
    return samples


def main():
    buf = sys.stdin.buffer.read()
    img_arr = np.frombuffer(buf, dtype='uint8')
    img_arr = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
    img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)
    samples = median_sampling(img_arr, COLORSCHEME_COUNT)

    image = Image(blob=buf)
    if image.height > MAX_IMAGE_HEIGHT:
        mult = image.height / MAX_IMAGE_HEIGHT
        new_width = round(image.width / mult)
        new_height = round(image.height / mult)
        image.resize(new_width, new_height)

    canvas = Image(width=image.width, height=image.height + COLORSCHEME_HEIGHT)
    sample_width = round(canvas.width / COLORSCHEME_COUNT)
    with Drawing() as draw:
        draw.composite(operator='over',
                       left=0,
                       top=COLORSCHEME_HEIGHT,
                       width=image.width,
                       height=image.height,
                       image=image)
        draw.push()
        for i, sample in enumerate(samples):
            color_str = 'rgb(%d, %d, %d)' % (sample[0], sample[1], sample[2])
            left = i * sample_width
            draw.fill_color = Color(color_str)
            draw.rectangle(left=left,
                           top=0,
                           right=canvas.width,
                           bottom=COLORSCHEME_HEIGHT)
            draw.push()
        draw(canvas)
    result = canvas.make_blob(format='png')
    sys.stdout.buffer.write(result)


if __name__ == '__main__':
    main()