From aac2d892db745f5609c894602770260c6a75328c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 10 Jun 2026 09:41:33 -0600
Subject: [PATCH 1/6] Enhance evaluation suite to support visualization as grid
 of multi-input and multi-target dataset and predictions

---
 .../evaluation/evaluation_utils.py            |  56 +++---
 .../evaluation/predict_utils.py               |  51 +++--
 .../evaluation/visualization.py               | 178 +++++++++++++++---
 3 files changed, 218 insertions(+), 67 deletions(-)

diff --git a/src/virtual_stain_flow/evaluation/evaluation_utils.py b/src/virtual_stain_flow/evaluation/evaluation_utils.py
index 8d1d285..bb2e4a2 100644
--- a/src/virtual_stain_flow/evaluation/evaluation_utils.py
+++ b/src/virtual_stain_flow/evaluation/evaluation_utils.py
@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Any
 
 import numpy as np
 import pandas as pd
@@ -10,13 +10,25 @@
 from virtual_stain_flow.datasets.base_wrapper_dataset import BaseWrapperDataset
 
 
+def _to_numpy_image(value: Any) -> np.ndarray:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().numpy()
+    return np.asarray(value)
+
+
+def _normalize_to_list(sample: Any) -> List[np.ndarray]:
+    if isinstance(sample, (list, tuple)):
+        return [_to_numpy_image(item) for item in sample]
+    return [_to_numpy_image(sample)]
+
+
 def extract_samples_from_dataset(
     dataset: Union[BaseImageDataset, CropImageDataset, BaseWrapperDataset],
     indices: List[int],
 ) -> Tuple[
-    List[np.ndarray],
-    List[np.ndarray],
-    Optional[List[np.ndarray]],
+    List[Union[np.ndarray, List[np.ndarray]]],
+    List[Union[np.ndarray, List[np.ndarray]]],
+    Optional[List[Union[np.ndarray, List[np.ndarray]]]],
     Optional[List[Tuple[int, int]]],
 ]:
     """
@@ -26,13 +38,15 @@ def extract_samples_from_dataset(
     (x, y) coordinates of each crop for visualization with bounding boxes.
 
     :param dataset: A BaseImageDataset or CropImageDataset instance.
-    :param indices: List of dataset indices to extract.
-    :return: Tuple of (inputs, targets, raw_images, patch_coords).
-        - inputs: List of numpy arrays, each with shape (C, H, W) or (H, W).
-        - targets: List of numpy arrays, each with shape (C, H, W) or (H, W).
-        - raw_images: List of numpy arrays for CropImageDataset (original uncropped images),
-          or None for BaseImageDataset.
-        - patch_coords: List of (x, y) tuples for CropImageDataset, or None for BaseImageDataset.
+        :param indices: List of dataset indices to extract.
+        :return: Tuple of (inputs, targets, raw_images, patch_coords).
+                - inputs: List of numpy arrays, each with shape (C, H, W) or (H, W).
+                    Multi-input samples can be provided as a list of arrays per sample.
+                - targets: List of numpy arrays, each with shape (C, H, W) or (H, W).
+                    Multi-target samples can be provided as a list of arrays per sample.
+                - raw_images: List of numpy arrays for CropImageDataset (original uncropped images),
+                    or None for BaseImageDataset.
+                - patch_coords: List of (x, y) tuples for CropImageDataset, or None for BaseImageDataset.
     """
     is_wrapper_dataset = False
     if isinstance(dataset, BaseWrapperDataset):
@@ -55,9 +69,9 @@ def extract_samples_from_dataset(
             f"max index requested: {max(indices)}"
         )
 
-    inputs: List[np.ndarray] = []
-    targets: List[np.ndarray] = []
-    raw_images: Optional[List[np.ndarray]] = [] if is_crop_dataset else None
+    inputs: List[Union[np.ndarray, List[np.ndarray]]] = []
+    targets: List[Union[np.ndarray, List[np.ndarray]]] = []
+    raw_images: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = [] if is_crop_dataset else None
     patch_coords: Optional[List[Tuple[int, int]]] = [] if is_crop_dataset else None
 
     for idx in indices:
