Add support for parsing auxiliary images (#297)

added support for AUX images (hdrgainmap)

pillow_heif/__init__.py

@@ -16,6 +16,7 @@
     HeifTransferCharacteristics,
 )
 from .heif import (
+    HeifAuxImage,
     HeifDepthImage,
     HeifFile,
     HeifImage,

pillow_heif/_pillow_heif.c

@@ -731,6 +731,75 @@ static struct PyMethodDef _CtxWrite_methods[] = {
     {NULL, NULL}
 };
 
+/* =========== CtxAuxImage ======== */
+
+static const char* _colorspace_to_str(enum heif_colorspace colorspace) {
+    switch (colorspace) {
+        case heif_colorspace_undefined:
+            return "undefined";
+        case heif_colorspace_monochrome:
+            return "monochrome";
+        case heif_colorspace_RGB:
+            return "RGB";
+        case heif_colorspace_YCbCr:
+            return "YCbCr";
+        default:  // note: this means the upstream API has changed
+            return "unknown";
+    }
+}
+
+PyObject* _CtxAuxImage(struct heif_image_handle* main_handle, heif_item_id aux_image_id,
+                       int remove_stride, int hdr_to_16bit, PyObject* file_bytes) {
+    struct heif_image_handle* aux_handle;
+    if (check_error(heif_image_handle_get_auxiliary_image_handle(main_handle, aux_image_id, &aux_handle))) {
+        return NULL;
+    }
+    int luma_bits = heif_image_handle_get_luma_bits_per_pixel(aux_handle);
+    enum heif_colorspace colorspace;
+    enum heif_chroma chroma;
+    if (check_error(heif_image_handle_get_preferred_decoding_colorspace(aux_handle, &colorspace, &chroma))) {
+        heif_image_handle_release(aux_handle);
+        return NULL;
+    }
+    if (luma_bits != 8 || colorspace != heif_colorspace_monochrome) {
+        const char* colorspace_str = _colorspace_to_str(colorspace);
+        PyErr_Format(
+            PyExc_NotImplementedError,
+            "Only 8-bit monochrome auxiliary images are currently supported. Got %d-bit %s image. "
+            "Please consider filing an issue with an example HEIF file.",
+            luma_bits, colorspace_str);
+        heif_image_handle_release(aux_handle);
+        return NULL;
+    }
+    CtxImageObject *ctx_image = PyObject_New(CtxImageObject, &CtxImage_Type);
+    if (!ctx_image) {
+        heif_image_handle_release(aux_handle);
+        return NULL;
+    }
+    ctx_image->depth_metadata = NULL;
+    ctx_image->image_type = PhHeifImage;
+    ctx_image->width = heif_image_handle_get_width(aux_handle);
+    ctx_image->height = heif_image_handle_get_height(aux_handle);
+    ctx_image->alpha = 0;
+    ctx_image->n_channels = 1;
+    ctx_image->bits = 8;
+    strcpy(ctx_image->mode, "L");
+    ctx_image->hdr_to_8bit = 0;
+    ctx_image->bgr_mode = 0;
+    ctx_image->colorspace = heif_colorspace_monochrome;
+    ctx_image->chroma = heif_chroma_monochrome;
+    ctx_image->handle = aux_handle;
+    ctx_image->heif_image = NULL;
+    ctx_image->data = NULL;
+    ctx_image->remove_stride = remove_stride;
+    ctx_image->hdr_to_16bit = hdr_to_16bit;
+    ctx_image->reload_size = 1;
+    ctx_image->file_bytes = file_bytes;
+    ctx_image->stride = get_stride(ctx_image);
+    Py_INCREF(file_bytes);
+    return (PyObject*)ctx_image;
+}
+
 /* =========== CtxDepthImage ======== */
 
 PyObject* _CtxDepthImage(struct heif_image_handle* main_handle, heif_item_id depth_image_id,
@@ -1203,6 +1272,57 @@ static PyObject* _CtxImage_depth_image_list(CtxImageObject* self, void* closure)
     return images_list;
 }
 
+static PyObject* _CtxImage_aux_image_ids(CtxImageObject* self, void* closure) {
+    int aux_filter = LIBHEIF_AUX_IMAGE_FILTER_OMIT_ALPHA | LIBHEIF_AUX_IMAGE_FILTER_OMIT_DEPTH;
+    int n_images = heif_image_handle_get_number_of_auxiliary_images(self->handle, aux_filter);
+    if (n_images == 0)
+        return PyList_New(0);
+    heif_item_id* images_ids = (heif_item_id*)malloc(n_images * sizeof(heif_item_id));
+    if (!images_ids)
+        return PyErr_NoMemory();
+
+    n_images = heif_image_handle_get_list_of_auxiliary_image_IDs(self->handle, aux_filter, images_ids, n_images);
+    PyObject* images_list = PyList_New(n_images);
+    if (!images_list) {
+        free(images_ids);
+        return PyErr_NoMemory();
+    }
+    for (int i = 0; i < n_images; i++) {
+        PyList_SET_ITEM(images_list, i, PyLong_FromUnsignedLong(images_ids[i]));
+    }
+    free(images_ids);
+    return images_list;
+}
+
+static PyObject* _CtxImage_get_aux_image(CtxImageObject* self, PyObject* arg_image_id) {
+    heif_item_id aux_image_id = (heif_item_id)PyLong_AsUnsignedLong(arg_image_id);
+    return _CtxAuxImage(
+        self->handle, aux_image_id, self->remove_stride, self->hdr_to_16bit, self->file_bytes
+    );
+}
+
+static PyObject* _get_aux_type(const struct heif_image_handle* aux_handle) {
+    const char* aux_type_c = NULL;
+    struct heif_error error = heif_image_handle_get_auxiliary_type(aux_handle, &aux_type_c);
+    if (check_error(error))
+        return NULL;
+    PyObject *aux_type = PyUnicode_FromString(aux_type_c);
+    heif_image_handle_release_auxiliary_type(aux_handle, &aux_type_c);
+    return aux_type;
+}
+
+static PyObject* _CtxImage_get_aux_type(CtxImageObject* self, PyObject* arg_image_id) {
+    heif_item_id aux_image_id = (heif_item_id)PyLong_AsUnsignedLong(arg_image_id);
+    struct heif_image_handle* aux_handle;
+    if (check_error(heif_image_handle_get_auxiliary_image_handle(self->handle, aux_image_id, &aux_handle)))
+        return NULL;
+    PyObject* aux_type = _get_aux_type(aux_handle);
+    if (!aux_type)
+        return NULL;
+    heif_image_handle_release(aux_handle);
+    return aux_type;
+}
+
 /* =========== CtxImage Experimental Part ======== */
 
 static PyObject* _CtxImage_camera_intrinsic_matrix(CtxImageObject* self, void* closure) {
@@ -1265,11 +1385,18 @@ static struct PyGetSetDef _CtxImage_getseters[] = {
     {"stride", (getter)_CtxImage_stride, NULL, NULL, NULL},
     {"data", (getter)_CtxImage_data, NULL, NULL, NULL},
     {"depth_image_list", (getter)_CtxImage_depth_image_list, NULL, NULL, NULL},
+    {"aux_image_ids", (getter)_CtxImage_aux_image_ids, NULL, NULL, NULL},
     {"camera_intrinsic_matrix", (getter)_CtxImage_camera_intrinsic_matrix, NULL, NULL, NULL},
     {"camera_extrinsic_matrix_rot", (getter)_CtxImage_camera_extrinsic_matrix_rot, NULL, NULL, NULL},
     {NULL, NULL, NULL, NULL, NULL}
 };
 
+static struct PyMethodDef _CtxImage_methods[] = {
+    {"get_aux_image", (PyCFunction)_CtxImage_get_aux_image, METH_O},
+    {"get_aux_type", (PyCFunction)_CtxImage_get_aux_type, METH_O},
+    {NULL, NULL}
+};
+
 /* =========== Functions ======== */
 
 static PyObject* _CtxWrite(PyObject* self, PyObject* args) {
@@ -1517,6 +1644,7 @@ static PyTypeObject CtxImage_Type = {
     .tp_dealloc = (destructor)_CtxImage_destructor,
     .tp_flags = Py_TPFLAGS_DEFAULT,
     .tp_getset = _CtxImage_getseters,
+    .tp_methods = _CtxImage_methods,
 };
 
 static int setup_module(PyObject* m) {

pillow_heif/as_plugin.py

@@ -226,6 +226,8 @@ def __options_update(**kwargs):
             options.THUMBNAILS = v
         elif k == "depth_images":
             options.DEPTH_IMAGES = v
+        elif k == "aux_images":
+            options.AUX_IMAGES = v
         elif k == "quality":
             options.QUALITY = v
         elif k == "save_to_12bit":

pillow_heif/heif.py

@@ -38,7 +38,7 @@
 
 
 class BaseImage:
-    """Base class for :py:class:`HeifImage` and :py:class:`HeifDepthImage`."""
+    """Base class for :py:class:`HeifImage`, :py:class:`HeifDepthImage` and :py:class:`HeifAuxImage`."""
 
     size: tuple[int, int]
     """Width and height of the image."""
@@ -127,7 +127,6 @@ def __init__(self, c_image):
         save_colorspace_chroma(c_image, self.info)
 
     def __repr__(self):
-        _bytes = f"{len(self.data)} bytes" if self._data or isinstance(self._c_image, MimCImage) else "no"
         return f"<{self.__class__.__name__} {self.size[0]}x{self.size[1]} {self.mode}>"
 
     def to_pillow(self) -> Image.Image:
@@ -140,6 +139,13 @@ def to_pillow(self) -> Image.Image:
         return image
 
 
+class HeifAuxImage(BaseImage):
+    """Class representing the auxiliary image associated with the :py:class:`~pillow_heif.HeifImage` class."""
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {self.size[0]}x{self.size[1]} {self.mode}>"
+
+
 class HeifImage(BaseImage):
     """One image in a :py:class:`~pillow_heif.HeifFile` container."""
 
@@ -152,7 +158,6 @@ def __init__(self, c_image):
         _depth_images: list[HeifDepthImage | None] = (
             [HeifDepthImage(i) for i in c_image.depth_image_list if i is not None] if options.DEPTH_IMAGES else []
         )
-        _heif_meta = _get_heif_meta(c_image)
         self.info = {
             "primary": bool(c_image.primary),
             "bit_depth": int(c_image.bit_depth),
@@ -161,6 +166,15 @@ def __init__(self, c_image):
             "thumbnails": _thumbnails,
             "depth_images": _depth_images,
         }
+        if options.AUX_IMAGES:
+            _ctx_aux_info = {}
+            for aux_id in c_image.aux_image_ids:
+                aux_type = c_image.get_aux_type(aux_id)
+                if aux_type not in _ctx_aux_info:
+                    _ctx_aux_info[aux_type] = []
+                _ctx_aux_info[aux_type].append(aux_id)
+            self.info["aux"] = _ctx_aux_info
+        _heif_meta = _get_heif_meta(c_image)
         if _xmp:
             self.info["xmp"] = _xmp
         if _heif_meta:
@@ -206,6 +220,14 @@ def to_pillow(self) -> Image.Image:
         image.info["original_orientation"] = set_orientation(image.info)
         return image
 
+    def get_aux_image(self, aux_id: int) -> HeifAuxImage:
+        """Method to retrieve the auxiliary image at the given ID.
+
+        :returns: a :py:class:`~pillow_heif.HeifAuxImage` class instance.
+        """
+        aux_image = self._c_image.get_aux_image(aux_id)
+        return HeifAuxImage(aux_image)
+
 
 class HeifFile:
     """Representation of the :py:class:`~pillow_heif.HeifImage` classes container.
@@ -481,6 +503,13 @@ def __copy(self):
         _im_copy.primary_index = self.primary_index
         return _im_copy
 
+    def get_aux_image(self, aux_id):
+        """`get_aux_image`` method of the primary :class:`~pillow_heif.HeifImage` in the container.
+
+        :exception IndexError: If there are no images.
+        """
+        return self._images[self.primary_index].get_aux_image(aux_id)
+
     __copy__ = __copy
 
 

pillow_heif/misc.py

@@ -478,6 +478,7 @@ def __init__(self, mode: str, size: tuple[int, int], data: bytes, **kwargs):
         self.color_profile = None
         self.thumbnails: list[int] = []
         self.depth_image_list: list = []
+        self.aux_image_ids: list[int] = []
         self.primary = False
         self.chroma = HeifChroma.UNDEFINED.value
         self.colorspace = HeifColorspace.UNDEFINED.value

pillow_heif/options.py

@@ -18,6 +18,12 @@
 When use pillow_heif as a plugin you can set it with: `register_*_opener(depth_images=False)`"""
 
 
+AUX_IMAGES = True
+"""Option to enable/disable auxiliary image support
+
+When use pillow_heif as a plugin you can set it with: `register_*_opener(aux_images=False)`"""
+
+
 QUALITY = None
 """Default encoding quality
 

tests/options_test.py

@@ -32,6 +32,7 @@ def test_options_change_from_plugin_registering(register_opener):
             save_to_12bit=True,
             decode_threads=3,
             depth_images=False,
+            aux_images=False,
             save_nclx_profile=False,
             preferred_encoder={"HEIF": "id1", "AVIF": "id2"},
             preferred_decoder={"HEIF": "id3", "AVIF": "id4"},
@@ -41,6 +42,7 @@ def test_options_change_from_plugin_registering(register_opener):
         assert options.SAVE_HDR_TO_12_BIT
         assert options.DECODE_THREADS == 3
         assert options.DEPTH_IMAGES is False
+        assert options.AUX_IMAGES is False
         assert options.SAVE_NCLX_PROFILE is False
         assert options.PREFERRED_ENCODER == {"HEIF": "id1", "AVIF": "id2"}
         assert options.PREFERRED_DECODER == {"HEIF": "id3", "AVIF": "id4"}
@@ -50,6 +52,7 @@ def test_options_change_from_plugin_registering(register_opener):
         options.SAVE_HDR_TO_12_BIT = False
         options.DECODE_THREADS = 4
         options.DEPTH_IMAGES = True
+        options.AUX_IMAGES = True
         options.SAVE_NCLX_PROFILE = True
         options.PREFERRED_ENCODER = {"HEIF": "", "AVIF": ""}
         options.PREFERRED_DECODER = {"HEIF": "", "AVIF": ""}

tests/read_test.py

@@ -496,6 +496,19 @@ def test_depth_image():
     assert im_pil.info == depth_image.info
 
 
+def test_aux_image():
+    im = pillow_heif.open_heif("images/heif_other/pug.heic")
+    assert len(im.info["aux"]) == 1
+    assert "urn:com:apple:photo:2020:aux:hdrgainmap" in im.info["aux"]
+    assert len(im.info["aux"]["urn:com:apple:photo:2020:aux:hdrgainmap"]) == 1
+    aux_id = im.info["aux"]["urn:com:apple:photo:2020:aux:hdrgainmap"][0]
+    aux_image = im.get_aux_image(aux_id)
+    assert isinstance(aux_image, pillow_heif.HeifAuxImage)
+    aux_pil = aux_image.to_pillow()
+    assert aux_pil.size == (2016, 1512)
+    assert aux_pil.mode == "L"
+
+
 @pytest.mark.skipif(
     parse_version(pillow_heif.libheif_version()) < parse_version("1.18.0"), reason="requires LibHeif 1.18+"
 )