Merge pull request #6 from caelestia-dots/python-rework

feat: rewrite in python

1 files changed, 16 insertions, 18 deletions

diff --git a/scheme/score.py b/src/caelestia/utils/material/score.py
index 18ddc21..7765050 100755..100644
--- a/scheme/score.py
+++ b/src/caelestia/utils/material/score.py
@@ -20,9 +20,8 @@ class Score:
         pass
 
     @staticmethod
-    def score(colors_to_population: dict) -> tuple[list[Hct], list[Hct]]:
+    def score(colors_to_population: dict, filter_enabled: bool = True) -> tuple[list[Hct], list[Hct]]:
         desired = 14
-        filter_enabled = False
         dislike_filter = True
 
         colors_hct = []
@@ -50,18 +49,11 @@ class Score:
             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
-            ):
+            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_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})
@@ -75,8 +67,7 @@ class Score:
             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
+                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_colors
                 )
                 if not duplicate_hue:
                     chosen_colors.append(hct)
@@ -102,8 +93,7 @@ class Score:
             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
+                    difference_degrees(hct.hue, chosen_hct.hue) < difference_degrees_ for chosen_hct in chosen_primaries
                 )
                 if not duplicate_hue:
                     chosen_primaries.append(hct)
@@ -119,16 +109,24 @@ class Score:
             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 l in colours:
+    for t in colours:
         if mode != "hex":
-            print("".join(["\x1b[48;2;{};{};{}m   \x1b[0m".format(*c.to_rgba()[:3]) for c in l]))
+            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 l]))
+            print(" ".join(["{:02X}{:02X}{:02X}".format(*c.to_rgba()[:3]) for c in t]))