Repository: qTipTip/Pylette

Issue:
Feature - Add Image/bytes or IO support
On my app I already have a image in memory. To use extract_colors I need to save it on a temp file, will be great if I can pass the bytes just.

Something like this:
```python
class ColorPaletteExtractor(Protocol):
    async def __call__(self, asset: bytes) -> List[str]:
        palette = extract_colors(asset, palette_size=10, resize=True)
        ...
```

Hints from issue comments:
@qTipTip: Great suggestion! 

I'll try to get this implemented! @qTipTip: Can you look at #28 and tell me what you think? You can pass your bytes in the `image_bytes`-argument. 

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 484bd86..b13adf3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Released
 
+## [2.1.1] 09/06/2024
+
+### Changed
+- Added `random_state` to the `KMeans` algorithm to ensure reproducibility. Two separate KMeans extractions now yield the same palettes. Thanks to @opeyemibami for fixing.
+
 ## [2.1.0] 07/06/2024
 
 ### Added
diff --git a/Pylette/src/color_extraction.py b/Pylette/src/color_extraction.py
index a399454..5212405 100644
--- a/Pylette/src/color_extraction.py
+++ b/Pylette/src/color_extraction.py
@@ -42,15 +42,17 @@ def median_cut_extraction(
 
 
 def extract_colors(
-    image=None,
+    image_bytes: bytes | None = None,
+    image: str | None = None,
     image_url: str | None = None,
     palette_size: int = 5,
     resize: bool = True,
     mode: Literal["KM"] | Literal["MC"] = "KM",
-    sort_mode: Literal["luminance"] | None = None,
+    sort_mode: Literal["luminance", "frequency"] | None = None,
 ) -> Palette:
     """
     Extracts a set of 'palette_size' colors from the given image.
+    :param image_bytes: bytes representing the image data
     :param image: path to Image file
     :param image_url: url to the image-file
     :param palette_size: number of colors to extract
@@ -60,10 +62,12 @@ def extract_colors(
     :return: a list of the extracted colors
     """
 
-    if image is None and image_url is None:
+    if image_bytes is None and image is None and image_url is None:
         raise ValueError("No image provided")
 
-    if image is None and image_url is not None:
+    if image_bytes is not None:
+        img = Image.open(BytesIO(image_bytes)).convert("RGB")
+    elif image is None and image_url is not None:
         response = requests.get(image_url)
         # Check if the request was successful and content type is an image
         if response.status_code == 200 and "image" in response.headers.get(
@@ -108,7 +112,9 @@ def k_means_extraction(
     :return: 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++")
+    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)
diff --git a/README.md b/README.md
index f477ef3..1e396c3 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,7 @@ palette = extract_colors(image='image.jpg', palette_size=10, resize=True)
 palette = extract_colors(image_url='https://path.to.image', palette_size=10, resize=True, mode='MC', sort_mode='luminance')
 ```
 
-This yields a palette of ten colors, and the `resize` flag tells Pylette to resize the image to a more manageable size before
+This yields a palette of ten colors, and the `resize` flag tells Pylette to resize the image to a more manageable size (256 x 256) before
 beginning color extraction. This significantly speeds up the extraction, but reduces the faithfulness of the color palette.
 One can choose between color quantization using K-Means (default) or Median-Cut algorithms, by setting in the `mode`-parameter. One can also specify to alternatively sort the color palette by the luminance (percieved brightness).
 
@@ -117,11 +117,10 @@ options:
 
 ## Under the hood
 
-Currently, Pylette uses KMeans for the color quantization. There are plans for implementing other color quantization schemes, like:
+Currently, Pylette can use KMeans or Median-cut for the color quantization. There are plans for implementing other color quantization schemes, like:
 
-1. Median-cut [Implemented]
-2. Octree
-3. Modified minmax
+1. Octree
+2. Modified minmax
 
 The article [*Improving the Performance of K-Means for Color Quantization*](https://arxiv.org/pdf/1101.0395.pdf) gives a
 nice overview of available methods.
diff --git a/pyproject.toml b/pyproject.toml
index 29fdf37..cd7c673 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pylette"
-version = "2.1.0"
+version = "2.1.1"
 description = "A Python library for extracting color palettes from images."
 authors = ["Ivar Stangeby"]
 license = "MIT"


Test patch:
diff --git a/tests/integration/test_colorspaces.py b/tests/integration/test_colorspaces.py
index 82a7b41..05443b4 100644
--- a/tests/integration/test_colorspaces.py
+++ b/tests/integration/test_colorspaces.py
@@ -1,6 +1,5 @@
 import pathlib
 
-import numpy as np
 import pytest
 from numpy.testing import assert_approx_equal
 
@@ -14,6 +13,14 @@ def test_image_path_as_str():
     yield str(test_image.absolute().resolve())
 
 
+@pytest.fixture
+def test_image_as_bytes():
+    test_image = pathlib.Path(__file__).parent.parent / "data/test_image.png"
+
+    with open(test_image, "rb") as f:
+        yield f.read()
+
+
 @pytest.fixture()
 def test_kmean_extracted_palette(test_image_path_as_str):
     return extract_colors(
@@ -70,6 +77,55 @@ def test_palette_invariants_with_image_path(
     )
 
 
+@pytest.mark.parametrize("palette_size", [1, 5, 10, 100])
+@pytest.mark.parametrize(
+    "extraction_mode",
+    [
+        "KM",
+        pytest.param(
+            "MC",
+            marks=pytest.mark.skip(
+                "Currently a bug in the MC algorithm, causing frequencies not summing to one"
+            ),
+        ),
+    ],
+)
+def test_palette_invariants_with_image_bytes(
+    test_image_as_bytes, palette_size, extraction_mode
+):
+    palette = extract_colors(
+        image_bytes=test_image_as_bytes,
+        palette_size=palette_size,
+        mode=extraction_mode,
+        resize=True,
+    )
+
+    assert (
+        len(palette) == palette_size
+    ), f"Expected {palette_size} colors in palette, got {len(palette)}"
+    assert (
+        palette.number_of_colors == palette_size
+    ), f"Expected {palette_size} colors in palette, got {palette.number_of_colors}"
+    assert (
+        len(palette.colors) == palette_size
+    ), f"Expected {palette_size} colors in palette, got {len(palette.colors)}"
+    assert (
+        palette.colors[0].freq >= palette.colors[-1].freq
+    ), "Expected colors to be sorted by frequency in descending order"
+    assert (
+        palette.colors[0].freq > 0.0
+    ), "Expected the most frequent color to have a frequency greater than 0.0"
+    assert (
+        palette.colors[0].freq <= 1.0
+    ), "Expected the most frequent color to have a frequency less than or equal to 1.0"
+
+    assert_approx_equal(
+        sum(c.freq for c in palette.colors),
+        1.0,
+        err_msg="Expected the sum of all frequencies to be 1.0",
+    )
+
+
 def test_colorspace_invariants_hls(test_kmean_extracted_palette):
     for color in test_kmean_extracted_palette:
         h, l, s = color.get_colors(colorspace="hls")
@@ -92,3 +148,33 @@ def test_colorspace_invariants_rgb(test_kmean_extracted_palette):
         assert 0 <= r <= 255, f"Expected 0 <= r <= 255, got {r}"
         assert 0 <= g <= 255, f"Expected 0 <= g <= 255, got {g}"
         assert 0 <= b <= 255, f"Expected 0 <= b <= 255, got {b}"
+
+
+@pytest.mark.parametrize(
+    "resize, sort_mode", [(True, "luminance"), (False, "frequency")]
+)
+def test_color_extraction_deterministic_kmeans(
+    test_image_path_as_str, resize, sort_mode
+):
+    palette1 = extract_colors(
+        image=test_image_path_as_str,
+        palette_size=5,
+        mode="KM",
+        resize=resize,
+        sort_mode=sort_mode,
+    )
+    palette2 = extract_colors(
+        image=test_image_path_as_str,
+        palette_size=5,
+        mode="KM",
+        resize=resize,
+        sort_mode=sort_mode,
+    )
+    for c1, c2 in zip(palette1.colors, palette2.colors):
+        r, g, b, freq = c1.rgb[0], c1.rgb[1], c1.rgb[2], c1.freq
+        r2, g2, b2, freq2 = c2.rgb[0], c2.rgb[1], c2.rgb[2], c2.freq
+
+        assert r == r2, f"Expected r1 == r2, got {r} != {r2}"
+        assert g == g2, f"Expected g1 == g2, got {g} != {g2}"
+        assert b == b2, f"Expected b1 == b2, got {b} != {b2}"
+        assert freq == freq2, f"Expected freq1 == freq2, got {freq} != {freq2}"