Add an entry for "Notes from Small Planets"
[books.alexwlchan.net] / scripts / tint_colors.py
1 import collections
2 import colorsys
3 import json
4 import math
5 import os
6
7 from PIL import Image, UnidentifiedImageError
8 from sklearn.cluster import KMeans
9 import wcag_contrast_ratio as contrast
10
11
12 def _get_colors_from_im(im):
13     # Resizing means less pixels to handle, so the *k*-means clustering converges
14     # faster.  Small details are lost, but the main details will be preserved.
15     if im.size > (100, 100):
16         resize_ratio = min([100 / im.width, 100 / im.height])
17
18         new_width = int(im.width * resize_ratio)
19         new_height = int(im.height * resize_ratio)
20
21         im = im.resize((new_width, new_height))
22
23     # Ensure the image is RGB for consistency.
24     im = im.convert("RGB")
25
26     return list(im.getdata())
27
28
29 def get_colors_from(path):
30     """
31     Returns a list of the colors in the image at ``path``.
32     """
33     im = Image.open(str(path))
34
35     if getattr(im, "is_animated", False):
36         result = []
37
38         frame_count = im.n_frames
39
40         # Don't get all the frames from an animated GIF; if it has hundreds of
41         # frames this massively increases computation required for little gain.
42         # Take a sample and work from that.
43         for frame in range(0, frame_count, int(math.ceil(frame_count / 25))):
44             im.seek(frame)
45             result.extend(_get_colors_from_im(im))
46         return result
47     else:
48         return _get_colors_from_im(im)
49
50
51 def choose_tint_color_from_dominant_colors(dominant_colors, background_color):
52     """
53     Given a set of dominant colors (say, from a k-means algorithm) and the
54     background against which they'll be displayed, choose a tint color.
55
56     Both ``dominant_colors`` and ``background_color`` should be tuples in [0,1].
57     """
58     # Clamp colours to the range 0.0 - 1.0; occasionally sklearn has spat out
59     # numbers outside this range.
60     dominant_colors = [
61         (min(max(col[0], 0), 1), min(max(col[1], 0), 1), min(max(col[2], 0), 1))
62         for col in dominant_colors
63     ]
64
65     # The minimum contrast ratio for text and background to meet WCAG AA
66     # is 4.5:1, so discard any dominant colours with a lower contrast.
67     sufficient_contrast_colors = [
68         col for col in dominant_colors if contrast.rgb(col, background_color) >= 4.5
69     ]
70
71     # If none of the dominant colours meet WCAG AA with the background,
72     # try again with black and white -- every colour in the RGB space
73     # has a contrast ratio of 4.5:1 with at least one of these, so we'll
74     # get a tint colour, even if it's not a good one.
75     #
76     # Note: you could modify the dominant colours until one of them
77     # has sufficient contrast, but that's omitted here because it adds
78     # a lot of complexity for a relatively unusual case.
79     if not sufficient_contrast_colors:
80         return choose_tint_color_from_dominant_colors(
81             dominant_colors=dominant_colors + [(0, 0, 0), (1, 1, 1)],
82             background_color=background_color,
83         )
84
85     # Of the colors with sufficient contrast, pick the one with the
86     # highest saturation.  This is meant to optimise for colors that are
87     # more colourful/interesting than simple greys and browns.
88     hsv_candidates = {
89         tuple(rgb_col): colorsys.rgb_to_hsv(*rgb_col)
90         for rgb_col in sufficient_contrast_colors
91     }
92
93     return max(hsv_candidates, key=lambda rgb_col: hsv_candidates[rgb_col][2])
94
95
96 def choose_tint_color(p, *, background_color):
97     try:
98         background_color = {"black": (0, 0, 0), "white": (1, 1, 1)}[background_color]
99     except KeyError:  # pragma: no cover
100         raise ValueError(f"Unrecognised background color: {background_color!r}")
101
102     colors = get_colors_from(p)
103
104     # Normalise to [0, 1]
105     colors = [(r / 255, g / 255, b / 255) for (r, g, b) in colors]
106
107     pixel_tally = collections.Counter(colors)
108     most_common, most_common_count = pixel_tally.most_common(1)[0]
109     if (
110         most_common_count >= len(colors) * 0.15
111         and contrast.rgb(most_common, background_color) >= 4.5
112     ):
113         return most_common
114
115     dominant_colors = KMeans(n_clusters=12).fit(colors).cluster_centers_
116
117     return choose_tint_color_from_dominant_colors(
118         dominant_colors=dominant_colors, background_color=background_color
119     )
120
121
122 def get_tint_color_data():
123     try:
124         return json.load(open(os.path.join("src", "tint_colors.json")))
125     except FileNotFoundError:
126         return {}
127
128
129 def get_tint_colors():
130     return {path: data["color"] for (path, data) in get_tint_color_data().items()}
131
132
133 def store_tint_color(cover_path):
134     tint_colors = get_tint_color_data()
135
136     # If the size of a file has changed since the previous run, we need to
137     # recompute the tint colour.
138     try:
139         if tint_colors[os.path.basename(cover_path)]["size"] == os.stat(cover_path).st_size:
140             return
141     except KeyError as err:
142         print(f"Recomputing tint color for {cover_path}")
143
144     cover_color = choose_tint_color(cover_path, background_color="white")
145     tint_colors[os.path.basename(cover_path)] = {
146         "color": cover_color,
147         "size": os.stat(cover_path).st_size,
148     }
149
150     with open(os.path.join("src", "tint_colors.json"), "w") as outfile:
151         outfile.write(json.dumps(tint_colors, indent=2, sort_keys=True))