Optimize OME-TIFF reader performance

src/tilefusion/core.py

@@ -6,6 +6,7 @@
 
 import gc
 import json
+import threading
 from concurrent.futures import ThreadPoolExecutor
 from multiprocessing import cpu_count
 from pathlib import Path
@@ -217,6 +218,94 @@ def __init__(
         self.fused_ts = None
         self.center = None
 
+        # Thread-local storage for TiffFile handles (thread-safe concurrent access)
+        self._thread_local = threading.local()
+        self._handles_lock = threading.Lock()
+        self._all_handles: List = []  # Track all handles for cleanup
+
+    def close(self) -> None:
+        """
+        Close any open file handles to release resources.
+
+        This should be called when finished using a TileFusion instance,
+        or use it as a context manager (``with TileFusion(...) as tf:``)
+        for automatic cleanup. Important for OME-TIFF inputs where file
+        handles are kept open for performance.
+        """
+        # Close all thread-local handles
+        with self._handles_lock:
+            for handle in self._all_handles:
+                try:
+                    handle.close()
+                except Exception:
+                    pass  # Best-effort cleanup
+            self._all_handles.clear()
+
+        # Clear this thread's handle reference
+        if hasattr(self._thread_local, "tiff_handle"):
+            self._thread_local.tiff_handle = None
+
+        # Close the original metadata handle (backward compatibility)
+        if self._metadata is not None and "tiff_handle" in self._metadata:
+            handle = self._metadata.pop("tiff_handle", None)
+            if handle is not None:
+                try:
+                    handle.close()
+                except Exception:
+                    pass  # Best-effort cleanup
+
+    def __enter__(self) -> "TileFusion":
+        """Enter the runtime context."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the runtime context and close file handles."""
+        self.close()
+
+    def __del__(self):
+        """Destructor to ensure file handles are closed."""
+        try:
+            self.close()
+        except AttributeError:
+            pass  # Object may be partially initialized
+
+    def _get_thread_local_handle(self):
+        """
+        Get or create a thread-local TiffFile handle for the current thread.
+
+        Each thread gets its own file handle to ensure thread-safe concurrent
+        reads. This avoids race conditions that can occur when multiple threads
+        share a single file descriptor (seek + read is not atomic on Windows).
+
+        Returns
+        -------
+        tifffile.TiffFile or None
+            Thread-local handle for OME-TIFF files, None for other formats.
+        """
+        # Only applies to OME-TIFF format (not zarr, individual tiffs, etc.)
+        if (
+            self._is_zarr_format
+            or self._is_individual_tiffs_format
+            or self._is_ome_tiff_tiles_format
+        ):
+            return None
+
+        # Check if this thread already has a handle
+        if hasattr(self._thread_local, "tiff_handle") and self._thread_local.tiff_handle is not None:
+            return self._thread_local.tiff_handle
+
+        # Create a new handle for this thread
+        import tifffile
+
+        handle = tifffile.TiffFile(self.tiff_path)
+        self._thread_local.tiff_handle = handle
+
+        # Track for cleanup
+        with self._handles_lock:
+            self._all_handles.append(handle)
+
+        return handle
+
     # -------------------------------------------------------------------------
     # Properties
     # -------------------------------------------------------------------------
@@ -319,7 +408,9 @@ def _read_tile(self, tile_idx: int, z_level: int = None, time_idx: int = 0) -> n
                 time_idx=time_idx,
             )
         else:
