From 83aa8e074425c027ddd8bbb9245db6de61c1e1a6 Mon Sep 17 00:00:00 2001 From: John Charankattu Date: Tue, 15 Oct 2024 13:18:07 +0530 Subject: [PATCH] Add support for parsing auxiliary images (#297) added support for AUX images (hdrgainmap) --- pillow_heif/__init__.py | 1 + pillow_heif/_pillow_heif.c | 128 +++++++++++++++++++++++++++++++++++++ pillow_heif/as_plugin.py | 2 + pillow_heif/heif.py | 35 +++++++++- pillow_heif/misc.py | 1 + pillow_heif/options.py | 6 ++ tests/options_test.py | 3 + tests/read_test.py | 13 ++++ 8 files changed, 186 insertions(+), 3 deletions(-) diff --git a/pillow_heif/__init__.py b/pillow_heif/__init__.py index 4fd5be4a..41f814af 100644 --- a/pillow_heif/__init__.py +++ b/pillow_heif/__init__.py @@ -16,6 +16,7 @@ HeifTransferCharacteristics, ) from .heif import ( + HeifAuxImage, HeifDepthImage, HeifFile, HeifImage, diff --git a/pillow_heif/_pillow_heif.c b/pillow_heif/_pillow_heif.c index 82067a24..f5722f13 100644 --- a/pillow_heif/_pillow_heif.c +++ b/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) { diff --git a/pillow_heif/as_plugin.py b/pillow_heif/as_plugin.py index fbe3786a..a1c85409 100644 --- a/pillow_heif/as_plugin.py +++ b/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": diff --git a/pillow_heif/heif.py b/pillow_heif/heif.py index b9df9adc..ddecd1b4 100644 --- a/pillow_heif/heif.py +++ b/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 diff --git a/pillow_heif/misc.py b/pillow_heif/misc.py index 375ced85..08c01e73 100644 --- a/pillow_heif/misc.py +++ b/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 diff --git a/pillow_heif/options.py b/pillow_heif/options.py index 8c023822..0bbd90b5 100644 --- a/pillow_heif/options.py +++ b/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 diff --git a/tests/options_test.py b/tests/options_test.py index 670c6251..cbfbade3 100644 --- a/tests/options_test.py +++ b/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": ""} diff --git a/tests/read_test.py b/tests/read_test.py index bcbf4ca8..182b3a4e 100644 --- a/tests/read_test.py +++ b/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+" )