Adding PNG support and enabling experimental reader in kwargs (#654)

* Adding PNG support and enabling experimental reader in kwargs * PR comments reader arg * Make reader type an Enum
TileDB-Inc · Oct 9, 2024 · 3a8d3e6 · 3a8d3e6
1 parent 3e9ff2b
commit 3a8d3e6
Showing 1 changed file with 20 additions and 2 deletions.
diff --git a/src/tiledb/cloud/bioimg/ingestion.py b/src/tiledb/cloud/bioimg/ingestion.py
@@ -1,5 +1,6 @@
 import logging
 import warnings
+from enum import Enum
 from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple, Union
 
 import tiledb
@@ -15,8 +16,13 @@
 DEFAULT_RESOURCES = {"cpu": "8", "memory": "4Gi"}
 DEFAULT_IMG_NAME = "3.9-imaging-dev"
 DEFAULT_DAG_NAME = "bioimg-ingestion"
-_SUPPORTED_EXTENSIONS = (".tiff", ".tif", ".svs", ".ndpi")
-_SUPPORTED_CONVERTERS = ("tiff", "zarr", "osd")
+_SUPPORTED_EXTENSIONS = (".tiff", ".tif", ".svs", ".ndpi", ".png")
+_SUPPORTED_CONVERTERS = ("tiff", "zarr", "osd", "png")
+
+
+class ReaderType(Enum):
+    PRODUCTION = "production"
+    EXPERIMENTAL = "experimental"
 
 
 def ingest(
@@ -74,6 +80,8 @@ def ingest(
         registered in TileDB Cloud (ARN type) if ``acn`` is not set.
     :param dest_config: dict configuration to pass on tiledb.VFS for the destination's
         resolution
+    :param reader: The selected reader backend implementation either "experimental"
+        or "production". Default["production"]
     """
 
     logger = get_logger_wrapper(verbose)
@@ -203,8 +211,10 @@ def ingest_tiff_udf(
         user_converter = {
             "zarr": Converters.OMEZARR,
             "osd": Converters.OSD,
+            "png": Converters.PNG,
         }.get(converter, Converters.OMETIFF)
 
+        experimental_reader = kwargs.get("experimental_reader", False)
         compressor = kwargs.get("compressor", None)
         if compressor:
             compressor_args = dict(compressor)
@@ -229,8 +239,10 @@ def ingest_tiff_udf(
                 tile_scale=tile_scale,
                 source_config=config,
                 dest_config=kwargs.get("dest_config", None),
+                experimental_reader=experimental_reader,
                 **kwargs,
             )
+
         return io_uris
 
     def register_dataset_udf(
@@ -350,6 +362,11 @@ def register_dataset_udf(
 
     # serialize udf arguments
     compressor = kwargs.pop("compressor", None)
+
+    # Get either the new experimental or default reader
+    reader = kwargs.pop("reader", ReaderType.PRODUCTION.value)
+    experimental_reader = reader == ReaderType.EXPERIMENTAL.value
+
     logger.debug("Compressor: %r", compressor)
     compressor_serial = serialize_filter(compressor) if compressor else None
 
@@ -368,6 +385,7 @@ def register_dataset_udf(
         image_name=DEFAULT_IMG_NAME,
         max_workers=threads,
         compressor=compressor_serial,
+        experimental_reader=experimental_reader,
         access_credentials_name=acn,
         **kwargs,
     )