PIL: YCbCR encoding support (#169)

* PIL: YCbCR encoding support

CHANGELOG.md

@@ -6,6 +6,7 @@ All notable changes to this project will be documented in this file.
 
 - Synonym for `chroma` encoder parameter: `subsampling`(usage is the same as in Pillow JPEG). #161 #165
 - Passing `image_orientation` value to libheif, instead of manually rotating image according to EXIF before encoding. #168
+- Pillow plugin: support for images in `YCbCr` mode for saving without converting to `RGB`. #169
 - Pi-Heif: Python3.12 32-bit `armv7` wheels. #160
 
 ### Changed

pillow_heif/_pillow_heif.c

@@ -414,11 +414,11 @@ static PyObject* _CtxWriteImage_add_plane_la(CtxWriteImageObject* self, PyObject
 
 static PyObject* _CtxWriteImage_add_plane_l(CtxWriteImageObject* self, PyObject* args) {
     /* (size), depth: int, depth_in: int, data: bytes */
-    int width, height, depth, depth_in, stride_out, stride_in, real_stride;
+    int width, height, depth, depth_in, stride_out, stride_in, real_stride, target_heif_channel;
     Py_buffer buffer;
     uint8_t *plane_data;
 
-    if (!PyArg_ParseTuple(args, "(ii)iiy*i", &width, &height, &depth, &depth_in, &buffer, &stride_in))
+    if (!PyArg_ParseTuple(args, "(ii)iiy*ii", &width, &height, &depth, &depth_in, &buffer, &stride_in, &target_heif_channel))
         return NULL;
 
     real_stride = width;
@@ -432,12 +432,12 @@ static PyObject* _CtxWriteImage_add_plane_l(CtxWriteImageObject* self, PyObject*
         return NULL;
     }
 
-    if (check_error(heif_image_add_plane(self->image, heif_channel_Y, width, height, depth))) {
+    if (check_error(heif_image_add_plane(self->image, target_heif_channel, width, height, depth))) {
         PyBuffer_Release(&buffer);
         return NULL;
     }
 
-    plane_data = heif_image_get_plane(self->image, heif_channel_Y, &stride_out);
+    plane_data = heif_image_get_plane(self->image, target_heif_channel, &stride_out);
     if (!plane_data) {
         PyBuffer_Release(&buffer);
         PyErr_SetString(PyExc_RuntimeError, "heif_image_get_plane(Y) failed");

pillow_heif/as_plugin.py

@@ -270,7 +270,7 @@ def __save_all(im, fp, compression_format: HeifCompressionFormat):
     ctx_write.save(fp)
 
 
-def _pil_encode_image(ctx: CtxEncode, img: Image.Image, primary: bool, **kwargs):
+def _pil_encode_image(ctx: CtxEncode, img: Image.Image, primary: bool, **kwargs) -> None:
     if img.size[0] <= 0 or img.size[1] <= 0:
         raise ValueError("Empty images are not supported.")
     _info = img.info.copy()
@@ -279,5 +279,10 @@ def _pil_encode_image(ctx: CtxEncode, img: Image.Image, primary: bool, **kwargs)
     if primary:
         _info.update(**kwargs)
     _info["primary"] = primary
-    _img = _pil_to_supported_mode(img)
-    ctx.add_image(_img.size, _img.mode, _img.tobytes(), image_orientation=_get_orientation_for_encoder(_info), **_info)
+    if img.mode == "YCbCr":
+        ctx.add_image_ycbcr(img, image_orientation=_get_orientation_for_encoder(_info), **_info)
+    else:
+        _img = _pil_to_supported_mode(img)
+        ctx.add_image(
+            _img.size, _img.mode, _img.tobytes(), image_orientation=_get_orientation_for_encoder(_info), **_info
+        )

pillow_heif/constants.py

@@ -176,3 +176,24 @@ class HeifDepthRepresentationType(IntEnum):
     """Unknown"""
     NON_UNIFORM_DISPARITY = 3
     """Unknown"""
+
+
+class HeifChannel(IntEnum):
+    """Internal libheif values, used in ``CtxEncode``."""
+
+    CHANNEL_Y = 0
+    """Monochrome or YCbCR"""
+    CHANNEL_CB = 1
+    """Only for YCbCR"""
+    CHANNEL_CR = 2
+    """Only for YCbCR"""
+    CHANNEL_R = 3
+    """RGB or RGBA"""
+    CHANNEL_G = 4
+    """RGB or RGBA"""
+    CHANNEL_B = 5
+    """RGB or RGBA"""
+    CHANNEL_ALPHA = 6
+    """Monochrome or RGBA"""
+    CHANNEL_INTERLEAVED = 10
+    """RGB or RGBA"""

pillow_heif/misc.py

@@ -15,7 +15,7 @@
 from PIL import Image
 
 from . import options
-from .constants import HeifChroma, HeifColorspace, HeifCompressionFormat
+from .constants import HeifChannel, HeifChroma, HeifColorspace, HeifCompressionFormat
 
 try:
     import _pillow_heif
@@ -69,6 +69,7 @@
     "LA": (2, 8, HeifColorspace.MONOCHROME, HeifChroma.MONOCHROME),
     "La": (2, 8, HeifColorspace.MONOCHROME, HeifChroma.MONOCHROME),
     "L": (1, 8, HeifColorspace.MONOCHROME, HeifChroma.MONOCHROME),
+    "YCbCr": (3, 8, HeifColorspace.YCBCR, HeifChroma.CHROMA_444),
 }
 
 SUBSAMPLING_CHROMA_MAP = {
@@ -265,6 +266,7 @@ def _xmp_from_pillow(img: Image.Image) -> Optional[bytes]:
 
 
 def _pil_to_supported_mode(img: Image.Image) -> Image.Image:
+    # We support "YCbCr" for encoding in Pillow plugin mode and do not call this function.
     if img.mode == "P":
         mode = "RGBA" if img.info.get("transparency") else "RGB"
         img = img.convert(mode=mode)
@@ -274,7 +276,7 @@ def _pil_to_supported_mode(img: Image.Image) -> Image.Image:
         img = img.convert(mode="L")
     elif img.mode == "CMYK":
         img = img.convert(mode="RGBA")
-    elif img.mode == "YCbCr":  # note: libheif supports native `YCbCr`.
+    elif img.mode == "YCbCr":
         img = img.convert(mode="RGB")
     return img
 
@@ -349,11 +351,23 @@ def add_image(self, size: tuple, mode: str, data, **kwargs) -> None:
         im_out = self.ctx_write.create_image(size, MODE_INFO[mode][2], MODE_INFO[mode][3], premultiplied_alpha)
         # image data
         if MODE_INFO[mode][0] == 1:
-            im_out.add_plane_l(size, bit_depth_out, bit_depth_in, data, kwargs.get("stride", 0))
+            im_out.add_plane_l(size, bit_depth_out, bit_depth_in, data, kwargs.get("stride", 0), HeifChannel.CHANNEL_Y)
         elif MODE_INFO[mode][0] == 2:
             im_out.add_plane_la(size, bit_depth_out, bit_depth_in, data, kwargs.get("stride", 0))
         else:
             im_out.add_plane(size, bit_depth_out, bit_depth_in, data, mode.find("BGR") != -1, kwargs.get("stride", 0))
+        self._finish_add_image(im_out, size, **kwargs)
+
+    def add_image_ycbcr(self, img: Image.Image, **kwargs) -> None:
+        """Adds image in `YCbCR` mode to the encoder."""
+        # creating image
+        im_out = self.ctx_write.create_image(img.size, MODE_INFO[img.mode][2], MODE_INFO[img.mode][3], 0)
+        # image data
+        for i in (HeifChannel.CHANNEL_Y, HeifChannel.CHANNEL_CB, HeifChannel.CHANNEL_CR):
+            im_out.add_plane_l(img.size, 8, 8, bytes(img.getdata(i)), kwargs.get("stride", 0), i)
+        self._finish_add_image(im_out, img.size, **kwargs)
+
+    def _finish_add_image(self, im_out, size: tuple, **kwargs):
         # color profile
         __icc_profile = kwargs.get("icc_profile", None)
         if __icc_profile is not None:

tests/write_test.py

@@ -268,14 +268,39 @@ def test_CMYK_color_mode():  # noqa
     helpers.compare_hashes([im, im_heif], hash_size=16)
 
 
-def test_YCbCr_color_mode():  # noqa
+@pytest.mark.parametrize("subsampling, expected_max_difference", (("4:4:4", 0.0004), ("4:2:2", 0.11), ("4:2:0", 1.33)))
+@pytest.mark.parametrize("save_format", ("HEIF", "AVIF"))
+def test_YCbCr_color_mode(
+    save_format,
+    subsampling,
+    expected_max_difference,
+):
+    im_original = helpers.gradient_rgb()
+    buf_jpeg = BytesIO()
+    im_original.save(buf_jpeg, format="JPEG", subsampling=subsampling, quality=-1)
+    im_jpeg = Image.open(buf_jpeg)
+    im_jpeg.draft("YCbCr", im_jpeg.size)
+    im_jpeg.load()
+    assert im_jpeg.mode == "YCbCr"
+    buf_heif = BytesIO()
+    im_jpeg.save(buf_heif, format=save_format, subsampling=subsampling, quality=-1)
+    im_heif = Image.open(buf_heif)
+    assert im_heif.mode == "RGB"
+    helpers.assert_image_similar(Image.open(buf_jpeg), im_heif, expected_max_difference)
+
+
+def test_heif_YCbCr_color_mode():  # noqa
+    # we support YCbCr for PIL only.
+    # in this test case, the image will be converted to "RGB" during "from_pillow".
     im = helpers.gradient_rgb().convert("YCbCr")
     assert im.mode == "YCbCr"
-    out_heif = BytesIO()
-    im.save(out_heif, format="HEIF", quality=-1)
-    im_heif = Image.open(out_heif)
+    im_heif = pillow_heif.from_pillow(im)
     assert im_heif.mode == "RGB"
-    helpers.compare_hashes([im, im_heif], hash_size=16)
+    out_heif = BytesIO()
+    im_heif.save(out_heif, format="HEIF", quality=-1)
+    im_out = Image.open(out_heif)
+    assert im_out.mode == "RGB"
+    helpers.compare_hashes([im, im_out], hash_size=32)
 
 
 @pytest.mark.parametrize("enc_bits", (10, 12))