From 45a819ebc80359b22efdf7918f9d65d79a4646aa Mon Sep 17 00:00:00 2001
From: Leonid Kostrykin <void@evoid.de>
Date: Wed, 20 Mar 2024 09:35:50 +0100
Subject: [PATCH] Adds tifffile support for image assertions

---
 lib/galaxy/tool_util/verify/__init__.py      | 11 ++-
 lib/galaxy/tool_util/verify/asserts/image.py | 88 +++++++++++---------
 test/functional/tools/validation_image.xml   | 12 +++
 3 files changed, 70 insertions(+), 41 deletions(-)

diff --git a/lib/galaxy/tool_util/verify/__init__.py b/lib/galaxy/tool_util/verify/__init__.py
index da654645d1b9..dd40e8cacf62 100644
--- a/lib/galaxy/tool_util/verify/__init__.py
+++ b/lib/galaxy/tool_util/verify/__init__.py
@@ -501,14 +501,19 @@ def get_image_metric(
 
 
 def _load_image(filepath: str) -> "numpy.typing.NDArray":
+    """
+    Reads the given image, trying tifffile and Pillow for reading.
+    """
+    # Try reading with tifffile first. It fails if the file is not a TIFF.
     try:
-        # Try reading with tifffile first. It fails if the file is not a TIFF.
         arr = tifffile.imread(filepath)
-    except tifffile.TiffFileError:
 
-        # If tifffile failed, then the file is not a tifffile. In that case, try with Pillow.
+    # If tifffile failed, then the file is not a tifffile. In that case, try with Pillow.
+    except tifffile.TiffFileError:
         with Image.open(filepath) as im:
             arr = numpy.array(im)
+
+    # Return loaded image
     return arr
 
 
diff --git a/lib/galaxy/tool_util/verify/asserts/image.py b/lib/galaxy/tool_util/verify/asserts/image.py
index f0fbc6e8326e..7513be700f5e 100644
--- a/lib/galaxy/tool_util/verify/asserts/image.py
+++ b/lib/galaxy/tool_util/verify/asserts/image.py
@@ -18,6 +18,10 @@
     from PIL import Image
 except ImportError:
     pass
+try:
+    import tifffile
+except ImportError:
+    pass
 
 if TYPE_CHECKING:
     import numpy.typing
@@ -58,18 +62,17 @@ def assert_has_image_width(
     """
     Asserts the specified output is an image and has a width of the specified value.
     """
-    buf = io.BytesIO(output_bytes)
-    with Image.open(buf) as im:
-        _assert_number(
-            im.size[0],
-            width,
-            delta,
-            min,
-            max,
-            negate,
-            "{expected} width {n}+-{delta}",
-            "{expected} width to be in [{min}:{max}]",
-        )
+    im_arr = _get_image(output_bytes)
+    _assert_number(
+        im_arr.shape[1],
+        width,
+        delta,
+        min,
+        max,
+        negate,
+        "{expected} width {n}+-{delta}",
+        "{expected} width to be in [{min}:{max}]",
+    )
 
 
 def assert_has_image_height(
@@ -83,18 +86,17 @@ def assert_has_image_height(
     """
     Asserts the specified output is an image and has a height of the specified value.
     """
-    buf = io.BytesIO(output_bytes)
-    with Image.open(buf) as im:
-        _assert_number(
-            im.size[1],
-            height,
-            delta,
-            min,
-            max,
-            negate,
-            "{expected} height {n}+-{delta}",
-            "{expected} height to be in [{min}:{max}]",
-        )
+    im_arr = _get_image(output_bytes)
+    _assert_number(
+        im_arr.shape[0],
+        height,
+        delta,
+        min,
+        max,
+        negate,
+        "{expected} height {n}+-{delta}",
+        "{expected} height to be in [{min}:{max}]",
+    )
 
 
 def assert_has_image_channels(
@@ -108,18 +110,18 @@ def assert_has_image_channels(
     """
     Asserts the specified output is an image and has the specified number of channels.
     """
-    buf = io.BytesIO(output_bytes)
-    with Image.open(buf) as im:
-        _assert_number(
-            len(im.getbands()),
-            channels,
-            delta,
-            min,
-            max,
-            negate,
-            "{expected} image channels {n}+-{delta}",
-            "{expected} image channels to be in [{min}:{max}]",
-        )
+    im_arr = _get_image(output_bytes)
+    n_channels = 1 if im_arr.ndim < 3 else im_arr.shape[2]  # we assume here that the image is a 2-D image
+    _assert_number(
+        n_channels,
+        channels,
+        delta,
+        min,
+        max,
+        negate,
+        "{expected} image channels {n}+-{delta}",
+        "{expected} image channels to be in [{min}:{max}]",
+    )
 
 
 def _compute_center_of_mass(im_arr: "numpy.typing.NDArray") -> Tuple[float, float]:
@@ -139,10 +141,20 @@ def _get_image(
 ) -> "numpy.typing.NDArray":
     """
     Returns the output image or a specific channel.
+
+    The function tries to read the image using tifffile and Pillow.
     """
     buf = io.BytesIO(output_bytes)
-    with Image.open(buf) as im:
-        im_arr = numpy.array(im)
+
+    # Try reading with tifffile first. It fails if the file is not a TIFF.
+    try:
+        im_arr = tifffile.imread(buf)
+
+    # If tifffile failed, then the file is not a tifffile. In that case, try with Pillow.
+    except tifffile.TiffFileError:
+        buf.seek(0)
+        with Image.open(buf) as im:
+            im_arr = numpy.array(im)
 
     # Select the specified channel (if any).
     if channel is not None:
diff --git a/test/functional/tools/validation_image.xml b/test/functional/tools/validation_image.xml
index a1a7791fa6ec..a3c93ea84f67 100644
--- a/test/functional/tools/validation_image.xml
+++ b/test/functional/tools/validation_image.xml
@@ -93,6 +93,18 @@
                 </assert_contents>
             </output>
         </test>
+        <!-- Tests with float TIFF images -->
+        <test>
+            <param name="input" value="im4_float.tif" />
+            <output name="output">
+                <assert_contents>
+                    <has_image_width width="25" />
+                    <has_image_height height="25" />
+                    <has_image_channels channels="1" />
+                    <has_image_center_of_mass center_of_mass="11.75, 11.75" />
+                </assert_contents>
+            </output>
+        </test>
         <!-- Tests with label images -->
         <test>
             <param name="input" value="im2_b.png" />