diff --git a/CHANGELOG.md b/CHANGELOG.md index 0101f68..375a0ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,18 @@ All notable changes 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.11] - 2026-10-15 + +### Refactored + +#### Visualization normalization/joint normalization modes and checks (`virtual_stain_flow/evluation/`) +- Added configurable image scaling for displaying the evaluation grids, including: + 1. target-referenced (prediction range is clamped in reference to target vmin and vmax to expose localized over-darkening and brightening) + 2. joint (both target and reference are clamped by the pooled vim and vmax to reveal contrast differences) + 3. independent (target and reference are independently optimized to display best effort contrast, for checking if texture and details and properly reconstructed correcting for illumination and contrast differences) + 4. fixed ranges (manual free form) +- Bumped minimal python version and pinned matplotlib version due to coupling with the normalization modes. + ## [0.4.10] - 2026-09-11 ### Added diff --git a/pyproject.toml b/pyproject.toml index 628269a..fe88f12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,15 +6,15 @@ build-backend = "setuptools.build_meta" [project] name = "virtual_stain_flow" -version = "0.4.10" +version = "0.4.11" description = "For developing virtual staining models" -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ "pyarrow", "numpy", "pandas>=2.1", "seaborn", - "matplotlib", + "matplotlib>=3.10,<3.12", "scikit-learn", "opencv-python", "typing-extensions", diff --git a/src/virtual_stain_flow/evaluation/display_normalization.py b/src/virtual_stain_flow/evaluation/display_normalization.py new file mode 100644 index 0000000..b58a38a --- /dev/null +++ b/src/virtual_stain_flow/evaluation/display_normalization.py @@ -0,0 +1,105 @@ +"""Display-only intensity scaling; never modifies image or metric data.""" + +from typing import List, Optional, Sequence, Tuple, Union, Literal + +import numpy as np +from matplotlib.colors import Normalize + +DisplayLimits = Union[Tuple[float, float], Sequence[Tuple[float, float]]] + + +class _ClippedNormalize(Normalize): + """Handle a constant reference without hiding brighter predictions. + + For a zero-width range, equal values are middle gray, lower values are + black and higher values are white (with the gray colormap). + """ + + def __call__(self, value, clip=None): + if self.vmin != self.vmax: + return super().__call__(value, clip=clip) + result, is_scalar = self.process_value(value) + normalized = np.ma.array( + np.where(result.data < self.vmin, 0.0, + np.where(result.data > self.vmax, 1.0, 0.5)), + mask=np.ma.getmaskarray(result) | np.isnan(result.data), + ) + return normalized[0] if is_scalar else normalized + + +def _fixed_limits(limits: Optional[DisplayLimits], channels: int, name: str) -> np.ndarray: + """Validate a common pair or one pair per *displayed* channel.""" + try: + values = np.asarray(limits, dtype=float) + except (TypeError, ValueError) as error: + raise ValueError(f"{name} must be a (min, max) pair or one pair per displayed channel.") from error + if values.shape == (2,): + values = np.tile(values, (channels, 1)) + if values.shape != (channels, 2) or not np.all(np.isfinite(values)): + raise ValueError(f"{name} must contain finite (min, max) pairs for {channels} displayed channels.") + if np.any(values[:, 0] >= values[:, 1]): + raise ValueError(f"{name} requires min < max.") + return values + + +def _finite_range(images: np.ndarray, other: Optional[np.ndarray] = None) -> Tuple[float, float]: + finite = images[np.isfinite(images)] + other_finite = other[np.isfinite(other)] if other is not None else np.array([]) + if not finite.size: + finite = other_finite + other_finite = np.array([]) + if not finite.size: + raise ValueError("Cannot determine display limits from an image/channel with no finite values.") + lower, upper = float(finite.min()), float(finite.max()) + if other_finite.size: + lower = min(lower, float(other_finite.min())) + upper = max(upper, float(other_finite.max())) + return lower, upper + + +def image_norms( + images: np.ndarray, + *, + scope: Literal["image", "channel"] = "image", + limits: Optional[DisplayLimits] = None, + other: Optional[np.ndarray] = None, + legacy: bool = False, +) -> List[List[Normalize]]: + """Construct per-row/channel norms, optionally pooling another image stack.""" + if scope not in ("image", "channel"): + raise ValueError("Display scale scope must be 'image' or 'channel'.") + fixed = _fixed_limits(limits, images.shape[1], "Display limits") if limits is not None else None + rows: List[List[Normalize]] = [] + for row in range(images.shape[0]): + if row and (scope == "channel" or fixed is not None): + rows.append(rows[0]) + continue + norms = [] + for channel in range(images.shape[1]): + if fixed is not None: + lower, upper = fixed[channel] + else: + reference = images[:, channel] if scope == "channel" else images[row, channel] + if other is not None: + comparison = other[:, channel] if scope == "channel" else other[row, channel] + lower, upper = _finite_range(reference, comparison) + else: + lower, upper = _finite_range(reference) + norm_type = Normalize if legacy else _ClippedNormalize + norms.append(norm_type(vmin=lower, vmax=upper, clip=True)) + rows.append(norms) + return rows + + +def unpaired_norms( + images: np.ndarray, scaling: Literal["independent", "channel", "fixed"], limits: Optional[DisplayLimits], name: str +) -> List[List[Normalize]]: + """Input/raw panels have their own independent, channel-shared or fixed scale.""" + if scaling not in ("independent", "channel", "fixed"): + raise ValueError(f"{name}_scaling must be 'independent', 'channel', or 'fixed'.") + if (scaling == "fixed") != (limits is not None): + raise ValueError(f"{name}_limits must be provided if and only if {name}_scaling='fixed'.") + return image_norms( + images, scope="channel" if scaling == "channel" else "image", + limits=limits, legacy=scaling == "independent", + ) diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py index 2628d07..080f921 100644 --- a/src/virtual_stain_flow/evaluation/visualization.py +++ b/src/virtual_stain_flow/evaluation/visualization.py @@ -5,13 +5,14 @@ Supports both BaseImageDataset and CropImageDataset with optional metrics display. """ -from typing import List, Optional, Tuple, Union +from typing import List, Literal, Optional, Tuple, Union import numpy as np import pandas as pd import torch import matplotlib.pyplot as plt from matplotlib.patches import Rectangle +from torch.utils.data import TensorDataset from ..datasets.base_dataset import BaseImageDataset from ..datasets.crop_dataset import CropImageDataset @@ -19,6 +20,7 @@ from .evaluation_utils import evaluate_per_image_metric from .predict_utils import predict_image from .visualization_utils import extract_samples_from_dataset +from .display_normalization import DisplayLimits, image_norms, unpaired_norms def _select_channels( @@ -88,6 +90,13 @@ def plot_predictions_grid( input_channel_indices: Optional[List[int]] = None, target_channel_indices: Optional[List[int]] = None, prediction_channel_indices: Optional[List[int]] = None, + target_scaling: Literal["target", "joint", "independent", "fixed"] = "target", + target_scale_scope: Literal["image", "channel"] = "image", + target_limits: Optional[DisplayLimits] = None, + input_scaling: Literal["independent", "channel", "fixed"] = "independent", + input_limits: Optional[DisplayLimits] = None, + raw_scaling: Literal["independent", "channel", "fixed"] = "independent", + raw_limits: Optional[DisplayLimits] = None, ) -> plt.Figure: """ Core visualization function for visualizing grid of input/target and predictions. @@ -127,6 +136,29 @@ def plot_predictions_grid( :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. + :param target_scaling: 'target' (default) uses only the target's range for + each target/prediction pair; 'joint' uses their combined range; + 'independent' scales them separately (legacy); 'fixed' uses target_limits. + Values outside shared limits saturate at the colormap endpoints (black + and white for gray). Only display mapping changes, never data or metrics. + :param target_scale_scope: 'image' computes ranges per row/channel; + 'channel' shares each displayed channel's range across all rows. + Independent mode still keeps target and prediction ranges separate. + Fixed mode is always shared across rows. + :param target_limits: Required only with target_scaling='fixed': a finite + (min, max) pair, or one pair per selected target channel in display order. + The corresponding prediction uses the same pair. Requires min < max. + :param input_scaling: 'independent' (default), 'channel' (shared across rows + per channel), or 'fixed'. Input scaling never shares target/raw limits. + :param input_limits: Fixed limits for inputs, in selected channel order. + :param raw_scaling: Like input_scaling, independently controls raw panels. + :param raw_limits: Fixed limits for raw images, in selected channel order. + + A constant shared reference maps equal values to the colormap midpoint, + lower predictions to its low endpoint, and higher predictions to its high + endpoint. Independent image scaling retains Matplotlib's legacy constant + image behavior. Automatic ranges ignore nonfinite values; a reference with + no finite values raises ValueError. """ if inputs.ndim != 4: raise ValueError(f"Inputs must have shape (N, C, H, W), received {inputs.shape}.") @@ -188,6 +220,25 @@ def plot_predictions_grid( "Target and prediction channel counts must match for paired display." ) + if target_scaling not in ("target", "joint", "independent", "fixed"): + raise ValueError("target_scaling must be 'target', 'joint', 'independent', or 'fixed'.") + if (target_scaling == "fixed") != (target_limits is not None): + raise ValueError("target_limits must be provided if and only if target_scaling='fixed'.") + target_norms = image_norms( + targets, scope=target_scale_scope, limits=target_limits, + other=predictions if target_scaling == "joint" else None, + legacy=target_scaling == "independent", + ) + prediction_norms = ( + image_norms(predictions, scope=target_scale_scope, legacy=True) + if has_predictions and target_scaling == "independent" else target_norms + ) + input_norms = unpaired_norms(inputs, input_scaling, input_limits, "input") + raw_norms = ( + unpaired_norms(raw_images, raw_scaling, raw_limits, "raw") + if has_raw_images else [[] for _ in range(num_samples)] + ) + 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) @@ -212,11 +263,7 @@ def plot_predictions_grid( # Create figure fig_width = panel_width * num_cols fig_height = panel_width * num_samples - fig, axes = plt.subplots(num_samples, num_cols, figsize=(fig_width, fig_height)) - - # Handle single-row case where axes is 1D - if num_samples == 1: - axes = axes.reshape(1, -1) + fig, axes = plt.subplots(num_samples, num_cols, figsize=(fig_width, fig_height), squeeze=False) for row_idx in range(num_samples): raw_row = list(raw_images[row_idx]) if has_raw_images else [] @@ -231,6 +278,12 @@ def plot_predictions_grid( ] if has_predictions else target_row img_set = raw_row + input_row + target_pred_row + paired_norms = [ + norm + for target_norm, prediction_norm in zip(target_norms[row_idx], prediction_norms[row_idx]) + for norm in (target_norm, prediction_norm) + ] if has_predictions else target_norms[row_idx] + norms = raw_norms[row_idx] + input_norms[row_idx] + paired_norms if len(img_set) != num_cols: raise ValueError( @@ -242,7 +295,7 @@ def plot_predictions_grid( # Squeeze to 2D for display (handles (1, H, W) or (H, W)) img_2d = np.squeeze(img) - ax.imshow(img_2d, cmap=cmap) + ax.imshow(img_2d, cmap=cmap, norm=norms[col_idx], interpolation="nearest") # Column title only on first row if row_idx == 0: @@ -271,12 +324,8 @@ def plot_predictions_grid( # Sample a small region (e.g., top-left 10% of image) sample_size = max(1, int(min(img_2d.shape) * 0.1)) corner_region = img_2d[:sample_size, :sample_size] - # Normalize to 0-1 range for brightness check - img_min, img_max = img_2d.min(), img_2d.max() - if img_max > img_min: - normalized_brightness = (corner_region.mean() - img_min) / (img_max - img_min) - else: - normalized_brightness = 0.5 + # Use the actual display mapping, including channel-shared limits. + normalized_brightness = norms[col_idx](corner_region).mean() text_color = "black" if normalized_brightness > 0.5 else "white" ax.text( 0.02, 0.98, # Top-left corner in axes coordinates @@ -332,7 +381,7 @@ 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, + Includes all display scaling/limits options, row_label_prefix, cmap, panel_width, show_plot, wspace, hspace, raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices, raw_channel_names, input_channel_names, target_channel_names, prediction_channel_names. """ @@ -371,9 +420,9 @@ def plot_predictions_grid_from_model( Plot predictions grid by running inference on a model. Performs the following steps: - 1. Run inference on the specified dataset indices. - 2. Compute per-image metrics. - 3. Extract samples and plot using `plot_predictions_grid`. + 1. Snapshot transformed samples and raw/crop metadata in one dataset pass. + 2. Run inference and compute metrics on that same snapshot. + 3. Plot the snapshot and predictions using `plot_predictions_grid`. :param model: PyTorch model for inference. :param dataset: BaseImageDataset, CropImageDataset, or BaseWrapperDataset to visualize. @@ -385,32 +434,23 @@ 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, + Includes all display scaling/limits options, row_label_prefix, cmap, panel_width, show_plot, wspace, hspace, 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( - dataset, model, indices=indices, device=device - ) + # Capture metadata immediately after each access, not in a second traversal + # that could resample stochastic transforms or replace mutable crop state. + ( + inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names + ) = extract_samples_from_dataset(dataset, indices) + snapshot = TensorDataset(torch.from_numpy(inputs.copy()), torch.from_numpy(targets.copy())) + targets_tensor, predictions_tensor, _ = predict_image(snapshot, model, device=device) # Step 2: Compute metrics (if any) metrics_df = None if metrics: metrics_df = evaluate_per_image_metric(predictions_tensor, targets_tensor, metrics) - # 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): - raise ValueError( - "Visualization requires a single batched input tensor with shape (N, C, H, W); " - "multi-input sequences are not supported." - ) - - inputs = inputs_tensor.detach().cpu().numpy() - targets = targets_tensor.detach().cpu().numpy() predictions = predictions_tensor.detach().cpu().numpy() # Step 4: Plot diff --git a/src/virtual_stain_flow/evaluation/visualization_utils.py b/src/virtual_stain_flow/evaluation/visualization_utils.py index 8083cc8..3cfa4be 100644 --- a/src/virtual_stain_flow/evaluation/visualization_utils.py +++ b/src/virtual_stain_flow/evaluation/visualization_utils.py @@ -33,18 +33,38 @@ def _to_numpy_image(value: Union[np.ndarray, torch.Tensor], name: str) -> np.nda f"{name} must have shape (C, H, W), received shape {value.shape}." ) + # The caller must snapshot this view before accessing the next sample. 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 +class _ImageBatch: + """Snapshot images directly into batch storage, promoting dtype if needed.""" + + def __init__(self, count: int, name: str): + self.count = count + self.name = name + self.batch: Optional[np.ndarray] = None + self._dtypes: List[np.dtype] = [] + self._position = 0 + + def append(self, image: np.ndarray) -> None: + if self.batch is not None and image.shape != self.batch.shape[1:]: + raise ValueError(f"{self.name} must have a consistent (C, H, W) shape.") + + # Use source dtypes rather than pairwise promotion to match np.stack. + if image.dtype not in self._dtypes: + self._dtypes.append(image.dtype) + dtype = np.result_type(*self._dtypes) + if self.batch is None: + self.batch = np.empty((self.count, *image.shape), dtype=dtype) + elif dtype != self.batch.dtype: + promoted = np.empty(self.batch.shape, dtype=dtype) + promoted[:self._position] = self.batch[:self._position] + self.batch = promoted + + # Freeze reused NumPy/tensor buffers before accessing the next sample. + self.batch[self._position] = image + self._position += 1 def _get_channel_names( @@ -73,7 +93,12 @@ def extract_samples_from_dataset( 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. + images and crop annotations when available. Each requested sample is + accessed once and copied immediately into preallocated batch storage, + including its raw image metadata. Mixed dtypes may require reallocation + to preserve NumPy stacking promotion semantics. + Wrappers around crop datasets must keep the underlying crop metadata + synchronized with the sample they return. :param dataset: A BaseImageDataset, CropImageDataset, or BaseWrapperDataset. :param indices: Dataset indices to extract, in the displayed order. @@ -106,9 +131,9 @@ def extract_samples_from_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 + inputs = _ImageBatch(len(indices), "Dataset inputs") + targets = _ImageBatch(len(indices), "Dataset targets") + raw_images = _ImageBatch(len(indices), "Crop raw images") 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: @@ -118,7 +143,7 @@ def extract_samples_from_dataset( inputs.append(_to_numpy_image(input_image, "Dataset input")) targets.append(_to_numpy_image(target_image, "Dataset target")) - if crop_dataset is not None: + if crop_dataset is not None and raw_images is not None and patch_coords 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: @@ -129,9 +154,11 @@ def extract_samples_from_dataset( (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 + input_batch = inputs.batch + target_batch = targets.batch + raw_batch = raw_images.batch if raw_images is not None else None + # Nonempty indices guarantee that both batches have been allocated. + assert input_batch is not None and target_batch is not None return ( input_batch, diff --git a/tests/evaluation/test_display_normalization.py b/tests/evaluation/test_display_normalization.py new file mode 100644 index 0000000..7f32f5b --- /dev/null +++ b/tests/evaluation/test_display_normalization.py @@ -0,0 +1,206 @@ +"""Regression tests for the actual Matplotlib image mappings, not just titles.""" + +import matplotlib +import numpy as np +import pytest + +matplotlib.use("Agg") + +from virtual_stain_flow.evaluation.visualization import plot_predictions_grid + + +def _images(figure): + return [axis.images[0] for axis in figure.axes] + + +def _grid(**kwargs): + targets = np.array([0.2, 0.3, 0.5, 0.6]).reshape(1, 1, 2, 2) + predictions = np.array([-0.1, 0.3, 0.5, 1.2]).reshape(1, 1, 2, 2) + return plot_predictions_grid( + targets * 10, targets, predictions, show_plot=False, **kwargs + ) + + +def test_default_is_target_only_and_clips_prediction_display_not_data(): + figure = _grid() + input_image, target, prediction = _images(figure) + assert target.norm is prediction.norm + assert target.get_clim() == (0.2, 0.6) + assert input_image.get_clim() == (2.0, 6.0) + assert prediction.norm(-0.1) == 0 + assert prediction.norm(1.2) == 1 + assert prediction.cmap(prediction.norm(-0.1))[:3] == (0, 0, 0) + assert prediction.cmap(prediction.norm(1.2))[:3] == (1, 1, 1) + assert prediction.get_array().min() == -0.1 + assert prediction.get_array().max() == 1.2 + # Force rendering too, including the shared norm's interaction with imshow. + figure.canvas.draw() + assert target.get_clim() == prediction.get_clim() == (0.2, 0.6) + + +@pytest.mark.parametrize( + "mode,limits,target_range,prediction_range", + [ + ("target", None, (0.2, 0.6), (0.2, 0.6)), + ("joint", None, (-0.1, 1.2), (-0.1, 1.2)), + ("independent", None, (0.2, 0.6), (-0.1, 1.2)), + ("fixed", (0, 1), (0, 1), (0, 1)), + ], +) +def test_selectable_target_modes(mode, limits, target_range, prediction_range): + _, target, prediction = _images(_grid(target_scaling=mode, target_limits=limits)) + assert target.get_clim() == target_range + assert prediction.get_clim() == prediction_range + + +@pytest.mark.parametrize("mode", ["target", "joint", "independent"]) +def test_channel_scope_shares_rows_but_not_channels(mode): + targets = np.array([ + [[[1, 2], [3, 4]], [[10, 20], [30, 40]]], + [[[2, 3], [4, 5]], [[20, 30], [40, 50]]], + ], dtype=float) + predictions = targets * 2 + figure = plot_predictions_grid( + targets[:, :1], targets, predictions, target_scaling=mode, + target_scale_scope="channel", show_plot=False, + ) + images = _images(figure) + target_ranges = [(1, 5), (10, 50)] + prediction_ranges = [(2, 10), (20, 100)] + if mode == "joint": + target_ranges = [(1, 10), (10, 100)] + if mode != "independent": + prediction_ranges = target_ranges + for channel in range(2): + target_column = 1 + 2 * channel + assert images[target_column].norm is images[target_column + 5].norm + assert images[target_column + 1].norm is images[target_column + 6].norm + assert images[target_column].get_clim() == target_ranges[channel] + assert images[target_column + 1].get_clim() == prediction_ranges[channel] + + +def test_image_scope_keeps_rows_independent_by_default(): + targets = np.arange(8, dtype=float).reshape(2, 1, 2, 2) + images = _images(plot_predictions_grid(targets, targets, targets * 2, show_plot=False)) + assert images[1].get_clim() == images[2].get_clim() == (0, 3) + assert images[4].get_clim() == images[5].get_clim() == (4, 7) + + +def test_fixed_limits_follow_selected_channel_order(): + images = np.arange(24, dtype=float).reshape(2, 3, 2, 2) + artists = _images(plot_predictions_grid( + images, images, images, raw_images=images, + input_channel_indices=[2, 0], raw_channel_indices=[1], + target_channel_indices=[2, 0], prediction_channel_indices=[0, 1], + target_scaling="fixed", target_limits=[(0, 100), (-1, 1)], + input_scaling="fixed", input_limits=[(0, 30), (0, 10)], + raw_scaling="fixed", raw_limits=(0, 65535), show_plot=False, + )) + expected = [(0, 65535), (0, 30), (0, 10), (0, 100), (0, 100), (-1, 1), (-1, 1)] + assert [image.get_clim() for image in artists] == expected * 2 + np.testing.assert_array_equal(artists[3].get_array(), images[0, 2]) + np.testing.assert_array_equal(artists[4].get_array(), images[0, 0]) + + +@pytest.mark.parametrize("scaling", ["independent", "channel"]) +def test_input_and_raw_scaling_are_separate_from_each_other_and_targets(scaling): + images = np.arange(16, dtype=float).reshape(2, 2, 2, 2) + artists = _images(plot_predictions_grid( + images, images[:, :1] / 100, raw_images=images * 100, + input_scaling=scaling, raw_scaling=scaling, show_plot=False, + )) + if scaling == "channel": + ranges = [(0, 1100), (400, 1500), (0, 11), (4, 15)] + for start in (0, 5): + assert [image.get_clim() for image in artists[start:start + 4]] == ranges + else: + assert artists[0].get_clim() == (0, 300) + assert artists[5].get_clim() == (800, 1100) + assert artists[2].get_clim() == (0, 3) + assert artists[7].get_clim() == (8, 11) + + +def test_constant_target_still_exposes_under_and_over_prediction(): + target = np.full((1, 1, 2, 2), 0.5) + prediction = np.array([0.4, 0.5, 0.6, np.nan]).reshape(1, 1, 2, 2) + figure = plot_predictions_grid(target, target, prediction, show_plot=False) + _, target_artist, pred_artist = _images(figure) + assert pred_artist.norm is target_artist.norm + assert pred_artist.get_clim() == (0.5, 0.5) + np.testing.assert_array_equal(pred_artist.norm([0.4, 0.5, 0.6]), [0, 0.5, 1]) + assert np.ma.is_masked(pred_artist.norm(np.nan)) + mapped = pred_artist.cmap(pred_artist.norm(prediction[0, 0])) + np.testing.assert_allclose( + mapped, + [[[0, 0, 0, 1], [0.5, 0.5, 0.5, 1]], + [[1, 1, 1, 1], [0, 0, 0, 0]]], + atol=1 / 255, + ) + + +def test_legacy_constant_images_keep_matplotlib_behavior(): + images = np.ones((1, 1, 2, 2)) + artists = _images(plot_predictions_grid( + images, images, images, target_scaling="independent", show_plot=False, + )) + assert artists[1].norm(1.0) == 0 + + +def test_auto_limits_ignore_nonfinite_values(): + targets = np.array([np.nan, np.inf, 0.2, 0.8]).reshape(1, 1, 2, 2) + artists = _images(plot_predictions_grid(np.zeros_like(targets), targets, show_plot=False)) + assert artists[1].get_clim() == (0.2, 0.8) + with pytest.raises(ValueError, match="no finite values"): + plot_predictions_grid(np.zeros_like(targets), targets * np.nan, show_plot=False) + + +def test_joint_limits_use_finite_target_when_prediction_channel_is_all_nonfinite(): + targets = np.array([np.nan, np.inf, 0.2, 0.8]).reshape(1, 1, 2, 2) + predictions = np.full_like(targets, np.nan) + artists = _images(plot_predictions_grid( + np.zeros_like(targets), targets, predictions, + target_scaling="joint", show_plot=False, + )) + assert artists[1].get_clim() == artists[2].get_clim() == (0.2, 0.8) + + predictions = np.array([-0.1, np.inf, 0.5, np.nan]).reshape(1, 1, 2, 2) + artists = _images(plot_predictions_grid( + np.zeros_like(targets), np.full_like(targets, np.nan), predictions, + target_scaling="joint", show_plot=False, + )) + assert artists[1].get_clim() == artists[2].get_clim() == (-0.1, 0.5) + + predictions = np.full_like(targets, np.nan) + with pytest.raises(ValueError, match="no finite values"): + plot_predictions_grid( + np.zeros_like(targets), predictions, predictions, + target_scaling="joint", show_plot=False, + ) + + +@pytest.mark.parametrize("kwargs", [ + {"target_scaling": "invalid"}, + {"target_scale_scope": "invalid"}, + {"target_scaling": "fixed"}, + {"target_limits": (0, 1)}, + {"target_scaling": "fixed", "target_limits": (1, 0)}, + {"target_scaling": "fixed", "target_limits": (0, 0)}, + {"target_scaling": "fixed", "target_limits": (0, np.inf)}, + {"target_scaling": "fixed", "target_limits": [(0, 1), (0, 2)]}, + {"input_scaling": "invalid"}, + {"input_scaling": "fixed"}, + {"input_limits": (0, 1)}, + {"raw_scaling": "invalid", "raw_images": np.ones((1, 1, 2, 2))}, +]) +def test_invalid_scaling_settings_raise(kwargs): + with pytest.raises(ValueError): + _grid(**kwargs) + + +def test_without_predictions_joint_uses_target_range_and_does_not_mutate_arrays(): + inputs = np.arange(4, dtype=float).reshape(1, 1, 2, 2) + targets = inputs / 10 + original = targets.copy() + artists = _images(plot_predictions_grid(inputs, targets, target_scaling="joint", show_plot=False)) + assert artists[1].get_clim() == (0, 0.3) + np.testing.assert_array_equal(targets, original) diff --git a/tests/evaluation/test_visualization.py b/tests/evaluation/test_visualization.py index 260f61d..193f387 100644 --- a/tests/evaluation/test_visualization.py +++ b/tests/evaluation/test_visualization.py @@ -79,3 +79,102 @@ def test_plot_predictions_grid_from_model_uses_tensor_batches(basic_dataset): ) assert len(figure.axes) == 8 + + +@pytest.mark.parametrize("scaling,limits", [("target", None), ("fixed", (0, 0.01)), ("independent", None)]) +def test_model_plot_uses_one_transformed_snapshot_for_metrics_and_metadata(crop_dataset, scaling, limits): + from virtual_stain_flow.datasets.base_wrapper_dataset import BaseWrapperDataset + from virtual_stain_flow.transforms.gamma import ContinuousGammaTransform + from virtual_stain_flow.transforms.normalizations import MaxScaleNormalize + + crop_dataset.input_transforms = [MaxScaleNormalize(10)] + crop_dataset.target_transforms = [MaxScaleNormalize(10), ContinuousGammaTransform(0.3)] + + class StatefulWrapper(BaseWrapperDataset): + def __init__(self, dataset): + super().__init__(dataset) + self.calls = [] + + def __getitem__(self, index): + self.calls.append(index) + inputs, targets = self._dataset[index] + # Simulate augmentation that changes with every read. + return inputs + len(self.calls), targets + + class RecordingModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.inputs = [] + + def forward(self, inputs): + self.inputs.append(inputs.clone()) + return inputs[:, :1] + 0.25 + + class RecordingL1(torch.nn.L1Loss): + def __init__(self): + super().__init__() + self.values = [] + + def forward(self, first, second): + result = super().forward(first, second) + self.values.append(result.item()) + return result + + dataset = StatefulWrapper(crop_dataset) + model, metric = RecordingModel(), RecordingL1() + figure = plot_predictions_grid_from_model( + model, dataset, indices=[3, 0], metrics=[metric], device="cpu", + target_scaling=scaling, target_limits=limits, show_plot=False, + ) + assert dataset.calls == [3, 0] + for row, (raw_value, crop_xy) in enumerate([(3, (5, 5)), (1, (0, 0))]): + axes = figure.axes[row * 6:(row + 1) * 6] + displayed_input = axes[2].images[0].get_array() + displayed_target = axes[4].images[0].get_array() + displayed_prediction = axes[5].images[0].get_array() + np.testing.assert_allclose(displayed_input, model.inputs[row][0, 0].numpy()) + np.testing.assert_allclose(displayed_input, raw_value / 10 + row + 1) + np.testing.assert_allclose(displayed_target, ((2 if row == 0 else 1) / 10) ** 0.3) + np.testing.assert_allclose(displayed_prediction, displayed_input + 0.25) + assert metric.values[row] == pytest.approx(np.abs(displayed_prediction - displayed_target).mean()) + assert f"RecordingL1: {metric.values[row]:.3f}" in axes[5].get_title() + np.testing.assert_array_equal(axes[0].images[0].get_array(), raw_value) + assert axes[0].patches[0].get_xy() == crop_xy + + +def test_dataset_wrapper_forwards_display_options(basic_dataset): + figure = plot_dataset_grid( + basic_dataset, [0, 1], target_scaling="fixed", target_limits=(0, 10), + input_scaling="channel", show_plot=False, + ) + assert figure.axes[2].images[0].get_clim() == (0, 10) + assert figure.axes[5].images[0].get_clim() == (0, 10) + assert figure.axes[0].images[0].get_clim() == figure.axes[3].images[0].get_clim() + + +def test_plot_callback_forwards_scaling_options(basic_dataset, tmp_path, monkeypatch): + from virtual_stain_flow.vsf_logging.callbacks.PlotCallback import PlotPredictionCallback + + callback = PlotPredictionCallback( + "test", basic_dataset, save_path=tmp_path, indices=[0], + target_scaling="fixed", target_limits=(0, 1), input_scaling="channel", + show_plot=False, panel_width=0.5, + ) + monkeypatch.setattr(callback, "get_model", lambda **kwargs: torch.nn.Conv2d(2, 1, 1)) + monkeypatch.setattr(callback, "get_epoch", lambda: 0) + assert callback._plot().is_file() + + +def test_inplace_model_does_not_change_displayed_snapshot(basic_dataset): + class InplaceModel(torch.nn.Module): + def forward(self, inputs): + inputs.add_(100) + return inputs[:, :1] + + figure = plot_predictions_grid_from_model( + InplaceModel(), basic_dataset, [0], [], device="cpu", show_plot=False, + ) + np.testing.assert_array_equal(figure.axes[0].images[0].get_array(), 1) + np.testing.assert_array_equal(figure.axes[1].images[0].get_array(), 2) + np.testing.assert_array_equal(figure.axes[2].images[0].get_array(), 1) + np.testing.assert_array_equal(figure.axes[3].images[0].get_array(), 101) diff --git a/tests/evaluation/test_visualization_utils.py b/tests/evaluation/test_visualization_utils.py index f842151..9d6d34d 100644 --- a/tests/evaluation/test_visualization_utils.py +++ b/tests/evaluation/test_visualization_utils.py @@ -43,3 +43,88 @@ def __getitem__(self, index): with pytest.raises(ValueError, match="single channel-first image"): extract_samples_from_dataset(MultiInputDataset(basic_dataset), [0]) + + +def test_snapshot_copies_reused_buffers(basic_dataset): + import torch + + class ReusingWrapper(BaseWrapperDataset): + def __init__(self, dataset): + super().__init__(dataset) + self.buffer = torch.empty((1, 2, 2)) + + def __getitem__(self, index): + self.buffer.fill_(index) + return self.buffer, self.buffer + + dataset = ReusingWrapper(basic_dataset) + inputs, targets, *_ = extract_samples_from_dataset(dataset, [2, 0, 2]) + dataset.buffer.fill_(-1) + for images in (inputs, targets): + np.testing.assert_array_equal(images[:, 0, 0, 0], [2, 0, 2]) + inputs.fill(10) + np.testing.assert_array_equal(targets[:, 0, 0, 0], [2, 0, 2]) + assert np.all(dataset.buffer.numpy() == -1) + + +def test_snapshot_copies_reused_raw_buffers(crop_dataset, monkeypatch): + dataset_type = type(crop_dataset) + original_getitem = dataset_type.__getitem__ + raw_buffer = np.empty((2, 10, 10)) + + def getitem(self, index): + result = original_getitem(self, index) + raw_buffer.fill(index) + return result + + monkeypatch.setattr(dataset_type, "__getitem__", getitem) + monkeypatch.setattr(dataset_type, "original_input_image", property(lambda self: raw_buffer)) + _, _, raw, coords, *_ = extract_samples_from_dataset(crop_dataset, [1, 2]) + raw_buffer.fill(-1) + np.testing.assert_array_equal(raw[:, 0, 0, 0], [1, 2]) + assert coords == [(5, 5, 4, 4), (0, 0, 4, 4)] + raw.fill(10) + assert np.all(raw_buffer == -1) + + +@pytest.mark.parametrize("dtypes", [(np.uint8, np.int8, np.float16), (np.int16, np.float32, np.float64)]) +def test_snapshot_matches_stack_dtype_promotion(basic_dataset, dtypes): + images = [np.full((1, 2, 2), index + 0.5, dtype=dtype) for index, dtype in enumerate(dtypes)] + + class MixedDtypeWrapper(BaseWrapperDataset): + def __getitem__(self, index): + return images[index], images[index] + + inputs, targets, *_ = extract_samples_from_dataset(MixedDtypeWrapper(basic_dataset), [0, 1, 2]) + expected = np.stack(images) + for batch in (inputs, targets): + assert batch.dtype == expected.dtype + np.testing.assert_array_equal(batch, expected) + + +def test_snapshot_rejects_inconsistent_shapes(basic_dataset): + class VaryingShapeWrapper(BaseWrapperDataset): + def __getitem__(self, index): + image = np.zeros((1, index + 1, 2)) + return image, image + + with pytest.raises(ValueError, match=r"Dataset inputs must have a consistent \(C, H, W\) shape"): + extract_samples_from_dataset(VaryingShapeWrapper(basic_dataset), [0, 1]) + + +@pytest.mark.parametrize("fixture_name", ["basic_dataset", "crop_dataset"]) +def test_snapshot_matches_dataloader_with_separate_transforms(request, fixture_name): + from torch.utils.data import DataLoader, Subset + from virtual_stain_flow.transforms.gamma import ContinuousGammaTransform + from virtual_stain_flow.transforms.normalizations import MaxScaleNormalize + + dataset = request.getfixturevalue(fixture_name) + dataset.transforms = [MaxScaleNormalize(10)] + dataset.input_transforms = [MaxScaleNormalize(2)] + dataset.target_transforms = [ContinuousGammaTransform(0.3)] + inputs, targets, *_ = extract_samples_from_dataset(dataset, [2, 0]) + loader_inputs, loader_targets = next(iter(DataLoader(Subset(dataset, [2, 0]), batch_size=2))) + np.testing.assert_array_equal(inputs, loader_inputs.numpy()) + np.testing.assert_array_equal(targets, loader_targets.numpy()) + np.testing.assert_allclose(inputs[1, 0], 1 / 10 / 2) + np.testing.assert_allclose(targets[1, 0], (1 / 10) ** 0.3)