Update MSR file wrapper specification (#139)

* Change msr sample collection for per

* Update msr sample collection: added optional args, merged with other class.

* Fix type annotations and docstring issues in MSRSampleCollection

- Use Optional/Tuple from typing instead of built-in generics
- Add type annotations to dtype and shape constructor params
- Fix StackSelection | None annotation on read_sample
- Fix Note section to accurately describe shape mismatch risk
- Fix Sequence[int]/Sequence[str] in prose (was tuple[int]/tuple[str])
- Clarify that file read is skipped only when both dtype and shape are given
- Revert E203 slice spacing change in MSRDataset.__getitem__

---------

Co-authored-by: Oleh.Bakumenko <oleh.bakumenko@itmp.fraunhofer.de>

Author: anwai98

elf/io/msr_wrapper.py

@@ -1,6 +1,6 @@
 import os
 from collections.abc import Mapping, Sequence
-from typing import Union
+from typing import Optional, Tuple, Union
 
 import numpy as np
 
@@ -41,7 +41,7 @@ def _read_msr_stack(msr_path: PathLike, stack_index: int = 0) -> np.ndarray:
     return image
 
 
-def _normalize_stack_selection(stack_selection: StackSelection) -> tuple[Union[int, str], ...]:
+def _normalize_stack_selection(stack_selection: StackSelection) -> Tuple[Union[int, str], ...]:
     if isinstance(stack_selection, (str, int)):
         return (stack_selection,)
     stack_selection = tuple(stack_selection)
@@ -50,7 +50,7 @@ def _normalize_stack_selection(stack_selection: StackSelection) -> tuple[Union[i
     return stack_selection
 
 
-def _normalize_stack_indices(stack_index: StackIndexSelection = 0) -> tuple[int, ...]:
+def _normalize_stack_indices(stack_index: StackIndexSelection = 0) -> Tuple[int, ...]:
     if isinstance(stack_index, int):
         return (stack_index,)
     stack_index = tuple(stack_index)
@@ -61,7 +61,7 @@ def _normalize_stack_indices(stack_index: StackIndexSelection = 0) -> tuple[int,
     return stack_index
 
 
-def _normalize_stack_names(stack_names: StackNameSelection) -> tuple[str, ...]:
+def _normalize_stack_names(stack_names: StackNameSelection) -> Tuple[str, ...]:
     if isinstance(stack_names, str):
         return (stack_names,)
     stack_names = tuple(stack_names)
@@ -72,7 +72,7 @@ def _normalize_stack_names(stack_names: StackNameSelection) -> tuple[str, ...]:
     return stack_names
 
 
-def _resolve_stack_indices(msr, stack_selection: StackSelection) -> tuple[int, ...]:
+def _resolve_stack_indices(msr, stack_selection: StackSelection) -> Tuple[int, ...]:
     resolved = []
     for stack in _normalize_stack_selection(stack_selection):
         if isinstance(stack, int):
@@ -82,27 +82,57 @@ def _resolve_stack_indices(msr, stack_selection: StackSelection) -> tuple[int, .
     return tuple(resolved)
 
 
-def _get_msr_stack_shape(msr_path: PathLike, stack_selection: StackSelection = 0) -> tuple[int, int]:
+def _get_msr_stack_shape(msr_path: PathLike, stack_selection: StackSelection = 0) -> Tuple[int, int]:
     _require_msr_reader()
     with OBFFile(os.fspath(msr_path)) as msr:
         stack_index = _resolve_stack_indices(msr, stack_selection)[0]
     return tuple(_read_msr_stack(msr_path, stack_index=stack_index).shape)
 
 
-def _resolve_stack_indices_for_path(msr_path: PathLike, stack_selection: StackSelection) -> tuple[int, ...]:
+def _resolve_stack_indices_for_path(msr_path: PathLike, stack_selection: StackSelection) -> Tuple[int, ...]:
     _require_msr_reader()
     with OBFFile(os.fspath(msr_path)) as msr:
         return _resolve_stack_indices(msr, stack_selection)
 
 
 class MSRSampleCollection:
-    """Collection-like helper for loading one sample per MSR file."""
+    """Collection-like helper for loading one sample per MSR file.
+
+    Each MSR file in image_paths is treated as one sample. Stacks to load are specified globally
+    via stack_index (integer | Sequence[int]) or stack_names (str | Sequence[str]);
+    these are mutually exclusive. The resolved selection is stored as self.stack_selection and used
+    by default in read_sample.
+
+    dtype and shape are inferred by loading the first file unless both are provided explicitly.
+
+    For datasets where stack order or naming is inconsistent across files, read_sample accepts a
+    per-call stack_selection that overrides self.stack_selection for that single read. This covers
+    the case where e.g. channel 0 in one file is a different marker than channel 0 in another, or
+    where stack names contain typos or extra suffixes that differ per file.
+
+    Args:
+        image_paths: Paths to the .msr files.
+        stack_index: Integer or sequence of integers selecting stacks by position. Default 0.
+        stack_names: String or sequence of strings selecting stacks by name. Mutually exclusive
+            with stack_index.
+        dtype: numpy dtype of the loaded arrays. Inferred from the first file if None.
+        shape: Shape of one sample as a tuple. Inferred from the first file if None.
+            For a single stack this is (H, W); for multiple stacks it is (C, H, W).
+
+    Note:
+        When read_sample is called with a stack_selection that selects a different number of
+        channels than the constructor, the returned array shape will not match self.shape.
+        Similarly, if dtype or shape are provided explicitly but do not match the actual file
+        contents, self.dtype and self.shape will silently diverge from what read_sample returns.
+    """
 
     def __init__(
         self,
         image_paths: Sequence[PathLike],
         stack_index: StackIndexSelection = 0,
         stack_names: StackNameSelection | None = None,
+        dtype: Optional[np.dtype] = None,
+        shape: Optional[Tuple[int, ...]] = None,
     ):
         self.image_paths = [os.fspath(path) for path in image_paths]
         if stack_names is not None and stack_index != 0:
@@ -111,10 +141,17 @@ def __init__(
             self.stack_selection = _normalize_stack_names(stack_names)
         else:
             self.stack_selection = _normalize_stack_indices(stack_index)
-        self.stack_indices = _resolve_stack_indices_for_path(self.image_paths[0], self.stack_selection)
-        sample = _read_msr_stack(self.image_paths[0], self.stack_indices[0])
-        self._dtype = sample.dtype
-        self._shape = sample.shape if len(self.stack_indices) == 1 else (len(self.stack_indices),) + tuple(sample.shape)
+
+        if dtype is None or shape is None:
+            stack_indices = _resolve_stack_indices_for_path(self.image_paths[0], self.stack_selection)
+            sample = _read_msr_stack(self.image_paths[0], stack_indices[0])
+            if dtype is None:
+                dtype = sample.dtype
+            if shape is None:
+                shape = sample.shape if len(stack_indices) == 1 else (len(stack_indices),) + tuple(sample.shape)
+
+        self._dtype = dtype
+        self._shape = shape
 
     @property
     def dtype(self):
@@ -140,8 +177,21 @@ def size(self):
     def attrs(self):
         return {}
 
-    def read_sample(self, index: int) -> np.ndarray:
-        stack_indices = _resolve_stack_indices_for_path(self.image_paths[index], self.stack_selection)
+    def read_sample(self, index: int, stack_selection: Optional[StackSelection] = None) -> np.ndarray:
+        """Load one sample from image_paths[index].
+
+        Args:
+            index: Position in image_paths to read.
+            stack_selection: Optional override for which stacks to load from this specific file.
+                Accepts the same forms as the constructor (int, str, or sequences thereof).
+                When provided, self.stack_selection is ignored for this call. Useful when stack
+                order or names differ per file.
+
+        Returns:
+            numpy array of shape (H, W) for a single stack, or (C, H, W) for multiple stacks.
+        """
+        selection = self.stack_selection if stack_selection is None else stack_selection
+        stack_indices = _resolve_stack_indices_for_path(self.image_paths[index], selection)
         if len(stack_indices) == 1:
             return _read_msr_stack(self.image_paths[index], stack_indices[0])
         data = [_read_msr_stack(self.image_paths[index], stack_index) for stack_index in stack_indices]