@@ -65,15 +79,11 @@ def extract_samples_from_dataset(
         input_tensor, target_tensor = dataset[idx]
 
         # Convert to numpy - handle both Tensor and ndarray inputs
-        if isinstance(input_tensor, torch.Tensor):
-            inputs.append(input_tensor.numpy())
-        else:
-            inputs.append(np.asarray(input_tensor))
-
-        if isinstance(target_tensor, torch.Tensor):
-            targets.append(target_tensor.numpy())
-        else:
-            targets.append(np.asarray(target_tensor))
+        input_list = _normalize_to_list(input_tensor)
+        target_list = _normalize_to_list(target_tensor)
+
+        inputs.append(input_list[0] if len(input_list) == 1 else input_list)
+        targets.append(target_list[0] if len(target_list) == 1 else target_list)
 
         if is_crop_dataset:
             # Access the original uncropped image and crop coordinates
diff --git a/src/virtual_stain_flow/evaluation/predict_utils.py b/src/virtual_stain_flow/evaluation/predict_utils.py
index b0eea34..4e5829a 100644
--- a/src/virtual_stain_flow/evaluation/predict_utils.py
+++ b/src/virtual_stain_flow/evaluation/predict_utils.py
@@ -1,18 +1,26 @@
-from typing import Optional, List, Tuple, Callable
+from typing import Optional, List, Tuple, Callable, Union, Any
 
 import torch
 import numpy as np
 from torch.utils.data import DataLoader, Dataset, Subset
 from albumentations import ImageOnlyTransform, Compose
 
+def _move_to_device(value: Any, device: Union[str, torch.device]) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.to(device)
+    if isinstance(value, (list, tuple)):
+        return type(value)(_move_to_device(item, device) for item in value)
+    return value
+
+
 def predict_image(
     dataset: Dataset,
     model: torch.nn.Module,
     batch_size: int = 1,
-    device: str = "cpu",
+    device: Union[str, torch.device] = "cpu",
     num_workers: int = 0,
-    indices: Optional[List[int]] = None
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    indices: Optional[List[int]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, Union[torch.Tensor, List[torch.Tensor]]]:
     """
     Runs a model on a dataset, performing a forward pass on all (or a subset of) input images 
     in evaluation mode and returning a stacked tensor of predictions.
@@ -27,6 +35,7 @@ def predict_image(
     :param indices: Optional list of dataset indices to subset the dataset before inference.
 
     :return: Tuple of stacked target, prediction, and input tensors.
+        For multi-input datasets, the third element is a list of stacked input tensors.
     """
     # Subset the dataset if indices are provided
     if indices is not None:
@@ -38,25 +47,41 @@ def predict_image(
     model.to(device)
     model.eval()
 
-    predictions, targets, inputs = [], [], []
+    predictions, targets = [], []
+    inputs: Union[List[torch.Tensor], List[List[torch.Tensor]]] = []
 
     with torch.no_grad():
         for input, target in dataloader:  # Unpacking (input_tensor, target_tensor)
-            input = input.to(device)  # Move input data to the specified device
+            input = _move_to_device(input, device)
 
             # Forward pass
-            prediction = model(input)
-            
+            if isinstance(input, (list, tuple)):
+                prediction = model(*input)
+            else:
+                prediction = model(input)
+
             # output both target and prediction tensors for metric
             targets.append(target.cpu())
             predictions.append(prediction.cpu())  # Move to CPU for stacking
-            inputs.append(input.cpu())
+
+            if isinstance(input, (list, tuple)):
+                if not inputs:
+                    inputs = [[] for _ in range(len(input))]
+                for idx, item in enumerate(input):
+                    inputs[idx].append(item.cpu())
+            else:
+                inputs.append(input.cpu())
+
+    if inputs and isinstance(inputs[0], list):
+        inputs_stacked = [torch.cat(batch_list, dim=0) for batch_list in inputs]  # type: ignore[arg-type]
+    else:
+        inputs_stacked = torch.cat(inputs, dim=0)  # type: ignore[arg-type]
 
     return (
-        torch.cat(targets, dim=0), 
-        torch.cat(predictions, dim=0), 
-        torch.cat(inputs, dim=0)
-    ) 
+        torch.cat(targets, dim=0),
+        torch.cat(predictions, dim=0),
+        inputs_stacked,
+    )
 
 def process_tensor_image(
     img_tensor: torch.Tensor,
diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py
index 271f54e..08da34e 100644
--- a/src/virtual_stain_flow/evaluation/visualization.py
+++ b/src/virtual_stain_flow/evaluation/visualization.py
@@ -5,7 +5,7 @@
 Supports both BaseImageDataset and CropImageDataset with optional metrics display.
 """
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Any
 
 import numpy as np
 import pandas as pd
@@ -20,14 +20,56 @@
 from .predict_utils import predict_image
 
 
+def _to_numpy_image(value: Any) -> np.ndarray:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().numpy()
+    return np.asarray(value)
+
+
+def _normalize_to_list(sample: Any) -> List[np.ndarray]:
+    if isinstance(sample, (list, tuple)):
+        return [_to_numpy_image(item) for item in sample]
+    return [_to_numpy_image(sample)]
+
+
+def _split_channels(img: np.ndarray, channel_indices: Optional[List[int]]) -> List[np.ndarray]:
+    if img.ndim == 2:
+        if channel_indices is not None and any(idx != 0 for idx in channel_indices):
+            raise ValueError("Channel indices out of range for 2D image.")
+        return [img]
+
+    if img.ndim == 3:
+        total_channels = img.shape[0]
+        indices = channel_indices if channel_indices is not None else list(range(total_channels))
+        for idx in indices:
+            if idx < 0 or idx >= total_channels:
+                raise ValueError(
+                    f"Channel index {idx} out of range for image with {total_channels} channels."
+                )
+        return [img[idx] for idx in indices]
+
+    raise ValueError(f"Unsupported image shape for visualization: {img.shape}")
+
+
+def _split_sample_channels(sample: Any, channel_indices: Optional[List[int]]) -> List[np.ndarray]:
+    split_images: List[np.ndarray] = []
+    for item in _normalize_to_list(sample):
+        split_images.extend(_split_channels(item, channel_indices))
+    return split_images
+
+
+def _build_titles(prefix: str, count: int) -> List[str]:
+    return [f"{prefix} {i + 1}" for i in range(count)]
+
+
 def plot_predictions_grid(
-    inputs: List[np.ndarray],
-    targets: List[np.ndarray],
-    predictions: Optional[List[np.ndarray]] = None,
+    inputs: List[Union[np.ndarray, List[np.ndarray]]],
+    targets: List[Union[np.ndarray, List[np.ndarray]]],
+    predictions: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = None,
     *,
     sample_indices: Optional[List[int]] = None,
     row_label_prefix: str = "",
-    raw_images: Optional[List[np.ndarray]] = None,
+    raw_images: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = None,
     patch_coords: Optional[List[Tuple[int, int]]] = None,
     metrics_df: Optional[pd.DataFrame] = None,
     save_path: Optional[str] = None,
@@ -36,6 +78,10 @@ def plot_predictions_grid(
     show_plot: bool = True,
     wspace: float = 0.05,
     hspace: float = 0.15,
+    raw_channel_indices: Optional[List[int]] = None,
+    input_channel_indices: Optional[List[int]] = None,
+    target_channel_indices: Optional[List[int]] = None,
+    prediction_channel_indices: Optional[List[int]] = None,
 ) -> plt.Figure:
     """
     Plot a grid of images comparing inputs, targets, and optionally predictions.
@@ -47,8 +93,11 @@ def plot_predictions_grid(
     - [Prediction] (only if predictions provided, with optional metrics in title)
 
     :param inputs: List of input images, each (C, H, W) or (H, W).
+        Multi-input samples can be provided as a list of arrays per sample.
     :param targets: List of target images, each (C, H, W) or (H, W).
+        Multi-target samples can be provided as a list of arrays per sample.
     :param predictions: Optional list of prediction images, each (C, H, W) or (H, W).
+        Multi-output samples can be provided as a list of arrays per sample.
         If None, only inputs and targets are displayed.
     :param sample_indices: Optional list of indices to display as row labels.
         If None, uses 0-based sequential indices.
@@ -66,6 +115,10 @@ def plot_predictions_grid(
     :param show_plot: Whether to display the plot (default: True).
     :param wspace: Horizontal spacing between subplots (default: 0.05).
     :param hspace: Vertical spacing between subplots (default: 0.15).
+    :param raw_channel_indices: Optional list of channel indices to display for raw images.
+    :param input_channel_indices: Optional list of channel indices to display for inputs.
+    :param target_channel_indices: Optional list of channel indices to display for targets.
+    :param prediction_channel_indices: Optional list of channel indices to display for predictions.
     """
     num_samples = len(inputs)
     if num_samples == 0:
@@ -84,24 +137,55 @@ def plot_predictions_grid(
         )
 
     has_raw_images = raw_images is not None and len(raw_images) > 0
+    if has_raw_images and len(raw_images) != num_samples:
+        raise ValueError(
+            f"Length mismatch: inputs ({num_samples}), raw_images ({len(raw_images)})"
+        )
+
+    if sample_indices is not None and len(sample_indices) != num_samples:
+        raise ValueError(
+            f"Length mismatch: inputs ({num_samples}), sample_indices ({len(sample_indices)})"
+        )
+
+    if has_predictions and prediction_channel_indices is None:
+        prediction_channel_indices = target_channel_indices
     
-    # Determine number of columns based on what's provided
-    # Base: Input, Target (2 cols) or Input, Target, Prediction (3 cols)
-    # With raw: Raw, Input, Target (3 cols) or Raw, Input, Target, Prediction (4 cols)
-    num_cols = 2 + (1 if has_raw_images else 0) + (1 if has_predictions else 0)
+    # Determine columns based on channel splits in the first sample
+    raw_first = _split_sample_channels(raw_images[0], raw_channel_indices) if has_raw_images else []
+    input_first = _split_sample_channels(inputs[0], input_channel_indices)
+    target_first = _split_sample_channels(targets[0], target_channel_indices)
+    pred_first = (
+        _split_sample_channels(predictions[0], prediction_channel_indices)
+        if has_predictions
+        else []
+    )
+
+    if has_predictions and len(pred_first) != len(target_first):
+        raise ValueError(
+            "Target and prediction channel counts must match for paired display."
+        )
+
+    raw_titles = _build_titles("Raw Input", len(raw_first))
+    input_titles = _build_titles("Input", len(input_first))
+    target_titles = _build_titles("Target", len(target_first))
+    pred_titles = _build_titles("Prediction", len(pred_first))
+
+    if has_predictions:
+        interleaved_titles = [
+            title
+            for i in range(len(target_titles))
+            for title in (target_titles[i], pred_titles[i])
+        ]
+    else:
+        interleaved_titles = target_titles
+
+    column_titles = raw_titles + input_titles + interleaved_titles
+    num_cols = len(column_titles)
 
     # Default sample indices if not provided
     if sample_indices is None:
         sample_indices = list(range(num_samples))
 
-    # Column titles - build dynamically
-    column_titles = []
-    if has_raw_images:
-        column_titles.append("Raw Input")
-    column_titles.extend(["Input", "Target"])
-    if has_predictions:
-        column_titles.append("Prediction")
-
     # Create figure
     fig_width = panel_width * num_cols
     fig_height = panel_width * num_samples
@@ -112,13 +196,35 @@ def plot_predictions_grid(
         axes = axes.reshape(1, -1)
 
     for row_idx in range(num_samples):
-        # Build image set for this row dynamically
-        img_set = []
-        if has_raw_images:
-            img_set.append(raw_images[row_idx])
-        img_set.extend([inputs[row_idx], targets[row_idx]])
+        raw_row = _split_sample_channels(raw_images[row_idx], raw_channel_indices) if has_raw_images else []
+        input_row = _split_sample_channels(inputs[row_idx], input_channel_indices)
+        target_row = _split_sample_channels(targets[row_idx], target_channel_indices)
+        pred_row = (
+            _split_sample_channels(predictions[row_idx], prediction_channel_indices)
+            if has_predictions
+            else []
+        )
+
+        if has_predictions and len(pred_row) != len(target_row):
+            raise ValueError(
+                f"Row {row_idx} has mismatched target/prediction channel counts."
+            )
+
         if has_predictions:
-            img_set.append(predictions[row_idx])
+            target_pred_row = [
+                img
+                for i in range(len(target_row))
+                for img in (target_row[i], pred_row[i])
+            ]
+        else:
+            target_pred_row = target_row
+
+        img_set = raw_row + input_row + target_pred_row
+
+        if len(img_set) != num_cols:
+            raise ValueError(
+                f"Row {row_idx} has {len(img_set)} columns, expected {num_cols}."
+            )
 
         for col_idx, img in enumerate(img_set):
             ax = axes[row_idx, col_idx]
@@ -134,10 +240,10 @@ def plot_predictions_grid(
             ax.axis("off")
 
             # Draw bounding box on raw image
-            if has_raw_images and col_idx == 0 and patch_coords is not None:
+            if has_raw_images and col_idx < len(raw_titles) and patch_coords is not None:
                 patch_x, patch_y = patch_coords[row_idx]
                 # Infer patch size from target shape
-                target_shape = np.squeeze(targets[row_idx]).shape
+                target_shape = np.squeeze(target_row[0]).shape
                 patch_h, patch_w = target_shape[-2], target_shape[-1]
                 rect = Rectangle(
                     (patch_x, patch_y),
@@ -149,8 +255,8 @@ def plot_predictions_grid(
                 )
                 ax.add_patch(rect)
 
-            # Row label on the Input column (first column if no raw images, second if raw images)
-            input_col_idx = 1 if has_raw_images else 0
+            # Row label on the first input column
+            input_col_idx = len(raw_titles)
             if col_idx == input_col_idx:
                 row_label = f"{row_label_prefix}{sample_indices[row_idx]}"
                 # Determine text color based on top-left corner brightness
@@ -175,7 +281,7 @@ def plot_predictions_grid(
                     horizontalalignment="left",
                 )
 
-        # Metrics on prediction column (last column, only if predictions provided)
+        # Metrics on last prediction column (only if predictions provided)
         if has_predictions and metrics_df is not None and row_idx < len(metrics_df):
             metric_values = metrics_df.iloc[row_idx]
             metric_text = "\n".join(
@@ -218,7 +324,8 @@ def plot_dataset_grid(
     :param indices: List of dataset indices to display.
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
-        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace.
+        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
+        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices.
     """
     # Extract samples from dataset
     (
@@ -266,7 +373,8 @@ def plot_predictions_grid_from_model(
     :param device: Device for inference ("cpu" or "cuda").
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
-        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace.
+        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
+        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices.
     """
     # Step 1: Run inference
     targets_tensor, predictions_tensor, inputs_tensor = predict_image(
@@ -282,7 +390,15 @@ def plot_predictions_grid_from_model(
     _, _, raw_images, patch_coords = extract_samples_from_dataset(dataset, indices)    
     # use the collected input and target stack at prediction time instead of
     # re-extract
-    inputs, targets = inputs_tensor.numpy(), targets_tensor.numpy()
+    if isinstance(inputs_tensor, list):
+        inputs = [
+            [inputs_tensor[i][row_idx].numpy() for i in range(len(inputs_tensor))]
+            for row_idx in range(len(indices))
+        ]
+    else:
+        inputs = inputs_tensor.numpy()
+
+    targets = targets_tensor.numpy()
 
     # Convert predictions tensor to list of numpy arrays
     predictions = [predictions_tensor[i].numpy() for i in range(len(indices))]

From 69a9ee8ec0930d9bb3abdc55adaa1b3b8fcc442c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 22 Jul 2026 11:15:49 -0600
Subject: [PATCH 2/6] Refactor docstring in evaluation_utils.py for clarity and
 completeness

---
 src/virtual_stain_flow/evaluation/evaluation_utils.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/virtual_stain_flow/evaluation/evaluation_utils.py b/src/virtual_stain_flow/evaluation/evaluation_utils.py
index bb2e4a2..55c28c5 100644
--- a/src/virtual_stain_flow/evaluation/evaluation_utils.py
+++ b/src/virtual_stain_flow/evaluation/evaluation_utils.py
@@ -1,3 +1,11 @@
+"""
+evaluation_utils.py
+
+Evaluation utilities 
+ - interaction with dataset to obtain normalized images and annotations
+ - evaluation of per-image metrics
+"""
+
 from typing import List, Optional, Tuple, Union, Any
 
 import numpy as np

From e589c91c588e904a268559f4de1aee1f6b4e9c4d Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 22 Jul 2026 12:41:58 -0600
Subject: [PATCH 3/6] Refactor evaluation utilities, modifies image tensor
 normalization to be stacked as N, C, H, W following extraction, which is more
 natural, as opposed to returning as list of C, H, W arrays. Also enrich the
 extraction function to best effort retrieve real channel names and title
 plots with that. Also separated metric eval as separate module as
 visualization utils.

---
 .../evaluation/evaluation_utils.py            |  99 +------
 .../evaluation/visualization.py               | 244 +++++++++---------
 .../evaluation/visualization_utils.py         | 143 ++++++++++
 3 files changed, 274 insertions(+), 212 deletions(-)
 create mode 100644 src/virtual_stain_flow/evaluation/visualization_utils.py

diff --git a/src/virtual_stain_flow/evaluation/evaluation_utils.py b/src/virtual_stain_flow/evaluation/evaluation_utils.py
index 55c28c5..003beb9 100644
--- a/src/virtual_stain_flow/evaluation/evaluation_utils.py
+++ b/src/virtual_stain_flow/evaluation/evaluation_utils.py
@@ -1,110 +1,15 @@
 """
 evaluation_utils.py
 
-Evaluation utilities 
- - interaction with dataset to obtain normalized images and annotations
- - evaluation of per-image metrics
+Utilities for evaluating model predictions.
 """
 
-from typing import List, Optional, Tuple, Union, Any
+from typing import List, Optional
 
-import numpy as np
 import pandas as pd
 import torch
 from torch.nn import Module
 
-from ..datasets.base_dataset import BaseImageDataset
-from ..datasets.crop_dataset import CropImageDataset
-from virtual_stain_flow.datasets.base_wrapper_dataset import BaseWrapperDataset
-
-
-def _to_numpy_image(value: Any) -> np.ndarray:
-    if isinstance(value, torch.Tensor):
-        return value.detach().cpu().numpy()
-    return np.asarray(value)
-
-
-def _normalize_to_list(sample: Any) -> List[np.ndarray]:
-    if isinstance(sample, (list, tuple)):
-        return [_to_numpy_image(item) for item in sample]
-    return [_to_numpy_image(sample)]
-
-
-def extract_samples_from_dataset(
-    dataset: Union[BaseImageDataset, CropImageDataset, BaseWrapperDataset],
-    indices: List[int],
-) -> Tuple[
-    List[Union[np.ndarray, List[np.ndarray]]],
-    List[Union[np.ndarray, List[np.ndarray]]],
-    Optional[List[Union[np.ndarray, List[np.ndarray]]]],
-    Optional[List[Tuple[int, int]]],
-]:
-    """
-    Extract input/target samples and optional raw images with patch coordinates from a dataset.
-
-    For CropImageDataset, also extracts the original (uncropped) input images and the
-    (x, y) coordinates of each crop for visualization with bounding boxes.
-
-    :param dataset: A BaseImageDataset or CropImageDataset instance.
-        :param indices: List of dataset indices to extract.
-        :return: Tuple of (inputs, targets, raw_images, patch_coords).
-                - inputs: List of numpy arrays, each with shape (C, H, W) or (H, W).
-                    Multi-input samples can be provided as a list of arrays per sample.
-                - targets: List of numpy arrays, each with shape (C, H, W) or (H, W).
-                    Multi-target samples can be provided as a list of arrays per sample.
-                - raw_images: List of numpy arrays for CropImageDataset (original uncropped images),
-                    or None for BaseImageDataset.
-                - patch_coords: List of (x, y) tuples for CropImageDataset, or None for BaseImageDataset.
-    """
-    is_wrapper_dataset = False
-    if isinstance(dataset, BaseWrapperDataset):
-        is_crop_dataset = isinstance(dataset.original, CropImageDataset)
-        is_wrapper_dataset = True
-    elif isinstance(dataset, CropImageDataset):
-        is_crop_dataset = True
-    elif isinstance(dataset, BaseImageDataset):
-        is_crop_dataset = False
-    else:
-        raise ValueError(
-            "Unsupported dataset type. Expected BaseImageDataset, CropImageDataset, or BaseWrapperDataset.")
-    
-    if not indices:
-        raise ValueError("Indices list cannot be empty.")
-    
-    if max(indices) >= len(dataset):
-        raise IndexError(
-            f"Index out of range. Dataset length: {len(dataset)}, "
-            f"max index requested: {max(indices)}"
-        )
-
-    inputs: List[Union[np.ndarray, List[np.ndarray]]] = []
-    targets: List[Union[np.ndarray, List[np.ndarray]]] = []
-    raw_images: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = [] if is_crop_dataset else None
-    patch_coords: Optional[List[Tuple[int, int]]] = [] if is_crop_dataset else None
-
-    for idx in indices:
-        # Access dataset item to trigger lazy loading and state update
-        input_tensor, target_tensor = dataset[idx]
-
-        # Convert to numpy - handle both Tensor and ndarray inputs
-        input_list = _normalize_to_list(input_tensor)
-        target_list = _normalize_to_list(target_tensor)
-
-        inputs.append(input_list[0] if len(input_list) == 1 else input_list)
-        targets.append(target_list[0] if len(target_list) == 1 else target_list)
-
-        if is_crop_dataset:
-            # Access the original uncropped image and crop coordinates
-            # These are populated after __getitem__ is called
-            if is_wrapper_dataset:
-                raw_images.append(dataset.original.original_input_image[0])
-                patch_coords.append((dataset.original.crop_info.x, dataset.original.crop_info.y))
-            else:
-                raw_images.append(dataset.original_input_image[0])
-                patch_coords.append((dataset.crop_info.x, dataset.crop_info.y))
-
-    return inputs, targets, raw_images, patch_coords
-
 
 def evaluate_per_image_metric(
     predictions: torch.Tensor,
diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py
index 08da34e..2628d07 100644
--- a/src/virtual_stain_flow/evaluation/visualization.py
+++ b/src/virtual_stain_flow/evaluation/visualization.py
@@ -5,7 +5,7 @@
 Supports both BaseImageDataset and CropImageDataset with optional metrics display.
 """
 
-from typing import List, Optional, Tuple, Union, Any
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -16,61 +16,63 @@
 from ..datasets.base_dataset import BaseImageDataset
 from ..datasets.crop_dataset import CropImageDataset
 from ..datasets.base_wrapper_dataset import BaseWrapperDataset
-from .evaluation_utils import evaluate_per_image_metric, extract_samples_from_dataset
+from .evaluation_utils import evaluate_per_image_metric
 from .predict_utils import predict_image
+from .visualization_utils import extract_samples_from_dataset
 
 
-def _to_numpy_image(value: Any) -> np.ndarray:
-    if isinstance(value, torch.Tensor):
-        return value.detach().cpu().numpy()
-    return np.asarray(value)
-
-
-def _normalize_to_list(sample: Any) -> List[np.ndarray]:
-    if isinstance(sample, (list, tuple)):
-        return [_to_numpy_image(item) for item in sample]
-    return [_to_numpy_image(sample)]
-
-
-def _split_channels(img: np.ndarray, channel_indices: Optional[List[int]]) -> List[np.ndarray]:
-    if img.ndim == 2:
-        if channel_indices is not None and any(idx != 0 for idx in channel_indices):
-            raise ValueError("Channel indices out of range for 2D image.")
-        return [img]
+def _select_channels(
+        images: np.ndarray, channel_indices: Optional[List[int]], name: str
+    ) -> Tuple[np.ndarray, List[int]]:
+    """
+    Helper for selecting channel by index working with (N, C, H, W) image stacks.
 
-    if img.ndim == 3:
-        total_channels = img.shape[0]
-        indices = channel_indices if channel_indices is not None else list(range(total_channels))
-        for idx in indices:
-            if idx < 0 or idx >= total_channels:
-                raise ValueError(
-                    f"Channel index {idx} out of range for image with {total_channels} channels."
-                )
-        return [img[idx] for idx in indices]
+    :param images: Image stack with shape (N, C, H, W).
+    :param channel_indices: Optional list of channel indices to select.
+        If None, all channels are selected.
+    :param name: Name of the image stack for error messages.
+    :return: Image stack with selected channels and their original channel indices.
+    """
 
-    raise ValueError(f"Unsupported image shape for visualization: {img.shape}")
+    if images.ndim != 4:
+        raise ValueError(f"{name} must have shape (N, C, H, W), received {images.shape}.")
 
+    indices = channel_indices if channel_indices is not None else list(range(images.shape[1]))
+    if not indices:
+        raise ValueError(f"{name} channel indices cannot be empty.")
+    if any(index < 0 or index >= images.shape[1] for index in indices):
+        raise ValueError(
+            f"Channel indices out of range for {name} with {images.shape[1]} channels."
+        )
 
-def _split_sample_channels(sample: Any, channel_indices: Optional[List[int]]) -> List[np.ndarray]:
-    split_images: List[np.ndarray] = []
-    for item in _normalize_to_list(sample):
-        split_images.extend(_split_channels(item, channel_indices))
-    return split_images
+    return images[:, indices, :, :], indices
 
 
-def _build_titles(prefix: str, count: int) -> List[str]:
-    return [f"{prefix} {i + 1}" for i in range(count)]
+def _build_titles(
+    prefix: str,
+    channel_indices: List[int],
+    channel_names: Optional[List[str]],
+) -> List[str]:
+    """
+    Build display titles with optional dataset channel names.
+    """
+    return [
+        f"{channel_names[channel_index]} ({prefix} {channel_index + 1})"
+        if channel_names is not None and channel_index < len(channel_names)
+        else f"{prefix} {position + 1}"
+        for position, channel_index in enumerate(channel_indices)
+    ]
 
 
 def plot_predictions_grid(
-    inputs: List[Union[np.ndarray, List[np.ndarray]]],
-    targets: List[Union[np.ndarray, List[np.ndarray]]],
-    predictions: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = None,
+    inputs: np.ndarray,
+    targets: np.ndarray,
+    predictions: Optional[np.ndarray] = None,
     *,
     sample_indices: Optional[List[int]] = None,
     row_label_prefix: str = "",
-    raw_images: Optional[List[Union[np.ndarray, List[np.ndarray]]]] = None,
-    patch_coords: Optional[List[Tuple[int, int]]] = None,
+    raw_images: Optional[np.ndarray] = None,
+    patch_coords: Optional[List[Tuple[int, int, int, int]]] = None,
     metrics_df: Optional[pd.DataFrame] = None,
     save_path: Optional[str] = None,
     cmap: str = "gray",
@@ -78,12 +80,17 @@ def plot_predictions_grid(
     show_plot: bool = True,
     wspace: float = 0.05,
     hspace: float = 0.15,
+    raw_channel_names: Optional[List[str]] = None,
+    input_channel_names: Optional[List[str]] = None,
+    target_channel_names: Optional[List[str]] = None,
+    prediction_channel_names: Optional[List[str]] = None,
     raw_channel_indices: Optional[List[int]] = None,
     input_channel_indices: Optional[List[int]] = None,
     target_channel_indices: Optional[List[int]] = None,
     prediction_channel_indices: Optional[List[int]] = None,
 ) -> plt.Figure:
     """
+    Core visualization function for visualizing grid of input/target and predictions. 
     Plot a grid of images comparing inputs, targets, and optionally predictions.
 
     Each row represents one sample. Columns are:
@@ -92,19 +99,16 @@ def plot_predictions_grid(
     - Target
     - [Prediction] (only if predictions provided, with optional metrics in title)
 
-    :param inputs: List of input images, each (C, H, W) or (H, W).
-        Multi-input samples can be provided as a list of arrays per sample.
-    :param targets: List of target images, each (C, H, W) or (H, W).
-        Multi-target samples can be provided as a list of arrays per sample.
-    :param predictions: Optional list of prediction images, each (C, H, W) or (H, W).
-        Multi-output samples can be provided as a list of arrays per sample.
+    :param inputs: Input images with shape (N, C, H, W).
+    :param targets: Target images with shape (N, C, H, W).
+    :param predictions: Optional prediction images with shape (N, C, H, W).
         If None, only inputs and targets are displayed.
     :param sample_indices: Optional list of indices to display as row labels.
         If None, uses 0-based sequential indices.
     :param row_label_prefix: Prefix for row labels (default: "").
         E.g., "Sample " would give labels "Sample 0", "Sample 1", etc.
-    :param raw_images: Optional list of original uncropped images for CropImageDataset.
-    :param patch_coords: Optional list of (x, y) crop coordinates for bounding boxes.
+    :param raw_images: Optional original uncropped images with shape (N, C, H, W).
+    :param patch_coords: Optional list of (x, y, width, height) crop rectangles.
         Only used when raw_images is provided.
     :param metrics_df: Optional DataFrame with per-image metrics to display.
         Each row corresponds to a sample; column names become metric labels.
@@ -115,31 +119,45 @@ def plot_predictions_grid(
     :param show_plot: Whether to display the plot (default: True).
     :param wspace: Horizontal spacing between subplots (default: 0.05).
     :param hspace: Vertical spacing between subplots (default: 0.15).
+    :param raw_channel_names: Optional names for raw image channels.
+    :param input_channel_names: Optional names for input image channels.
+    :param target_channel_names: Optional names for target image channels.
+    :param prediction_channel_names: Optional names for prediction image channels.
     :param raw_channel_indices: Optional list of channel indices to display for raw images.
     :param input_channel_indices: Optional list of channel indices to display for inputs.
     :param target_channel_indices: Optional list of channel indices to display for targets.
     :param prediction_channel_indices: Optional list of channel indices to display for predictions.
     """
-    num_samples = len(inputs)
+    if inputs.ndim != 4:
+        raise ValueError(f"Inputs must have shape (N, C, H, W), received {inputs.shape}.")
+
+    num_samples = inputs.shape[0]
     if num_samples == 0:
         raise ValueError("No samples provided to plot.")
 
-    # Validate lengths match
-    if len(targets) != num_samples:
+    if targets.ndim != 4 or targets.shape[0] != num_samples:
         raise ValueError(
-            f"Length mismatch: inputs ({num_samples}), targets ({len(targets)})"
+            f"Inputs and targets must have matching batch sizes; received "
+            f"{inputs.shape} and {targets.shape}."
         )
-    
-    has_predictions = predictions is not None and len(predictions) > 0
-    if has_predictions and len(predictions) != num_samples:
+
+    has_predictions = predictions is not None
+    if has_predictions and (predictions.ndim != 4 or predictions.shape[0] != num_samples):
         raise ValueError(
-            f"Length mismatch: inputs ({num_samples}), predictions ({len(predictions)})"
+            f"Inputs and predictions must have matching batch sizes; received "
+            f"{inputs.shape} and {predictions.shape}."
         )
 
-    has_raw_images = raw_images is not None and len(raw_images) > 0
-    if has_raw_images and len(raw_images) != num_samples:
+    has_raw_images = raw_images is not None
+    if has_raw_images and (raw_images.ndim != 4 or raw_images.shape[0] != num_samples):
         raise ValueError(
-            f"Length mismatch: inputs ({num_samples}), raw_images ({len(raw_images)})"
+            f"Inputs and raw images must have matching batch sizes; received "
+            f"{inputs.shape} and {raw_images.shape}."
+        )
+
+    if has_raw_images and patch_coords is not None and len(patch_coords) != num_samples:
+        raise ValueError(
+            f"Length mismatch: inputs ({num_samples}), patch_coords ({len(patch_coords)})"
         )
 
     if sample_indices is not None and len(sample_indices) != num_samples:
@@ -149,26 +167,31 @@ def plot_predictions_grid(
 
     if has_predictions and prediction_channel_indices is None:
         prediction_channel_indices = target_channel_indices
-    
-    # Determine columns based on channel splits in the first sample
-    raw_first = _split_sample_channels(raw_images[0], raw_channel_indices) if has_raw_images else []
-    input_first = _split_sample_channels(inputs[0], input_channel_indices)
-    target_first = _split_sample_channels(targets[0], target_channel_indices)
-    pred_first = (
-        _split_sample_channels(predictions[0], prediction_channel_indices)
+    if has_predictions and prediction_channel_names is None:
+        prediction_channel_names = target_channel_names
+
+    raw_images, raw_indices = (
+        _select_channels(raw_images, raw_channel_indices, "Raw images")
+        if has_raw_images
+        else (None, [])
+    )
+    inputs, input_indices = _select_channels(inputs, input_channel_indices, "Inputs")
+    targets, target_indices = _select_channels(targets, target_channel_indices, "Targets")
+    predictions, prediction_indices = (
+        _select_channels(predictions, prediction_channel_indices, "Predictions")
         if has_predictions
-        else []
+        else (None, [])
     )
 
-    if has_predictions and len(pred_first) != len(target_first):
+    if has_predictions and predictions.shape[1] != targets.shape[1]:
         raise ValueError(
             "Target and prediction channel counts must match for paired display."
         )
 
-    raw_titles = _build_titles("Raw Input", len(raw_first))
-    input_titles = _build_titles("Input", len(input_first))
-    target_titles = _build_titles("Target", len(target_first))
-    pred_titles = _build_titles("Prediction", len(pred_first))
+    raw_titles = _build_titles("Raw Input", raw_indices, raw_channel_names)
+    input_titles = _build_titles("Input", input_indices, input_channel_names)
+    target_titles = _build_titles("Target", target_indices, target_channel_names)
+    pred_titles = _build_titles("Prediction", prediction_indices, prediction_channel_names)
 
     if has_predictions:
         interleaved_titles = [
@@ -196,28 +219,16 @@ def plot_predictions_grid(
         axes = axes.reshape(1, -1)
 
     for row_idx in range(num_samples):
-        raw_row = _split_sample_channels(raw_images[row_idx], raw_channel_indices) if has_raw_images else []
-        input_row = _split_sample_channels(inputs[row_idx], input_channel_indices)
-        target_row = _split_sample_channels(targets[row_idx], target_channel_indices)
-        pred_row = (
-            _split_sample_channels(predictions[row_idx], prediction_channel_indices)
-            if has_predictions
-            else []
-        )
-
-        if has_predictions and len(pred_row) != len(target_row):
-            raise ValueError(
-                f"Row {row_idx} has mismatched target/prediction channel counts."
-            )
+        raw_row = list(raw_images[row_idx]) if has_raw_images else []
+        input_row = list(inputs[row_idx])
+        target_row = list(targets[row_idx])
+        pred_row = list(predictions[row_idx]) if has_predictions else []
 
-        if has_predictions:
-            target_pred_row = [
-                img
-                for i in range(len(target_row))
-                for img in (target_row[i], pred_row[i])
-            ]
-        else:
-            target_pred_row = target_row
+        target_pred_row = [
+            image
+            for target_image, prediction_image in zip(target_row, pred_row)
+            for image in (target_image, prediction_image)
+        ] if has_predictions else target_row
 
         img_set = raw_row + input_row + target_pred_row
 
@@ -241,10 +252,7 @@ def plot_predictions_grid(
 
             # Draw bounding box on raw image
             if has_raw_images and col_idx < len(raw_titles) and patch_coords is not None:
-                patch_x, patch_y = patch_coords[row_idx]
-                # Infer patch size from target shape
-                target_shape = np.squeeze(target_row[0]).shape
-                patch_h, patch_w = target_shape[-2], target_shape[-1]
+                patch_x, patch_y, patch_w, patch_h = patch_coords[row_idx]
                 rect = Rectangle(
                     (patch_x, patch_y),
                     patch_w,
@@ -325,11 +333,12 @@ def plot_dataset_grid(
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
         Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
-        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices.
+        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices,
+        raw_channel_names, input_channel_names, target_channel_names, prediction_channel_names.
     """
     # Extract samples from dataset
     (
-        inputs, targets, raw_images, patch_coords
+        inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names
     ) = extract_samples_from_dataset(dataset, indices)
 
     # Plot without predictions
@@ -340,6 +349,9 @@ def plot_dataset_grid(
         sample_indices=indices,
         raw_images=raw_images,
         patch_coords=patch_coords,
+        raw_channel_names=input_channel_names,
+        input_channel_names=input_channel_names,
+        target_channel_names=target_channel_names,
         metrics_df=None,
         save_path=save_path,
         **kwargs,
@@ -374,7 +386,8 @@ def plot_predictions_grid_from_model(
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
         Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
-        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices.
+        raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices,
+        raw_channel_names, input_channel_names, target_channel_names, prediction_channel_names.
     """
     # Step 1: Run inference
     targets_tensor, predictions_tensor, inputs_tensor = predict_image(
@@ -386,22 +399,19 @@ def plot_predictions_grid_from_model(
     if metrics:
         metrics_df = evaluate_per_image_metric(predictions_tensor, targets_tensor, metrics)
 
-    # Step 3: Extract samples from dataset (need to re-access for raw images in CropImageDataset)
-    _, _, raw_images, patch_coords = extract_samples_from_dataset(dataset, indices)    
-    # use the collected input and target stack at prediction time instead of
-    # re-extract
+    # Step 3: Re-access the dataset for CropImageDataset raw images and crop metadata.
+    (
+        _, _, raw_images, patch_coords, input_channel_names, target_channel_names
+    ) = extract_samples_from_dataset(dataset, indices)
     if isinstance(inputs_tensor, list):
-        inputs = [
-            [inputs_tensor[i][row_idx].numpy() for i in range(len(inputs_tensor))]
-            for row_idx in range(len(indices))
-        ]
-    else:
-        inputs = inputs_tensor.numpy()
-
-    targets = targets_tensor.numpy()
+        raise ValueError(
+            "Visualization requires a single batched input tensor with shape (N, C, H, W); "
+            "multi-input sequences are not supported."
+        )
 
-    # Convert predictions tensor to list of numpy arrays
-    predictions = [predictions_tensor[i].numpy() for i in range(len(indices))]
+    inputs = inputs_tensor.detach().cpu().numpy()
+    targets = targets_tensor.detach().cpu().numpy()
+    predictions = predictions_tensor.detach().cpu().numpy()
 
     # Step 4: Plot
     return plot_predictions_grid(
@@ -411,6 +421,10 @@ def plot_predictions_grid_from_model(
         sample_indices=indices,
         raw_images=raw_images,
         patch_coords=patch_coords,
+        raw_channel_names=input_channel_names,
+        input_channel_names=input_channel_names,
+        target_channel_names=target_channel_names,
+        prediction_channel_names=target_channel_names,
         metrics_df=metrics_df,
         save_path=save_path,
         **kwargs,
diff --git a/src/virtual_stain_flow/evaluation/visualization_utils.py b/src/virtual_stain_flow/evaluation/visualization_utils.py
new file mode 100644
index 0000000..8083cc8
--- /dev/null
+++ b/src/virtual_stain_flow/evaluation/visualization_utils.py
@@ -0,0 +1,143 @@
+"""
+visualization_utils.py
+
+Utilities for extracting batch-form image data and crop metadata for visualization.
+"""
+
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+
+from ..datasets.base_dataset import BaseImageDataset
+from ..datasets.base_wrapper_dataset import BaseWrapperDataset
+from ..datasets.crop_dataset import CropImageDataset
+
+PatchCoords = Tuple[int, int, int, int]
+DatasetType = Union[BaseImageDataset, CropImageDataset, BaseWrapperDataset]
+
+
+def _to_numpy_image(value: Union[np.ndarray, torch.Tensor], name: str) -> np.ndarray:
+    if isinstance(value, (list, tuple)):
+        raise ValueError(
+            f"{name} must be a single channel-first image, not a sequence of images."
+        )
+
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().numpy()
+    else:
+        value = np.asarray(value)
+
+    if value.ndim != 3:
+        raise ValueError(
+            f"{name} must have shape (C, H, W), received shape {value.shape}."
+        )
+
+    return value
+
+
+def _stack_images(images: List[np.ndarray], name: str) -> np.ndarray:
+    """
+    Helper function to stack a list of images into a single numpy array 
+        with shape (N, C, H, W).
+    """
+    try:
+        return np.stack(images)
+    except ValueError as error:
+        raise ValueError(f"{name} must have a consistent (C, H, W) shape.") from error
+
+
+def _get_channel_names(
+    channel_keys: Optional[Union[str, Sequence[str]]], expected_count: int
+) -> Optional[List[str]]:
+    if channel_keys is None:
+        return None
+
+    names = [channel_keys] if isinstance(channel_keys, str) else list(channel_keys)
+    return names if len(names) == expected_count else None
+
+
+def extract_samples_from_dataset(
+    dataset: DatasetType,
+    indices: List[int],
+) -> Tuple[
+    np.ndarray,
+    np.ndarray,
+    Optional[np.ndarray],
+    Optional[List[PatchCoords]],
+    Optional[List[str]],
+    Optional[List[str]],
+]:
+    """
+    Extract image batches and optional crop metadata for visualization.
+    Primary function of this abstraction is to provide a consistent data
+        access interface between Dataset objects and plotting functions by
+        extracting input/target with __get_item__ and also accessing raw
+        images and crop annotations when available. 
+
+    :param dataset: A BaseImageDataset, CropImageDataset, or BaseWrapperDataset.
+    :param indices: Dataset indices to extract, in the displayed order.
+    :return: Tuple of inputs, targets, raw images, crop coordinates, input channel
+        names, and target channel names. Image arrays have shape (N, C, H, W).
+        Channel names are None when unavailable or inconsistent with the image
+        channels. Raw images and patch coordinates are None unless the dataset is
+        a CropImageDataset or wraps one. Patch coordinates are (x, y, width, height).
+    """
+    if isinstance(dataset, BaseWrapperDataset):
+        metadata_dataset = dataset.original
+        crop_dataset = metadata_dataset if isinstance(metadata_dataset, CropImageDataset) else None
+    elif isinstance(dataset, CropImageDataset):
+        metadata_dataset = dataset
+        crop_dataset = dataset
+    elif isinstance(dataset, BaseImageDataset):
+        metadata_dataset = dataset
+        crop_dataset = None
+    else:
+        raise ValueError(
+            "Unsupported dataset type. Expected BaseImageDataset, CropImageDataset, "
+            "or BaseWrapperDataset."
+        )
+
+    if not indices:
+        raise ValueError("Indices list cannot be empty.")
+    if min(indices) < 0 or max(indices) >= len(dataset):
+        raise IndexError(
+            f"Index out of range. Dataset length: {len(dataset)}, "
+            f"requested indices: {indices}"
+        )
+
+    inputs: List[np.ndarray] = []
+    targets: List[np.ndarray] = []
+    raw_images: Optional[List[np.ndarray]] = [] if crop_dataset is not None else None
+    patch_coords: Optional[List[PatchCoords]] = [] if crop_dataset is not None else None
+
+    for index in indices:
+        # Need to access and cast single image samples of
+        # dimension (C, H, W) from torch tensor to np arrays
+        input_image, target_image = dataset[index]
+        inputs.append(_to_numpy_image(input_image, "Dataset input"))
+        targets.append(_to_numpy_image(target_image, "Dataset target"))
+
+        if crop_dataset is not None:
+            raw_image = crop_dataset.original_input_image
+            crop_info = crop_dataset.crop_info
+            if raw_image is None or crop_info is None:
+                raise RuntimeError("Crop image metadata was not populated after dataset access.")
+
+            raw_images.append(_to_numpy_image(raw_image, "Crop raw image"))
+            patch_coords.append(
+                (crop_info.x, crop_info.y, crop_info.width, crop_info.height)
+            )
+
+    input_batch = _stack_images(inputs, "Dataset inputs")
+    target_batch = _stack_images(targets, "Dataset targets")
+    raw_batch = _stack_images(raw_images, "Crop raw images") if raw_images is not None else None
+
+    return (
+        input_batch,
+        target_batch,
+        raw_batch,
+        patch_coords,
+        _get_channel_names(metadata_dataset.input_channel_keys, input_batch.shape[1]),
+        _get_channel_names(metadata_dataset.target_channel_keys, target_batch.shape[1]),
+    )

From 15e882834462ed0203ca2756c8aa9e4b937af901 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 22 Jul 2026 12:43:31 -0600
Subject: [PATCH 4/6] Add tests for visualization, fix test breaks following
 refactor

---
 .../evaluation/predict_utils.py               | 170 +++++++++++-------
 .../vsf_logging/callbacks/PlotCallback.py     |   5 +-
 tests/conftest.py                             |  57 ++++++
 tests/evaluation/test_evaluation_utils.py     |  27 +++
 tests/evaluation/test_predict_utils.py        | 117 ++++++++++++
 tests/evaluation/test_visualization.py        |  81 +++++++++
 tests/evaluation/test_visualization_utils.py  |  45 +++++
 7 files changed, 432 insertions(+), 70 deletions(-)
 create mode 100644 tests/evaluation/test_evaluation_utils.py
 create mode 100644 tests/evaluation/test_predict_utils.py
 create mode 100644 tests/evaluation/test_visualization.py
 create mode 100644 tests/evaluation/test_visualization_utils.py

diff --git a/src/virtual_stain_flow/evaluation/predict_utils.py b/src/virtual_stain_flow/evaluation/predict_utils.py
index 4e5829a..fcc6680 100644
--- a/src/virtual_stain_flow/evaluation/predict_utils.py
+++ b/src/virtual_stain_flow/evaluation/predict_utils.py
@@ -1,9 +1,8 @@
-from typing import Optional, List, Tuple, Callable, Union, Any
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
-import numpy as np
 from torch.utils.data import DataLoader, Dataset, Subset
-from albumentations import ImageOnlyTransform, Compose
+
 
 def _move_to_device(value: Any, device: Union[str, torch.device]) -> Any:
     if isinstance(value, torch.Tensor):
@@ -13,6 +12,48 @@ def _move_to_device(value: Any, device: Union[str, torch.device]) -> Any:
     return value
 
 
+def _validate_image_tensor(value: Any, name: str) -> torch.Tensor:
+    if not isinstance(value, torch.Tensor):
+        raise TypeError(f"{name} must be a torch.Tensor, received {type(value).__name__}.")
+    if value.ndim != 4:
+        raise ValueError(
+            f"{name} must have shape (N, C, H, W), received {tuple(value.shape)}."
+        )
+    if value.shape[0] == 0:
+        raise ValueError(f"{name} cannot have an empty batch dimension.")
+    return value
+
+
+def _validate_batch(
+    input_batch: Any,
+    target_batch: Any,
+    expected_multi_input: Optional[bool],
+    expected_input_count: Optional[int],
+) -> Tuple[bool, int]:
+    _validate_image_tensor(target_batch, "Target batch")
+
+    if isinstance(input_batch, (list, tuple)):
+        if not input_batch:
+            raise ValueError("Multi-input batches cannot be empty.")
+        if expected_multi_input is False:
+            raise ValueError("Input structure changed from single-input to multi-input during prediction.")
+        if expected_input_count is not None and expected_input_count != len(input_batch):
+            raise ValueError("Multi-input batch count changed during prediction.")
+
+        for input_index, input_tensor in enumerate(input_batch):
+            _validate_image_tensor(input_tensor, f"Input batch {input_index}")
+            if input_tensor.shape[0] != target_batch.shape[0]:
+                raise ValueError("Input and target batch sizes must match.")
+        return True, len(input_batch)
+
+    _validate_image_tensor(input_batch, "Input batch")
+    if input_batch.shape[0] != target_batch.shape[0]:
+        raise ValueError("Input and target batch sizes must match.")
+    if expected_multi_input is True:
+        raise ValueError("Input structure changed from multi-input to single-input during prediction.")
+    return False, 1
+
+
 def predict_image(
     dataset: Dataset,
     model: torch.nn.Module,
@@ -34,12 +75,29 @@ def predict_image(
     :param num_workers: Number of workers for the DataLoader (default is 0).
     :param indices: Optional list of dataset indices to subset the dataset before inference.
 
-    :return: Tuple of stacked target, prediction, and input tensors.
-        For multi-input datasets, the third element is a list of stacked input tensors.
+    :return: Tuple of stacked target, prediction, and input tensors. For multi-input
+        datasets, the third element is a list of stacked input tensors.
+    :raises ValueError: If the selected dataset is empty, a batch is malformed, or
+        its input structure changes during prediction.
+    :raises TypeError: If inputs, targets, or model predictions are not tensors.
     """
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive.")
+    if num_workers < 0:
+        raise ValueError("num_workers cannot be negative.")
+
     # Subset the dataset if indices are provided
     if indices is not None:
+        if not indices:
+            raise ValueError("indices cannot be empty.")
+        if min(indices) < 0 or max(indices) >= len(dataset):
+            raise IndexError(
+                f"Index out of range. Dataset length: {len(dataset)}, "
+                f"requested indices: {indices}"
+            )
         dataset = Subset(dataset, indices)
+    elif len(dataset) == 0:
+        raise ValueError("dataset cannot be empty.")
 
     # Create DataLoader for efficient batch processing
     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
@@ -47,80 +105,56 @@ def predict_image(
     model.to(device)
     model.eval()
 
-    predictions, targets = [], []
-    inputs: Union[List[torch.Tensor], List[List[torch.Tensor]]] = []
+    predictions: List[torch.Tensor] = []
+    targets: List[torch.Tensor] = []
+    input_batches: List[torch.Tensor] = []
+    multi_input_batches: Optional[List[List[torch.Tensor]]] = None
+    expected_multi_input: Optional[bool] = None
+    expected_input_count: Optional[int] = None
 
     with torch.no_grad():
-        for input, target in dataloader:  # Unpacking (input_tensor, target_tensor)
-            input = _move_to_device(input, device)
+        for input_batch, target_batch in dataloader:
+            is_multi_input, input_count = _validate_batch(
+                input_batch, target_batch, expected_multi_input, expected_input_count
+            )
+            expected_multi_input = is_multi_input
+            if is_multi_input:
+                expected_input_count = input_count
+
+            input_batch = _move_to_device(input_batch, device)
 
             # Forward pass
-            if isinstance(input, (list, tuple)):
-                prediction = model(*input)
+            if is_multi_input:
+                prediction = model(*input_batch)
             else:
-                prediction = model(input)
-
-            # output both target and prediction tensors for metric
-            targets.append(target.cpu())
-            predictions.append(prediction.cpu())  # Move to CPU for stacking
-
-            if isinstance(input, (list, tuple)):
-                if not inputs:
-                    inputs = [[] for _ in range(len(input))]
-                for idx, item in enumerate(input):
-                    inputs[idx].append(item.cpu())
+                prediction = model(input_batch)
+            _validate_image_tensor(prediction, "Model prediction")
+            if prediction.shape[0] != target_batch.shape[0]:
+                raise ValueError("Prediction and target batch sizes must match.")
+
+            targets.append(target_batch.cpu())
+            predictions.append(prediction.cpu())
+
+            if is_multi_input:
+                if multi_input_batches is None:
+                    multi_input_batches = [[] for _ in range(input_count)]
+                for input_index, input_tensor in enumerate(input_batch):
+                    multi_input_batches[input_index].append(input_tensor.cpu())
             else:
-                inputs.append(input.cpu())
+                input_batches.append(input_batch.cpu())
+
+    if not targets or not predictions:
+        raise RuntimeError("Prediction did not retrieve any batches.")
 
-    if inputs and isinstance(inputs[0], list):
-        inputs_stacked = [torch.cat(batch_list, dim=0) for batch_list in inputs]  # type: ignore[arg-type]
+    if multi_input_batches is not None:
+        inputs_stacked = [torch.cat(batch, dim=0) for batch in multi_input_batches]
     else:
-        inputs_stacked = torch.cat(inputs, dim=0)  # type: ignore[arg-type]
+        if not input_batches:
+            raise RuntimeError("Prediction did not accumulate any input batches.")
+        inputs_stacked = torch.cat(input_batches, dim=0)
 
     return (
         torch.cat(targets, dim=0),
         torch.cat(predictions, dim=0),
         inputs_stacked,
     )
-
-def process_tensor_image(
-    img_tensor: torch.Tensor,
-    dtype: Optional[np.dtype] = None,
-    dataset: Optional[Dataset] = None,
-    invert_function: Optional[Callable] = None
-) -> np.ndarray:
-    """
-    Processes model output/other image tensor by casting to numpy, applying an optional dtype casting, 
-    and inverting target transformations if a dataset with `target_transform` is provided.
-
-    :param img_tensor: Tensor stack of model-predicted images with shape (N, C, H, W).
-    :type img_tensor: torch.Tensor
-    :param dtype: Optional numpy dtype to cast the output array (default: None).
-    :type dtype: Optional[np.dtype], optional
-    :param dataset: Optional dataset object with `target_transform` to invert transformations.
-    :type dataset: Optional[torch.utils.data.Dataset], optional
-    :param invert_function: Optional function to invert transformations applied to the images.
-        If provided, overrides the invert function call from dataset transform.
-    :type invert_function: Optional[Callable], optional
-
-    :return: Processed numpy array of images with shape (N, C, H, W).
-    :rtype: np.ndarray
-    """
-    # Convert img_tensor to CPU and NumPy
-    output_images = img_tensor.cpu().numpy()
-
-    # Optionally cast to specified dtype
-    if dtype is not None:
-        output_images = output_images.astype(dtype)
-
-    # Apply invert function when supplied or transformation if invert function is supplied
-    if invert_function is not None and isinstance(invert_function, Callable):
-        output_images = np.array([invert_function(img) for img in output_images])
-    elif dataset is not None and hasattr(dataset, "target_transform"):
-        # Apply inverted target transformation if available
-        target_transform = dataset.target_transform
-        if isinstance(target_transform, (ImageOnlyTransform, Compose)):
-            # Apply the transformation on each image
-            output_images = np.array([target_transform.invert(img) for img in output_images])
-
-    return output_images
diff --git a/src/virtual_stain_flow/vsf_logging/callbacks/PlotCallback.py b/src/virtual_stain_flow/vsf_logging/callbacks/PlotCallback.py
index eb7690e..74167c5 100644
--- a/src/virtual_stain_flow/vsf_logging/callbacks/PlotCallback.py
+++ b/src/virtual_stain_flow/vsf_logging/callbacks/PlotCallback.py
@@ -6,6 +6,7 @@
 import torch.nn as nn
 from torch.utils.data import Dataset
 
+from...datasets.base_wrapper_dataset import BaseWrapperDataset
 from .LoggerCallback import (
     AbstractLoggerCallback,
     log_artifact_type
@@ -54,8 +55,8 @@ def __init__(
             self.__tmpdir = tempfile.TemporaryDirectory()
             self._path = pathlib.Path(self.__tmpdir.name)
 
-        if not isinstance(dataset, Dataset):
-            raise TypeError(f"Expected Dataset, got {type(dataset)}")
+        if not isinstance(dataset, (Dataset, BaseWrapperDataset)):
+            raise TypeError(f"Expected Dataset or BaseWrapperDataset, got {type(dataset)}")
 
         self._dataset = dataset
 
diff --git a/tests/conftest.py b/tests/conftest.py
index b635c69..5842619 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -8,9 +8,13 @@
 from types import SimpleNamespace
 
 import pytest
+import numpy as np
+import pandas as pd
 import torch
 from torch.utils.data import DataLoader, Dataset
 
+from virtual_stain_flow.datasets.base_dataset import BaseImageDataset
+from virtual_stain_flow.datasets.crop_dataset import CropImageDataset
 from virtual_stain_flow.trainers.AbstractTrainer import AbstractTrainer
 from virtual_stain_flow.trainers.logging_trainer import SingleGeneratorTrainer
 from virtual_stain_flow.vsf_logging import MlflowLogger
@@ -175,6 +179,59 @@ def __getitem__(self, idx):
     return ImageDataset(num_samples=20)
 
 
+@pytest.fixture
+def file_index(tmp_path):
+    from PIL import Image
+
+    test_data_dir = tmp_path / "test_data"
+    test_data_dir.mkdir()
+    paths = {
+        "input_ch1": [test_data_dir / f"img_{index}_in1.tif" for index in range(3)],
+        "input_ch2": [test_data_dir / f"img_{index}_in2.tif" for index in range(3)],
+        "target_ch1": [test_data_dir / f"img_{index}_target.tif" for index in range(3)],
+    }
+
+    for index in range(3):
+        for channel, start_value in (("input_ch1", 1), ("input_ch2", 2), ("target_ch1", 1)):
+            value = start_value + index * (2 if channel.startswith("input") else 1)
+            Image.fromarray(np.full((10, 10), value, dtype=np.uint16), mode="I;16").save(
+                paths[channel][index]
+            )
+
+    return pd.DataFrame(paths)
+
+
+@pytest.fixture
+def basic_dataset(file_index):
+    return BaseImageDataset(
+        file_index=file_index,
+        pil_image_mode="I;16",
+        input_channel_keys=["input_ch1", "input_ch2"],
+        target_channel_keys="target_ch1",
+        cache_capacity=8,
+    )
+
+
+@pytest.fixture
+def crop_specs():
+    return {
+        index: [((0, 0), 4, 4), ((5, 5), 4, 4)]
+        for index in range(3)
+    }
+
+
+@pytest.fixture
+def crop_dataset(file_index, crop_specs):
+    return CropImageDataset(
+        file_index=file_index,
+        crop_specs=crop_specs,
+        pil_image_mode="I;16",
+        input_channel_keys=["input_ch1", "input_ch2"],
+        target_channel_keys="target_ch1",
+        cache_capacity=8,
+    )
+
+
 @pytest.fixture
 def minimal_model():
     """Create a minimal PyTorch model."""
diff --git a/tests/evaluation/test_evaluation_utils.py b/tests/evaluation/test_evaluation_utils.py
new file mode 100644
index 0000000..e759e7d
--- /dev/null
+++ b/tests/evaluation/test_evaluation_utils.py
@@ -0,0 +1,27 @@
+import pandas as pd
+import pytest
+import torch
+
+from virtual_stain_flow.evaluation.evaluation_utils import evaluate_per_image_metric
+
+
+class MeanAbsoluteError(torch.nn.Module):
+    def forward(self, target, prediction):
+        return torch.abs(target - prediction).mean()
+
+
+def test_evaluate_per_image_metric_returns_one_row_per_image():
+    targets = torch.tensor([[[[1.0]]], [[[3.0]]]])
+    predictions = torch.tensor([[[[2.0]]], [[[1.0]]]])
+
+    result = evaluate_per_image_metric(predictions, targets, [MeanAbsoluteError()])
+
+    pd.testing.assert_frame_equal(
+        result,
+        pd.DataFrame({"MeanAbsoluteError": [1.0, 2.0]}),
+    )
+
+
+def test_evaluate_per_image_metric_rejects_mismatched_shapes():
+    with pytest.raises(ValueError, match="Shape mismatch"):
+        evaluate_per_image_metric(torch.zeros(1, 1, 2, 2), torch.zeros(1, 1, 3, 3), [])
diff --git a/tests/evaluation/test_predict_utils.py b/tests/evaluation/test_predict_utils.py
new file mode 100644
index 0000000..e03496c
--- /dev/null
+++ b/tests/evaluation/test_predict_utils.py
@@ -0,0 +1,117 @@
+import pytest
+import torch
+from torch.utils.data import Dataset
+
+from virtual_stain_flow.evaluation.predict_utils import predict_image
+
+
+class ImageDataset(Dataset):
+    def __init__(self, inputs, targets=None):
+        self.inputs = inputs
+        self.targets = targets or [torch.zeros(1, 4, 4) for _ in inputs]
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, index):
+        return self.inputs[index], self.targets[index]
+
+
+class IdentityModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.forward_calls = 0
+
+    def forward(self, inputs):
+        self.forward_calls += 1
+        return inputs
+
+
+class AddInputsModel(torch.nn.Module):
+    def forward(self, first_input, second_input):
+        return first_input + second_input
+
+
+class InvalidOutputModel(torch.nn.Module):
+    def forward(self, inputs):
+        return "not a tensor"
+
+
+def test_predict_image_stacks_single_input_batches():
+    inputs = [torch.full((1, 4, 4), float(index)) for index in range(3)]
+
+    targets, predictions, stacked_inputs = predict_image(
+        ImageDataset(inputs), IdentityModel(), batch_size=2, device="cpu"
+    )
+
+    assert targets.shape == (3, 1, 4, 4)
+    assert torch.equal(predictions, torch.stack(inputs))
+    assert torch.equal(stacked_inputs, torch.stack(inputs))
+
+
+def test_predict_image_stacks_each_multi_input_component():
+    first_inputs = [torch.ones(1, 4, 4), torch.full((1, 4, 4), 2.0)]
+    second_inputs = [torch.full((1, 4, 4), 3.0), torch.full((1, 4, 4), 4.0)]
+    dataset = ImageDataset(list(zip(first_inputs, second_inputs)))
+
+    targets, predictions, stacked_inputs = predict_image(
+        dataset, AddInputsModel(), batch_size=2, device="cpu"
+    )
+
+    assert targets.shape == (2, 1, 4, 4)
+    assert isinstance(stacked_inputs, list)
+    assert len(stacked_inputs) == 2
+    assert torch.equal(stacked_inputs[0], torch.stack(first_inputs))
+    assert torch.equal(stacked_inputs[1], torch.stack(second_inputs))
+    assert torch.equal(predictions, torch.stack(first_inputs) + torch.stack(second_inputs))
+
+
+def test_predict_image_rejects_empty_dataset():
+    with pytest.raises(ValueError, match="dataset cannot be empty"):
+        predict_image(ImageDataset([]), IdentityModel(), device="cpu")
+
+
+def test_predict_image_rejects_empty_indices():
+    dataset = ImageDataset([torch.zeros(1, 4, 4)])
+
+    with pytest.raises(ValueError, match="indices cannot be empty"):
+        predict_image(dataset, IdentityModel(), indices=[], device="cpu")
+
+
+def test_predict_image_rejects_empty_multi_input_before_forward_pass():
+    model = IdentityModel()
+    dataset = ImageDataset([[]])
+
+    with pytest.raises(ValueError, match="Multi-input batches cannot be empty"):
+        predict_image(dataset, model, device="cpu")
+
+    assert model.forward_calls == 0
+
+
+def test_predict_image_rejects_non_tensor_input_before_forward_pass():
+    model = IdentityModel()
+    dataset = ImageDataset(["invalid input"])
+
+    with pytest.raises(TypeError, match="Input batch 0 must be a torch.Tensor"):
+        predict_image(dataset, model, device="cpu")
+
+    assert model.forward_calls == 0
+
+
+def test_predict_image_rejects_input_structure_change_before_second_forward_pass():
+    model = IdentityModel()
+    single_input = torch.ones(1, 4, 4)
+    multi_input = (torch.ones(1, 4, 4), torch.ones(1, 4, 4))
+    dataset = ImageDataset([single_input, multi_input])
+
+    with pytest.raises(ValueError, match="single-input to multi-input"):
+        predict_image(dataset, model, batch_size=1, device="cpu")
+
+    assert model.forward_calls == 1
+
+
+def test_predict_image_rejects_non_tensor_model_output():
+    dataset = ImageDataset([torch.zeros(1, 4, 4)])
+
+    with pytest.raises(TypeError, match="Model prediction must be a torch.Tensor"):
+        predict_image(dataset, InvalidOutputModel(), device="cpu")
diff --git a/tests/evaluation/test_visualization.py b/tests/evaluation/test_visualization.py
new file mode 100644
index 0000000..260f61d
--- /dev/null
+++ b/tests/evaluation/test_visualization.py
@@ -0,0 +1,81 @@
+import matplotlib
+import numpy as np
+import pytest
+import torch
+
+matplotlib.use("Agg")
+
+from virtual_stain_flow.evaluation.visualization import (
+    plot_dataset_grid,
+    plot_predictions_grid,
+    plot_predictions_grid_from_model,
+)
+
+
+def test_plot_predictions_grid_uses_batched_channel_arrays():
+    inputs = np.zeros((2, 2, 4, 4))
+    targets = np.ones((2, 1, 4, 4))
+    predictions = np.full((2, 1, 4, 4), 2.0)
+
+    figure = plot_predictions_grid(
+        inputs,
+        targets,
+        predictions,
+        input_channel_indices=[1],
+        input_channel_names=["first_input", "second_input"],
+        target_channel_names=["target"],
+        show_plot=False,
+    )
+
+    assert len(figure.axes) == 6
+    assert [axis.get_title() for axis in figure.axes[:3]] == [
+        "second_input (Input 2)",
+        "target (Target 1)",
+        "target (Prediction 1)",
+    ]
+
+
+def test_plot_predictions_grid_falls_back_to_numbered_titles():
+    figure = plot_predictions_grid(
+        np.zeros((1, 1, 4, 4)),
+        np.zeros((1, 1, 4, 4)),
+        show_plot=False,
+    )
+
+    assert [axis.get_title() for axis in figure.axes] == ["Input 1", "Target 1"]
+
+
+def test_plot_predictions_grid_rejects_mismatched_batches():
+    with pytest.raises(ValueError, match="matching batch sizes"):
+        plot_predictions_grid(
+            np.zeros((2, 1, 4, 4)),
+            np.zeros((1, 1, 4, 4)),
+            show_plot=False,
+        )
+
+
+def test_plot_dataset_grid_draws_recorded_crop_rectangle(crop_dataset):
+    figure = plot_dataset_grid(crop_dataset, [1], show_plot=False)
+
+    raw_axes = figure.axes[:2]
+    assert len(figure.axes) == 5
+    for axis in raw_axes:
+        rectangle = axis.patches[0]
+        assert rectangle.get_xy() == (5, 5)
+        assert rectangle.get_width() == 4
+        assert rectangle.get_height() == 4
+
+
+def test_plot_predictions_grid_from_model_uses_tensor_batches(basic_dataset):
+    model = torch.nn.Conv2d(2, 1, kernel_size=1)
+
+    figure = plot_predictions_grid_from_model(
+        model,
+        basic_dataset,
+        indices=[0, 1],
+        metrics=[],
+        device="cpu",
+        show_plot=False,
+    )
+
+    assert len(figure.axes) == 8
diff --git a/tests/evaluation/test_visualization_utils.py b/tests/evaluation/test_visualization_utils.py
new file mode 100644
index 0000000..f842151
--- /dev/null
+++ b/tests/evaluation/test_visualization_utils.py
@@ -0,0 +1,45 @@
+import numpy as np
+import pytest
+
+from virtual_stain_flow.datasets.base_wrapper_dataset import BaseWrapperDataset
+from virtual_stain_flow.evaluation.visualization_utils import extract_samples_from_dataset
+
+
+def test_extract_samples_returns_channel_first_batches(basic_dataset):
+    (
+        inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names
+    ) = extract_samples_from_dataset(basic_dataset, [2, 0])
+
+    assert inputs.shape == (2, 2, 10, 10)
+    assert targets.shape == (2, 1, 10, 10)
+    assert raw_images is None
+    assert patch_coords is None
+    assert input_channel_names == ["input_ch1", "input_ch2"]
+    assert target_channel_names == ["target_ch1"]
+    assert np.all(inputs[0, 0] == 5)
+    assert np.all(inputs[1, 1] == 2)
+
+
+def test_extract_crop_samples_preserves_raw_channel_stack(crop_dataset):
+    (
+        inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names
+    ) = extract_samples_from_dataset(crop_dataset, [1, 2])
+
+    assert inputs.shape == (2, 2, 4, 4)
+    assert targets.shape == (2, 1, 4, 4)
+    assert raw_images.shape == (2, 2, 10, 10)
+    assert patch_coords == [(5, 5, 4, 4), (0, 0, 4, 4)]
+    assert input_channel_names == ["input_ch1", "input_ch2"]
+    assert target_channel_names == ["target_ch1"]
+    assert np.all(raw_images[0, 0] == 1)
+    assert np.all(raw_images[0, 1] == 2)
+
+
+def test_extract_samples_rejects_multiple_input_images(basic_dataset):
+    class MultiInputDataset(BaseWrapperDataset):
+        def __getitem__(self, index):
+            input_image, target_image = self._dataset[index]
+            return [input_image, input_image], target_image
+
+    with pytest.raises(ValueError, match="single channel-first image"):
+        extract_samples_from_dataset(MultiInputDataset(basic_dataset), [0])

From 323192523503bdc20c775876bd9524fc238f6c47 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 22 Jul 2026 12:52:07 -0600
Subject: [PATCH 5/6] Enhance image loading robustness by adding fallback to
 tifffile for TIFF files in DatasetManifest and add imagecodecs support to
 account for different tiff compressions

---
 pyproject.toml                                |  3 +-
 .../datasets/ds_engine/manifest.py            | 31 ++++++++++++++++---
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 2f7b3dd..fcc0af0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,8 @@ dependencies = [
   "notebook",
   "tifffile",
   "pandera[pandas]",
-  "boto3"
+  "boto3",
+  "imagecodecs"
 ]
 
 [project.optional-dependencies]
diff --git a/src/virtual_stain_flow/datasets/ds_engine/manifest.py b/src/virtual_stain_flow/datasets/ds_engine/manifest.py
index cb7649c..39b3884 100644
--- a/src/virtual_stain_flow/datasets/ds_engine/manifest.py
+++ b/src/virtual_stain_flow/datasets/ds_engine/manifest.py
@@ -24,6 +24,7 @@
 
 import numpy as np
 import pandas as pd
+import tifffile as tiff
 from PIL import Image
 from pandera.errors import SchemaError, SchemaErrors
 
@@ -123,12 +124,24 @@ def read_image(self, path: Path) -> np.ndarray:
         """
         if not path.exists():
             raise FileNotFoundError(f"File not found: {path}")
+        
         try:
             with Image.open(path) as img:
                 img = img.convert(self.pil_image_mode)
                 arr = np.asarray(img)
-        except Exception as e:
-            raise RuntimeError(f"Error reading image from {path}: {e}") from e
+        
+        except Exception as e_pil:
+            
+            try:
+                # fallback to tifffile for potentially more robust reading of TIFF
+                arr = tiff.imread(path)
+                
+            except Exception:
+
+                raise RuntimeError(
+                    f"Error reading image from {path}: {e_pil}"
+                ) from e_pil
+
 
         if arr.ndim not in (2, 3):
             raise ValueError(f"Unsupported image shape {arr.shape} from {path}")
@@ -149,8 +162,18 @@ def get_image_dimensions(
             if not Path(path).exists():
                 dims.append(None)
             else:
-                with Image.open(path) as img:
-                    dims.append(img.size)
+                try:
+                    with Image.open(path) as img:
+                        dims.append(img.size)
+                except Exception as e_pil:
+                    try:
+                        with tiff.TiffFile(path) as tif:
+                            page = tif.pages[0]
+                            dims.append((page.imagelength, page.imagewidth))
+                    except Exception:
+                        raise RuntimeError(
+                            f"Error reading image dimensions from {path}: {e_pil}"
+                        ) from e_pil
 
         return tuple(dims)
     

From 1fbed99e76211db6131bf6d7d7af79e6d6786f3a Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 22 Jul 2026 13:03:09 -0600
Subject: [PATCH 6/6] Update CHANGELOG.md for version 0.4.6: Add MLflow
 enhancements, channel-specific normalization, multi-channel visualization
 support, and refactor crop generation and logger functionalities.

---
 CHANGELOG.md | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3f4fa68..140d4d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,34 @@ All notable chagnes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),  
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+---
+
+## [0.4.6] - 2026-07-22
+
+### Added 
+
+#### MLflow auto logging enhancements (`virtual_stain_flow/vsf_logging/`):
+- The logger now records model optimizer parameters.
+
+#### Channel-specific normalization/transformation specification (`virtual_stain_flow/transforms/`):
+- Allows separate normalizations applied to different channels.
+
+#### Support for visualization of multi-channel input/target/prediction (`virtual_stain_flow/evaluation/`):
+- Displays channels as additional columns in visualization grid.
+- Allows selection of which channels to display vias indexing.
+
+### Refactored
+
+#### Crop generation module as subpackage (`virtual_stain_flow/datasets/ds_engine/crop_generators/`):
+- Reduce size of module by breaking into subpackage.
+- Added tiling crop.
+
+#### Logger auto-logging functionalities isolated (`virtual_stain_flow/transforms/`):
+- Moved auto-logging of optimizer, model and loss group configs outside to reduce size of main logger module and facilitate testing.
+
+#### Data access during plotting (`virtual_stain_flow/evaluation/`):
+- Now images are retrieved from dataset and uniformly normalized as (N, C, H, W) before passed to plotting functions.
+
 
 ---