-            return read_ome_tiff_tile(self.tiff_path, tile_idx)
+            # Use thread-local handle for thread-safe concurrent reads
+            handle = self._get_thread_local_handle()
+            return read_ome_tiff_tile(self.tiff_path, tile_idx, handle)
 
     def _read_tile_region(
         self,
@@ -364,7 +455,9 @@ def _read_tile_region(
                 time_idx=time_idx,
             )
         else:
-            return read_ome_tiff_region(self.tiff_path, tile_idx, y_slice, x_slice)
+            # Use thread-local handle for thread-safe concurrent reads
+            handle = self._get_thread_local_handle()
+            return read_ome_tiff_region(self.tiff_path, tile_idx, y_slice, x_slice, handle)
 
     # -------------------------------------------------------------------------
     # Registration

src/tilefusion/io/ome_tiff.py

@@ -5,7 +5,7 @@
 """
 
 from pathlib import Path
-from typing import Dict, List, Tuple, Any
+from typing import Any, Dict, Optional
 import xml.etree.ElementTree as ET
 
 import numpy as np
@@ -31,8 +31,17 @@ def load_ome_tiff_metadata(tiff_path: Path) -> Dict[str, Any]:
         - channels: int
         - pixel_size: (py, px)
         - tile_positions: list of (y, x) tuples
+        - tiff_handle: tifffile.TiffFile
+            Open TIFF file handle kept for fast repeated access. The caller is
+            responsible for closing this handle by calling ``tiff_handle.close()``
+            when it is no longer needed to avoid resource leaks. The handle
+            remains valid until it is explicitly closed, or until a higher-level
+            context manager (if used) closes it on your behalf.
     """
-    with tifffile.TiffFile(tiff_path) as tif:
+    # Keep file handle open for fast repeated access
+    tif = tifffile.TiffFile(tiff_path)
+
+    try:
         if not tif.ome_metadata:
             raise ValueError("TIFF file does not contain OME metadata")
 
@@ -66,19 +75,25 @@ def load_ome_tiff_metadata(tiff_path: Path) -> Dict[str, Any]:
             else:
                 tile_positions.append((0.0, 0.0))
 
-    return {
-        "n_tiles": n_tiles,
-        "n_series": n_series,
-        "shape": (Y, X),
-        "channels": channels,
-        "time_dim": time_dim,
-        "position_dim": position_dim,
-        "pixel_size": pixel_size,
-        "tile_positions": tile_positions,
-    }
-
-
-def read_ome_tiff_tile(tiff_path: Path, tile_idx: int) -> np.ndarray:
+        return {
+            "n_tiles": n_tiles,
+            "n_series": n_series,
+            "shape": (Y, X),
+            "channels": channels,
+            "time_dim": time_dim,
+            "position_dim": position_dim,
+            "pixel_size": pixel_size,
+            "tile_positions": tile_positions,
+            "tiff_handle": tif,
+        }
+    except Exception:
+        tif.close()
+        raise
+
+
+def read_ome_tiff_tile(
+    tiff_path: Path, tile_idx: int, tiff_handle: Optional[tifffile.TiffFile] = None
+) -> np.ndarray:
     """
     Read a single tile from OME-TIFF (all channels).
 
@@ -88,14 +103,28 @@ def read_ome_tiff_tile(tiff_path: Path, tile_idx: int) -> np.ndarray:
         Path to the OME-TIFF file.
     tile_idx : int
         Index of the tile to read.
+    tiff_handle : TiffFile, optional
+        Cached TiffFile handle for faster access. For repeated reads,
+        keep the handle open and pass it here, or use TileFusion which
+        manages this automatically.
 
     Returns
     -------
     arr : ndarray of shape (C, Y, X)
         Tile data as float32.
+
+    Warning
+    -------
+    A single TiffFile handle is NOT thread-safe for concurrent reads.
+    On Windows, seek+read operations are not atomic, leading to data
+    corruption. Use separate handles per thread (TileFusion handles
+    this automatically via thread-local storage).
     """
-    with tifffile.TiffFile(tiff_path) as tif:
-        arr = tif.series[tile_idx].asarray()
+    if tiff_handle is not None:
+        arr = tiff_handle.series[tile_idx].asarray()
+    else:
+        with tifffile.TiffFile(tiff_path) as tif:
+            arr = tif.series[tile_idx].asarray()
     if arr.ndim == 2:
         arr = arr[np.newaxis, :, :]
     # Flip along Y axis to correct orientation
@@ -108,6 +137,7 @@ def read_ome_tiff_region(
     tile_idx: int,
     y_slice: slice,
     x_slice: slice,
+    tiff_handle: Optional[tifffile.TiffFile] = None,
 ) -> np.ndarray:
     """
     Read a region of a tile from OME-TIFF.
@@ -120,14 +150,28 @@ def read_ome_tiff_region(
         Index of the tile.
     y_slice, x_slice : slice
         Region to read.
+    tiff_handle : TiffFile, optional
+        Cached TiffFile handle for faster access. For repeated reads,
+        keep the handle open and pass it here, or use TileFusion which
+        manages this automatically.
 
     Returns
     -------
     arr : ndarray of shape (C, h, w)
         Tile region as float32.
+
+    Warning
+    -------
+    A single TiffFile handle is NOT thread-safe for concurrent reads.
+    On Windows, seek+read operations are not atomic, leading to data
+    corruption. Use separate handles per thread (TileFusion handles
+    this automatically via thread-local storage).
     """
-    with tifffile.TiffFile(tiff_path) as tif:
-        arr = tif.series[tile_idx].asarray()
+    if tiff_handle is not None:
+        arr = tiff_handle.series[tile_idx].asarray()
+    else:
+        with tifffile.TiffFile(tiff_path) as tif:
+            arr = tif.series[tile_idx].asarray()
     if arr.ndim == 2:
         arr = arr[np.newaxis, :, :]
     # Flip along Y axis to correct orientation