Implement KMeans and MedianCut extractors in terms of the extractor base class (#55)

* Extract methods to separate file

* Add a protocol and base class for color extractors

* Implement median cut with a colorextractor class

* Implement KMeans using the Extractor Base Class

* Fix misc typing in utils

* Bump patch version, and update Changelog

Author: qTipTip

CHANGELOG.md

@@ -10,6 +10,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Released
 
+## 4.0.1 27/01/2025
+
+### Added
+
+- A `ColorExtractor`-protocol that defines an interface for color extractors.
+- Create a `ColorExtractorBase` abstract class that extractors can inherit from to implement the interface.
+
+### Changed
+- The implementation of `median_cut_extraction` and `k_means_extraction` is now
+implemented as in terms of subclasses of the `ColorExtractorBase`
+
 ## [4.0.0] 08/10/2024
 
 ### Changed

Pylette/src/color_extraction.py

@@ -8,11 +8,10 @@
 import requests  # type: ignore
 from numpy.typing import NDArray
 from PIL import Image
-from sklearn.cluster import KMeans
 
-from Pylette.src.color import Color
+from Pylette.src.extractors.k_means import k_means_extraction
+from Pylette.src.extractors.median_cut import median_cut_extraction
 from Pylette.src.palette import Palette
-from Pylette.src.utils import ColorBox
 
 ImageType_T: TypeAlias = Union["os.PathLike[Any]", bytes, NDArray[float], str]
 
@@ -25,35 +24,6 @@ class ImageType(str, Enum):
     NONE = "none"
 
 
-def median_cut_extraction(arr: np.ndarray, height: int, width: int, palette_size: int) -> list[Color]:
-    """
-    Extracts a color palette using the median cut algorithm.
-
-    Parameters:
-        arr (np.ndarray): The input array.
-        height (int): The height of the image.
-        width (int): The width of the image.
-        palette_size (int): The number of colors to extract from the image.
-
-    Returns:
-        list[Color]: A list of colors extracted from the image.
-    """
-
-    arr = arr.reshape((width * height, -1))
-    c = [ColorBox(arr)]
-
-    # Each iteration, find the largest box, split it, remove original box from list of boxes, and add the two new boxes.
-    while len(c) < palette_size:
-        largest_c_idx = np.argmax(c)
-        # add the two new boxes to the list, while removing the split box.
-        c = c[:largest_c_idx] + c[largest_c_idx].split() + c[largest_c_idx + 1 :]
-
-    total_pixels = width * height
-    colors = [Color(tuple(map(int, box.average)), box.pixel_count / total_pixels) for box in c]
-
-    return colors
-
-
 def _parse_image_type(image: ImageType_T) -> ImageType:
     """
     Determines the type of the input image.
@@ -169,28 +139,3 @@ def request_image(image_url: str) -> Image.Image:
         return img
     else:
         raise ValueError("The URL did not point to a valid image.")
-
-
-def k_means_extraction(arr: NDArray[float], height: int, width: int, palette_size: int) -> list[Color]:
-    """
-    Extracts a color palette using KMeans.
-
-    Parameters:
-        arr (NDArray[float]): The input array.
-        height (int): The height of the image.
-        width (int): The width of the image.
-        palette_size (int): The number of colors to extract from the image.
-
-    Returns:
-        list[Color]: A palette of colors sorted by frequency.
-    """
-    arr = np.reshape(arr, (width * height, -1))
-    model = KMeans(n_clusters=palette_size, n_init="auto", init="k-means++", random_state=2024)
-    labels = model.fit_predict(arr)
-    palette = np.array(model.cluster_centers_, dtype=int)
-    color_count = np.bincount(labels)
-    color_frequency = color_count / float(np.sum(color_count))
-    colors = []
-    for color, freq in zip(palette, color_frequency):
-        colors.append(Color(color, freq))
-    return colors

Pylette/src/extractors/k_means.py

@@ -0,0 +1,48 @@
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.cluster import KMeans
+
+from Pylette.src.color import Color
+from Pylette.src.extractors.protocol import NP_T, ColorExtractorBase
+
+
+class KMeansExtractor(ColorExtractorBase):
+    def extract(self, arr: NDArray[NP_T], height: int, width: int, palette_size: int) -> list[Color]:
+        """
+        Extracts a color palette using KMeans.
+
+        Parameters:
+            arr (NDArray[float]): The input array.
+            height (int): The height of the image.
+            width (int): The width of the image.
+            palette_size (int): The number of colors to extract from the image.
+
+        Returns:
+            list[Color]: A palette of colors sorted by frequency.
+        """
+        arr = np.reshape(arr, (width * height, -1))
+        model = KMeans(n_clusters=palette_size, n_init="auto", init="k-means++", random_state=2024)
+        labels = model.fit_predict(arr)
+        palette = np.array(model.cluster_centers_, dtype=int)
+        color_count = np.bincount(labels)
+        color_frequency = color_count / float(np.sum(color_count))
+        colors = []
+        for color, freq in zip(palette, color_frequency):
+            colors.append(Color(color, freq))
+        return colors
+
+
+def k_means_extraction(arr: NDArray[NP_T], height: int, width: int, palette_size: int) -> list[Color]:
+    """
+    Extracts a color palette using KMeans.
+
+    Parameters:
+        arr (NDArray[float]): The input array.
+        height (int): The height of the image.
+        width (int): The width of the image.
+        palette_size (int): The number of colors to extract from the image.
+
+    Returns:
+        list[Color]: A palette of colors sorted by frequency.
+    """
+    return KMeansExtractor().extract(arr=arr, height=height, width=width, palette_size=palette_size)

Pylette/src/extractors/median_cut.py

@@ -0,0 +1,49 @@
+import numpy as np
+from numpy.typing import NDArray
+
+from Pylette.src.color import Color
+from Pylette.src.extractors.protocol import NP_T, ColorExtractorBase
+from Pylette.src.utils import ColorBox
+
+
+class MedianCutExtractor(ColorExtractorBase):
+    def extract(self, arr: NDArray[NP_T], height: int, width: int, palette_size: int) -> list[Color]:
+        """
+        Extracts a color palette using the median cut algorithm.
+
+        Parameters:
+            arr (np.ndarray): The input array.
+            height (int): The height of the image.
+            width (int): The width of the image.
+            palette_size (int): The number of colors to extract from the image.
+
+        Returns:
+            list[Color]: A list of colors extracted from the image.
+        """
+
+        arr = self._reshape_array(arr=arr, height=height, width=width)
+
+        boxes = [ColorBox(arr)]
+        while len(boxes) < palette_size:
+            largest_box_idx = np.argmax(boxes)  # type: ignore
+            boxes = boxes[:largest_box_idx] + boxes[largest_box_idx].split() + boxes[largest_box_idx + 1 :]
+
+        total_pixels = width * height
+        return [Color(tuple(map(int, box.average)), box.pixel_count / total_pixels) for box in boxes]
+
+
+def median_cut_extraction(arr: np.ndarray, height: int, width: int, palette_size: int) -> list[Color]:
+    """
+    Extracts a color palette using the median cut algorithm.
+
+    Parameters:
+        arr (np.ndarray): The input array.
+        height (int): The height of the image.
+        width (int): The width of the image.
+        palette_size (int): The number of colors to extract from the image.
+
+    Returns:
+        list[Color]: A list of colors extracted from the image.
+    """
+
+    return MedianCutExtractor().extract(arr, height, width, palette_size)

Pylette/src/extractors/protocol.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+from typing import Protocol, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from Pylette.src.color import Color
+
+NP_T = TypeVar("NP_T", bound=np.generic, covariant=True)
+
+
+class ColorExtractor(Protocol):
+    def extract(self, arr: NDArray[NP_T], height: int, width: int, palette_size: int) -> list[Color]: ...
+
+
+class ColorExtractorBase(ABC):
+    @abstractmethod
+    def extract(self, arr: NDArray[NP_T], height: int, width: int, palette_size: int) -> list[Color]:
+        pass
+
+    def _reshape_array(self, arr: NDArray[NP_T], height: int, width: int) -> NDArray[NP_T]:
+        return arr.reshape((height * width, -1))

Pylette/src/utils.py

@@ -23,8 +23,8 @@ def _get_min_max(self) -> None:
         """
         Calculates the minimum and maximum values for each color channel in the ColorBox.
         """
-        self.min_channel: NDArray[np.uint8, (3,)] = np.min(self.colors, axis=0)
-        self.max_channel: NDArray[np.uint8, (3,)] = np.max(self.colors, axis=0)
+        self.min_channel: NDArray[np.uint8] = np.min(self.colors, axis=0)
+        self.max_channel: NDArray[np.uint8] = np.max(self.colors, axis=0)
 
     def __lt__(self, other: "ColorBox") -> bool:
         """
@@ -36,10 +36,10 @@ def __lt__(self, other: "ColorBox") -> bool:
         Returns:
         bool: True if the volume of this ColorBox is less than the volume of the other ColorBox, False otherwise.
         """
-        return self.size < other.size
+        return bool(self.size < other.size)
 
     @property
-    def size(self) -> np.uint64:
+    def size(self) -> int:
         """
         Returns the volume of the ColorBox.
 
@@ -55,12 +55,12 @@ def _get_dominant_channel(self) -> int:
         Returns:
             int: The index of the dominant color channel.
         """
-        diff: NDArray[np.uint8, (3,)] = self.max_channel - self.min_channel
+        diff: NDArray[np.uint8] = self.max_channel - self.min_channel
         dominant_channel = np.argmax(diff)
-        return dominant_channel
+        return int(dominant_channel)
 
     @property
-    def average(self) -> np.ndarray:
+    def average(self) -> NDArray[np.uint8]:
         """
         Calculates the average color contained in the ColorBox.
 
@@ -80,7 +80,7 @@ def volume(self) -> int:
         Returns:
             int: The volume of the ColorBox.
         """
-        diff: NDArray[np.uint8, (3,)] = self.max_channel - self.min_channel
+        diff: NDArray[np.uint8] = self.max_channel - self.min_channel
         return np.prod(diff).item()
 
     def split(self) -> list["ColorBox"]:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pylette"
-version = "4.0.0"
+version = "4.0.1"
 description = "A Python library for extracting color palettes from images."
 authors = ["Ivar Stangeby"]
 license = "MIT"