1#!/usr/bin/env python3
2"""
3Given an image, this divides it into "slices" based on its dominant colour
4components.
5
6It starts by finding the dominant colours, then groups the pixels based
7on which colour they're closest to.  There's one slice per colour -- containing
8all the pixels which are closest to this colour.
9
10"""
11
12import argparse
13import os
14import subprocess
15
16from PIL import Image
17import tqdm
18
19
20def parse_hex(line):
21    return (int(line[1:3], 16), int(line[3:5], 16), int(line[5:7], 16))
22
23
24def parse_args():
25    parser = argparse.ArgumentParser(
26        description="Slice an image based on its dominant colours."
27    )
28    parser.add_argument("path", metavar="PATH", help="Path to the image to slice")
29    parser.add_argument(
30        "--max-colours",
31        dest="max_colours",
32        default=5,
33        help="how many colours to find",
34    )
35
36    return parser.parse_args()
37
38
39def get_dominant_colours(args):
40    """
41    Finds the dominant colours given the argumrnt.
42
43    Returns a list of RGB tuples, e.g.
44
45        [
46            (233, 228, 215),
47            (133, 139, 136),
48            (69, 118, 187),
49        ]
50
51    """
52    command = [
53        "dominant_colours",
54        "--no-palette",
55        args.path,
56        f"--max-colours={args.max_colours}",
57    ]
58    output = subprocess.check_output(command)
59
60    return [parse_hex(line) for line in output.decode("ascii").strip().splitlines()]
61
62
63if __name__ == "__main__":
64    args = parse_args()
65
66    dominant_colours = get_dominant_colours(args)
67
68    new_image_data = {colour: [] for colour in dominant_colours}
69
70    im = Image.open(args.path)
71
72    for pixel in tqdm.tqdm(im.getdata()):
73        if im.mode == "RGB":
74            r, g, b = pixel
75            alpha = 255
76        elif im.mode == "RGBA":
77            r, g, b, alpha = pixel
78        elif im.mode == "L":
79            r = g = b = pixel
80            alpha = 255
81        else:
82            raise ValueError(f"Unsupported image mode: {im.mode}")
83
84        closest_colour = min(
85            dominant_colours,
86            key=lambda c: (r - c[0]) ** 2 + (g - c[1]) ** 2 + (b - c[2]) ** 2,
87        )
88
89        for colour in new_image_data:
90            if colour == closest_colour:
91                new_image_data[colour].append((r, g, b, alpha))
92            else:
93                new_image_data[colour].append((255, 255, 255, 0))
94
95    for (r, g, b), data in new_image_data.items():
96        new_im = Image.new("RGBA", im.size)
97        new_im.putdata(data)
98
99        name, _ = os.path.splitext(os.path.basename(args.path))
100        out_path = f"{name}.{r:02x}{g:02x}{b:02x}.png"
101
102        new_im.save(out_path)