Filter an Image based on an array of colors

filter_color.py

import argparse
import os.path
from PIL import Image

def create_palette(colors):
    result = []
    for c in colors:
        rgb_values = c.split(',')

        # there must be 3
        result.append(int(rgb_values[0]))
        result.append(int(rgb_values[1]))
        result.append(int(rgb_values[2]))

    return result

def filter_image(image, colors):
    palette_image = Image.new("P", (1, 1))

    palette = create_palette(colors)
    # set the palette, set all other colors to 0
    palette_image.putpalette(palette + [0, 0, 0] * (256-len(colors)))

    if(image.mode != 'RGB'):
        # convert to RGB as quantize requires it
        image = image.convert(mode='RGB')

    # apply the palette
    return image.quantize(palette=palette_image)

parser = argparse.ArgumentParser(description='Image Color Filter Utility')
parser.add_argument('-i', '--input', required=True, type=str,
                    help="path to original color image")
parser.add_argument('-c', '--color', required=False, type=str,
                    nargs="*", default=["255,255,255", "0,0,0"],
                    help="colors to filter from source image as R,G,B values")
parser.add_argument('-s', '--split', action='store_true',
                    help="should the final image be split on these colors and saved separate")

args = parser.parse_args()

filename = os.path.splitext(os.path.basename(args.input))[0]

# open source image
image = Image.open(args.input)

# filter on the color palette
new_image = filter_image(image, args.color)


if(args.split):
    for i in range(0, len(args.color)):
        img_split = new_image.copy()

        palette = []
        for j in range(0, (i * 3)):
            palette.append(255)

        palette = palette + [0, 0, 0]

        print(palette)
        img_split.putpalette(palette)
        img_split.save(f"{filename}_{i}.png", "PNG")
else:
    # save
    new_image.save(f"{filename}_filtered.png", "PNG")