This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new da17cf0c71 [python] Decode palette image colors for training tensors
(#10096)
da17cf0c71 is described below
commit da17cf0c719918ff671e6a8ca125657dbe124456
Author: chaoyang <[email protected]>
AuthorDate: Thu Sep 24 14:38:07 2026 +0800
[python] Decode palette image colors for training tensors (#10096)
---
.../pypaimon/multimodal/window_transforms.py | 10 ++++--
.../pypaimon/tests/window_transforms_test.py | 41 ++++++++++++++++++++++
2 files changed, 48 insertions(+), 3 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/window_transforms.py
b/paimon-python/pypaimon/multimodal/window_transforms.py
index 3385b8525d..e8af830a3d 100644
--- a/paimon-python/pypaimon/multimodal/window_transforms.py
+++ b/paimon-python/pypaimon/multimodal/window_transforms.py
@@ -38,8 +38,9 @@ def to_tensor(values, dtype=None):
def images_to_tensor(values, return_uint8=False):
"""Decode a non-empty sequence of image bytes to a TCHW tensor.
- Apply EXIF orientation and preserve grayscale as one channel. Eight-bit
- pixels become float32 in [0, 1], or stay uint8 with ``return_uint8=True``.
+ Apply EXIF orientation, expand palettes to RGB/RGBA, and preserve grayscale
+ as one channel. Eight-bit pixels become float32 in [0, 1], or stay uint8
+ with ``return_uint8=True``.
Higher-bit-depth pixels always become float32 in their original units.
All decoded frames must have the same shape.
"""
@@ -70,7 +71,10 @@ def _decode_image(payload):
"or Pillow.") from error
with Image.open(io.BytesIO(payload)) as image:
- array = np.array(ImageOps.exif_transpose(image), copy=True)
+ image = ImageOps.exif_transpose(image)
+ if image.mode == "P":
+ image = image.convert("RGBA" if "transparency" in image.info else
"RGB")
+ array = np.array(image, copy=True)
if array.ndim == 2:
array = array[:, :, None]
return array
diff --git a/paimon-python/pypaimon/tests/window_transforms_test.py
b/paimon-python/pypaimon/tests/window_transforms_test.py
index 45df0766bb..7a74f38954 100644
--- a/paimon-python/pypaimon/tests/window_transforms_test.py
+++ b/paimon-python/pypaimon/tests/window_transforms_test.py
@@ -58,6 +58,47 @@ class WindowTransformsTest(unittest.TestCase):
gray = _png(np.array([[0, 255]], dtype=np.uint8))
self.assertEqual([[[[0., 1.]]]], images_to_tensor([gray]).tolist())
+ def test_palette_colors_and_transparency_are_expanded(self):
+ from pypaimon.multimodal.lerobot.dataset import _image_tensor
+
+ for transparency, alpha in (
+ (None, None), (0, [0, 255]), (bytes([255, 128]), [255, 128])):
+ with self.subTest(transparency=transparency):
+ image = Image.new("P", (2, 1))
+ image.putpalette([255, 0, 0, 0, 255, 0] + [0] * 762)
+ image.putdata([0, 1])
+ buffer = io.BytesIO()
+ options = {} if transparency is None else {"transparency":
transparency}
+ image.save(buffer, format="PNG", **options)
+ payload = buffer.getvalue()
+ channels = [[[255, 0]], [[0, 255]], [[0, 0]]]
+ if alpha is not None:
+ channels.append([alpha])
+ expected = torch.tensor([channels], dtype=torch.uint8)
+
+ torch.testing.assert_close(
+ images_to_tensor([payload], return_uint8=True), expected)
+ torch.testing.assert_close(
+ images_to_tensor([payload]), expected.float() / 255)
+ feature = {"dtype": "image", "shape": [1, 2, len(channels)]}
+ torch.testing.assert_close(
+ _image_tensor(payload, feature, return_uint8=True),
expected[0])
+
+ def test_palette_frames_stack_with_rgb_frames_after_exif_orientation(self):
+ image = Image.new("P", (2, 1))
+ image.putpalette([255, 0, 0, 0, 255, 0] + [0] * 762)
+ image.putdata([0, 1])
+ buffer = io.BytesIO()
+ exif = Image.Exif()
+ exif[274] = 6
+ image.save(buffer, format="PNG", exif=exif)
+ rgb = _png(np.array([[[255, 0, 0]], [[0, 255, 0]]], dtype=np.uint8))
+
+ result = images_to_tensor([buffer.getvalue(), rgb], return_uint8=True)
+
+ self.assertEqual((2, 3, 2, 1), tuple(result.shape))
+ torch.testing.assert_close(result[0], result[1])
+
def test_high_bit_depth_keeps_native_units(self):
payload = _png(np.array([[0, 1024, 65535]], dtype=np.uint16))
for return_uint8 in (False, True):
