summaryrefslogtreecommitdiff
path: root/src/caelestia/utils/material/score.py
blob: 7765050b5c420c4e2a44f5b75eaec929538f784d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#!/usr/bin/env python

import sys

from materialyoucolor.dislike.dislike_analyzer import DislikeAnalyzer
from materialyoucolor.hct import Hct
from materialyoucolor.quantize import ImageQuantizeCelebi
from materialyoucolor.utils.math_utils import difference_degrees, sanitize_degrees_int


class Score:
    TARGET_CHROMA = 48.0
    WEIGHT_PROPORTION = 0.7
    WEIGHT_CHROMA_ABOVE = 0.3
    WEIGHT_CHROMA_BELOW = 0.1
    CUTOFF_CHROMA = 5.0
    CUTOFF_EXCITED_PROPORTION = 0.01

    def __init__(self):
        pass

    @staticmethod
    def score(colors_to_population: dict, filter_enabled: bool = True) -> tuple[list[Hct], list[Hct]]:
        desired = 14
        dislike_filter = True

        colors_hct = []
        hue_population = [0] * 360
        population_sum = 0

        for rgb, population in colors_to_population.items():
            hct = Hct.from_int(rgb)
            colors_hct.append(hct)
            hue = int(hct.hue)
            hue_population[hue] += population
            population_sum += population

        hue_excited_proportions = [0.0] * 360

        for hue in range(360):
            proportion = hue_population[hue] / population_sum
            for i in range(hue - 14, hue + 16):
                neighbor_hue = int(sanitize_degrees_int(i))
                hue_excited_proportions[neighbor_hue] += proportion

        # Score colours
        scored_hct = []
        for hct in colors_hct:
            hue = int(sanitize_degrees_int(round(hct.hue)))
            proportion = hue_excited_proportions[hue]

            if filter_enabled and (hct.chroma < Score.CUTOFF_CHROMA or proportion <= Score.CUTOFF_EXCITED_PROPORTION):
                continue

            proportion_score = proportion * 100.0 * Score.WEIGHT_PROPORTION
            chroma_weight = Score.WEIGHT_CHROMA_BELOW if hct.chroma < Score.TARGET_CHROMA else Score.WEIGHT_CHROMA_ABOVE
            chroma_score = (hct.chroma - Score.TARGET_CHROMA) * chroma_weight
            score = proportion_score + chroma_score
            scored_hct.append({"hct": hct, "score": score})

        scored_hct.sort(key=lambda x: x["score"], reverse=True)

        # Choose distinct colours
        chosen_colors = []
        for difference_degrees_ in range(90, 0, -1):
            chosen_colors.clear()
            for item in scored_hct:
                hct = item["hct"]
                duplicate_hue = any(
                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_colors
                )
                if not duplicate_hue:
                    chosen_colors.append(hct)
                if len(chosen_colors) >= desired:
                    break
            if len(chosen_colors) >= desired:
                break

        # Get primary colour
        primary = None
        for cutoff in range(20, 0, -1):
            for item in scored_hct:
                if item["hct"].chroma > cutoff and item["hct"].tone > cutoff * 3:
                    primary = item["hct"]
                    break
            if primary:
                break

        # Choose distinct primaries
        chosen_primaries = [primary]
        for difference_degrees_ in range(90, 14, -1):
            chosen_primaries = [primary]
            for item in scored_hct:
                hct = item["hct"]
                duplicate_hue = any(
                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_primaries
                )
                if not duplicate_hue:
                    chosen_primaries.append(hct)
                if len(chosen_primaries) >= 3:
                    break
            if len(chosen_primaries) >= 3:
                break

        # Fix disliked colours
        if dislike_filter:
            for i, chosen_hct in enumerate(chosen_primaries):
                chosen_primaries[i] = DislikeAnalyzer.fix_if_disliked(chosen_hct)
            for i, chosen_hct in enumerate(chosen_colors):
                chosen_colors[i] = DislikeAnalyzer.fix_if_disliked(chosen_hct)

        # Ensure enough colours
        if len(chosen_colors) < desired:
            return Score.score(colors_to_population, False)

        return chosen_primaries, chosen_colors


def score(image: str) -> tuple[list[Hct], list[Hct]]:
    return Score.score(ImageQuantizeCelebi(image, 1, 128))


if __name__ == "__main__":
    img = sys.argv[1]
    mode = sys.argv[2] if len(sys.argv) > 2 else "hex"

    colours = Score.score(ImageQuantizeCelebi(img, 1, 128))
    for t in colours:
        if mode != "hex":
            print("".join(["\x1b[48;2;{};{};{}m   \x1b[0m".format(*c.to_rgba()[:3]) for c in t]))
        if mode != "swatch":
            print(" ".join(["{:02X}{:02X}{:02X}".format(*c.to_rgba()[:3]) for c in t]))