Extra pruning of Structured Annotations (#151)

* Extra pruning of Structured Annotations

* Adding asserts

* Expand test cases

* Typos

* Modify exclude_metadata argument to be a Callable for custom metadata pruning

* PR comments

* Adding a minimal test

* Adding minimal test

* Better callback exception handling

Author: ktsitsi

tests/__init__.py

@@ -115,7 +115,9 @@ def generate_test_case(num_axes, num_ranges, max_value):
     return input_ranges, expected_output
 
 
-def generate_xml(has_macro=True, has_label=True, root_tag="OME", num_images=1):
+def generate_xml(
+    has_macro=True, has_label=True, has_annotations=True, root_tag="OME", num_images=1
+):
     """Generate synthetic XML strings with options to include 'macro' and 'label' images."""
 
     # Create the root element
@@ -204,5 +206,37 @@ def generate_xml(has_macro=True, has_label=True, root_tag="OME", num_images=1):
         )
         tiffdata = ET.SubElement(pixels, "TiffData", IFD="2", PlaneCount="1")
 
+    if has_annotations:
+        annotations = ET.SubElement(ome, "StructuredAnnotations")
+        if has_macro:
+            macro_annot = ET.SubElement(
+                annotations,
+                "CommentAnnotation",
+                ID=f"Annotation:{random.randint}",
+                Namespace="",
+            )
+            description = ET.SubElement(macro_annot, "Description")
+
+            # Based on standard
+            description.text = "barcode_value"
+            value = ET.SubElement(macro_annot, "Value")
+            value.text = "random_text"
+        if has_label:
+            label_annot = ET.SubElement(
+                annotations,
+                "CommentAnnotation",
+                ID=f"Annotation:{random.randint}",
+                Namespace="",
+            )
+            description = ET.SubElement(
+                label_annot,
+                "Description",
+            )
+            value = ET.SubElement(label_annot, "Value")
+
+            # Based on standard
+            description.text = "label_text"
+            value.text = "random_text"
+
     # Convert the ElementTree to a string
     return ET.tostring(ome, encoding="unicode")

tests/integration/converters/test_ome_tiff.py

@@ -128,6 +128,31 @@ def test_ome_tiff_converter_exclude_original_metadata(
     with TileDBOpenSlide(str(tiledb_path)) as t:
         assert t.properties["original_metadata"] == "{}"
 
+    # Exclude with a custom function applies only to images with OME-XML metadata
+    if filename != "UTM2GTIF.tiff":
+
+        def custom_pruning_function(xml_string: str):
+            return "custom_metadata"
+
+        input_path = get_path(filename)
+        tiledb_path = tmp_path / "to_tiledb_custom_metadata"
+        OMETiffConverter.to_tiledb(
+            input_path,
+            str(tiledb_path),
+            preserve_axes=preserve_axes,
+            chunked=chunked,
+            max_workers=max_workers,
+            compressor=compressor,
+            log=False,
+            exclude_metadata=custom_pruning_function,
+        )
+
+        with TileDBOpenSlide(str(tiledb_path)) as t:
+            assert (
+                t.properties["original_metadata"]
+                == '{"ome_metadata": "custom_metadata"}'
+            )
+
 
 @pytest.mark.parametrize(
     "filename,num_series", [("CMU-1-Small-Region.ome.tiff", 3), ("UTM2GTIF.tiff", 1)]

tests/unit/test_helpers.py

@@ -63,14 +63,25 @@ def test_validate_ingestion(num_axes, num_ranges, max_value):
 
 @pytest.mark.parametrize("macro", [True, False])
 @pytest.mark.parametrize("has_label", [True, False])
+@pytest.mark.parametrize("annotations", [True, False])
 @pytest.mark.parametrize("num_images", [1, 2, 3])
 @pytest.mark.parametrize("root_tag", ["OME", "InvalidRoot"])
-def test_remove_ome_image_metadata(macro, has_label, num_images, root_tag):
+def test_remove_ome_image_metadata(macro, has_label, annotations, num_images, root_tag):
     original_xml_string = generate_xml(
-        has_macro=macro, has_label=has_label, num_images=1, root_tag=root_tag
+        has_macro=macro,
+        has_label=has_label,
+        has_annotations=annotations,
+        num_images=1,
+        root_tag=root_tag,
     )
 
     excluded_metadata = remove_ome_image_metadata(original_xml_string)
+
+    namespaces = {"ome": "http://www.openmicroscopy.org/Schemas/OME/2016-06"}
+
+    barcode_xpath = ".//ome:StructuredAnnotations/ome:CommentAnnotation[ome:Description='barcode_value']"
+    label_xpath = ".//ome:StructuredAnnotations/ome:CommentAnnotation[ome:Description='label_text']"
+
     if root_tag == "OME":
         parsed_excluded = ET.fromstring(excluded_metadata)
 
@@ -89,5 +100,10 @@ def test_remove_ome_image_metadata(macro, has_label, num_images, root_tag):
             )
             is None
         )
+
+        # Assert if "barcode_value" and "label_text" subelement is present
+        assert parsed_excluded.find(barcode_xpath, namespaces=namespaces) is None
+        assert parsed_excluded.find(label_xpath, namespaces=namespaces) is None
+
     else:
         assert remove_ome_image_metadata(original_xml_string) is None

tiledb/bioimg/converters/base.py

