Fixed palette index for zeroed color after FASTOCTREE quantize

python-pillow · Dec 5, 2021 · 695dd5f · 695dd5f
1 parent fa2bbf7
commit 695dd5f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/Tests/test_image_quantize.py b/Tests/test_image_quantize.py
@@ -85,3 +85,20 @@ def test_transparent_colors_equal():
     converted = im.quantize()
     converted_px = converted.load()
     assert converted_px[0, 0] == converted_px[0, 1]
+
+
+@pytest.mark.parametrize(
+    "method, color",
+    (
+        (Image.MEDIANCUT, (0, 0, 0)),
+        (Image.MAXCOVERAGE, (0, 0, 0)),
+        (Image.FASTOCTREE, (0, 0, 0)),
+        (Image.FASTOCTREE, (0, 0, 0, 0)),
+    ),
+)
+def test_palette(method, color):
+    im = Image.new("RGBA" if len(color) == 4 else "RGB", (1, 1), color)
+
+    converted = im.quantize(method=method)
+    converted_px = converted.load()
+    assert converted_px[0, 0] == converted.palette.colors[color]
diff --git a/src/PIL/Image.py b/src/PIL/Image.py
@@ -1112,6 +1112,13 @@ def quantize(self, colors=256, method=None, kmeans=0, palette=None, dither=1):
 
         mode = im.im.getpalettemode()
         im.palette = ImagePalette.ImagePalette(mode, im.im.getpalette(mode, mode))
+        if method == FASTOCTREE:
+            mode_len = len(mode)
+            color = (0,) * mode_len
+            if color in im.palette.colors and im.palette.palette.endswith(
+                b"\x00" * mode_len
+            ):
+                im.palette.colors[color] = 255
 
         return im