feat: support multispectral and float32 TIFF images via tifffile

labelme/_label_file.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 import PIL.Image
+import tifffile
 from loguru import logger
 from numpy.typing import NDArray
 
@@ -141,7 +142,7 @@ def __init__(self, filename=None):
     @staticmethod
     def load_image_file(filename):
         t0 = time.time()
-        image_pil = PIL.Image.open(filename)
+        image_pil = _imread(filename=filename)
 
         oriented: PIL.Image.Image = utils.apply_exif_orientation(image_pil)
         ext = osp.splitext(filename)[1].lower()
@@ -270,3 +271,48 @@ def save(
     @staticmethod
     def is_label_file(filename):
         return osp.splitext(filename)[1].lower() == LabelFile.suffix
+
+
+_DISPLAYABLE_MODES = {"1", "L", "P", "RGB", "RGBA", "LA", "PA"}
+
+
+def _imread(filename: str) -> PIL.Image.Image:
+    ext: str = osp.splitext(filename)[1].lower()
+    try:
+        image_pil = PIL.Image.open(filename)
+        if image_pil.mode not in _DISPLAYABLE_MODES:
+            raise PIL.UnidentifiedImageError
+        return image_pil
+    except PIL.UnidentifiedImageError:
+        if ext in (".tif", ".tiff"):
+            return _imread_tiff(filename)
+        raise
+
+
+def _imread_tiff(filename: str) -> PIL.Image.Image:
+    img_arr: NDArray = tifffile.imread(filename)
+
+    if img_arr.ndim == 2:
+        img_arr_normalized = _normalize_to_uint8(img_arr)
+    elif img_arr.ndim == 3:
+        if img_arr.shape[2] >= 3:
+            img_arr_normalized = np.stack(
+                [_normalize_to_uint8(img_arr[:, :, i]) for i in range(3)],
+                axis=2,
+            )
+        else:
+            img_arr_normalized = _normalize_to_uint8(img_arr[:, :, 0])
+    else:
+        raise OSError(f"Unsupported image shape: {img_arr.shape}")
+
+    return PIL.Image.fromarray(img_arr_normalized)
+
+
+def _normalize_to_uint8(arr: NDArray) -> NDArray[np.uint8]:
+    arr = arr.astype(np.float64)
+    min_val = np.nanmin(arr)
+    max_val = np.nanmax(arr)
+    if np.isnan(min_val) or np.isnan(max_val) or max_val - min_val == 0:
+        return np.zeros(arr.shape, dtype=np.uint8)
+    normalized = (arr - min_val) / (max_val - min_val) * 255
+    return np.clip(normalized, 0, 255).astype(np.uint8)

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
   "pyqt5-qt5!=5.15.11,!=5.15.12,!=5.15.13,!=5.15.14,!=5.15.15,!=5.15.16 ; sys_platform == 'win32'",
   "pyyaml",
   "scikit-image",
+  "tifffile",
 ]
 dynamic = ["readme", "version"]
 

tests/unit/load_image_file_test.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import io
 from pathlib import Path
 
 import numpy as np
 import PIL.Image
+import tifffile
 
 from labelme._label_file import LabelFile
 
@@ -38,3 +40,47 @@ def test_png_returns_raw_bytes(tmp_path):
     path = _make_image(tmp_path, "test.png")
     data = LabelFile.load_image_file(str(path))
     assert data == path.read_bytes()
+
+
+def test_multispectral_tiff_float32(tmp_path):
+    arr = np.random.rand(64, 64, 5).astype(np.float32) * 0.5
+    path = tmp_path / "multispectral.tif"
+    tifffile.imwrite(str(path), arr)
+
+    data = LabelFile.load_image_file(str(path))
+    assert data[:2] == b"\xff\xd8"
+
+    img = PIL.Image.open(io.BytesIO(data))
+    assert img.mode == "RGB"
+    assert img.size == (64, 64)
+
+
+def test_grayscale_tiff_float32(tmp_path):
+    arr = np.random.rand(64, 64).astype(np.float32)
+    path = tmp_path / "grayscale.tif"
+    tifffile.imwrite(str(path), arr)
+
+    data = LabelFile.load_image_file(str(path))
+    img = PIL.Image.open(io.BytesIO(data))
+    assert img.size == (64, 64)
+
+
+def test_constant_value_tiff_returns_black(tmp_path):
+    arr = np.full((64, 64), 42.0, dtype=np.float32)
+    path = tmp_path / "constant.tif"
+    tifffile.imwrite(str(path), arr)
+
+    data = LabelFile.load_image_file(str(path))
+    img = PIL.Image.open(io.BytesIO(data))
+    assert img.size == (64, 64)
+    assert np.array(img).max() == 0
+
+
+def test_two_band_tiff_falls_back_to_first_band(tmp_path):
+    arr = np.random.rand(64, 64, 2).astype(np.float32)
+    path = tmp_path / "twoband.tif"
+    tifffile.imwrite(str(path), arr)
+
+    data = LabelFile.load_image_file(str(path))
+    img = PIL.Image.open(io.BytesIO(data))
+    assert img.size == (64, 64)