@@ -9,6 +9,7 @@
 from operator import itemgetter
 from typing import (
     Any,
+    Callable,
     Dict,
     Generic,
     Iterator,
@@ -41,6 +42,7 @@
 
 from .. import ATTR_NAME
 from ..helpers import (
+    MetadataCallbackError,
     ReadWriteGroup,
     compute_channel_minmax,
     get_axes_mapper,
@@ -324,7 +326,7 @@ def to_tiledb(
         preserve_axes: bool = False,
         chunked: bool = False,
         max_workers: int = 0,
-        exclude_metadata: bool = False,
+        exclude_metadata: Union[bool, Callable[[str], str], None] = None,
         experimental_reader: bool = False,
         experimental_queue_limit: Tuple[int, int] = (10, 20),
         compressor: Optional[Union[Mapping[int, Any], Any]] = None,
@@ -350,7 +352,14 @@ def to_tiledb(
                     original ones.
                 :param max_workers: Maximum number of threads that can be used for conversion.
                     Applicable only if chunked=True.
-                :param exclude_metadata: If true, drop original metadata of the images and exclude them from being ingested.
+                :param exclude_metadata: An optional argument that specifies how to transform the original metadata.
+                    It can be one of the following:
+                    *   A callable (function, method, etc.) that takes an OME-XML string and returns it as a string, while removing
+                        some of the original metadata and excluding them from being ingested.
+                    *   A boolean value:
+                        *   ``True``: Indicates a specific built-in transformation should be applied. see: `remove_ome_image_metadata`
+                        *   ``False``: Indicates no transformation should be applied.
+                    *   ``None``: Indicates no transformation should be applied (same as ``False``).
                 :param experimental_reader: If true, use the experimental tiff reader optimized for s3 reads.
                     Experimental feature, use with caution
                 :param experimental_queue_limit: When using the experimental reader, define the minimum and maximum number of
@@ -536,7 +545,15 @@ def to_tiledb(
                 original_metadata = reader.original_metadata
             else:
                 if ome_xml := reader.original_metadata.get("ome_metadata"):
-                    pruned_metadata = remove_ome_image_metadata(ome_xml)
+                    if isinstance(exclude_metadata, bool):
+                        pruned_metadata = remove_ome_image_metadata(ome_xml)
+                    elif callable(exclude_metadata):
+                        try:
+                            pruned_metadata = exclude_metadata(ome_xml)
+                        except Exception as exc:
+                            raise MetadataCallbackError(str(exc))
+                    else:
+                        raise TypeError("exclude_metadata must be bool or callable")
                     original_metadata = (
                         {"ome_metadata": pruned_metadata} if pruned_metadata else {}
                     )

tiledb/bioimg/helpers.py

@@ -504,13 +504,27 @@ def remove_ome_image_metadata(xml_string: str) -> Union[str, Any]:
 
     # Find all images
     images = root.findall("ome:Image", ns)
+    structured_annotations = root.findall("ome:StructuredAnnotations", ns)
 
     # Iterate over images and remove those with Name 'macro' or 'label'
     for image in images:
         name = image.attrib.get("Name")
         if name in ["macro", "label"]:
             root.remove(image)
 
+    # Iterate over structured annotations and remove those with name `macro-text` and `label`
+    for sa in structured_annotations:
+        comment_annotations = sa.findall("ome:CommentAnnotation", ns)
+        comments_left = len(comment_annotations)
+        for comment_annot in sa.findall("ome:CommentAnnotation", ns):
+            desc = comment_annot.find("ome:Description", ns)
+            if desc is not None:
+                if "barcode" in str(desc.text) or "label" in str(desc.text):
+                    sa.remove(comment_annot)
+                    comments_left -= 1
+        if comments_left == 0:
+            root.remove(sa)
+
     # Return the modified XML as a string
     # Regular expression pattern to match 'ns0', 'ns0:', or ':ns0'
     pattern = r"ns0:|:ns0|ns0"
@@ -524,3 +538,16 @@ def remove_ome_image_metadata(xml_string: str) -> Union[str, Any]:
             encoding="unicode",
         ),
     )
+
+
+class MetadataCallbackError(Exception):
+    """
+    Custom exception class for specific error conditions.
+    """
+
+    def __init__(self, message: str = ""):
+        self.message = f"OME XML callback function failed {message}"
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        return f"MetadataCallbackError: {self.message}"

tiledb/bioimg/wrappers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import importlib.util
-from typing import Any, Optional
+from typing import Any, Callable, Optional, Union
 
 try:
     importlib.util.find_spec("tifffile")
@@ -30,7 +30,7 @@ def from_bioimg(
     converter: Converters = Converters.OMETIFF,
     *,
     verbose: bool = False,
-    exclude_metadata: bool = False,
+    exclude_metadata: Union[bool, Callable[[str], str], None] = None,
     tile_scale: int = 1,
     **kwargs: Any,
 ) -> Any:
@@ -41,6 +41,14 @@ def from_bioimg(
     :param dest: The destination path where the TileDB image will be stored
     :param converter: The converter type to be used (tentative) soon automatically detected
     :param verbose: verbose logging, defaults to False
+    :param exclude_metadata: An optional argument that specifies how to transform the original metadata.
+        It can be one of the following:
+        *   A callable (function, method, etc.) that takes an OME-XML string and returns it as a string, while removing
+            some of the original metadata and excluding them from being ingested.
+        *   A boolean value:
+            *   ``True``: Indicates a specific built-in transformation should be applied. see: `remove_ome_image_metadata`
+            *   ``False``: Indicates no transformation should be applied.
+        *   ``None``: Indicates no transformation should be applied (same as ``False``).
     :param kwargs: keyword arguments for custom ingestion behaviour
     :return: The converter class that was used for the ingestion
     """