From fe6770024d179ddaab8507745b96ce935a94aebc Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 18 Sep 2026 16:41:55 -0600 Subject: [PATCH 1/3] Add markdown checks to pre-commit configuration --- .pre-commit-config.yaml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 55c775a..1c2a2aa 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,3 +30,13 @@ repos: .*nbconverted.*| .*\.ipynb$ ) + +- repo: https://github.com/CU-DBMI/onesentence + rev: v0.1.1 + hooks: + # run checks + - id: check + types: [markdown] + # run checks and fixes where possible + - id: fix + types: [markdown] From 0d576a4bdd0d1b8e9c92d3bbf438b820c17c3faa Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 18 Sep 2026 16:45:47 -0600 Subject: [PATCH 2/3] Add display normalization utils to support several ways of scaling target and predicted images to faciliate effective comparison of prediction to ground truth. --- .../evaluation/display_normalization.py | 97 ++++++++++++++++ .../evaluation/visualization.py | 106 ++++++++++++------ .../evaluation/visualization_utils.py | 8 +- tests/evaluation/test_visualization_utils.py | 55 +++++++++ 4 files changed, 231 insertions(+), 35 deletions(-) create mode 100644 src/virtual_stain_flow/evaluation/display_normalization.py 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..0d3a0af --- /dev/null +++ b/src/virtual_stain_flow/evaluation/display_normalization.py @@ -0,0 +1,97 @@ +"""Display-only intensity scaling; never modifies image or metric data.""" + +from typing import List, Optional, Sequence, Tuple, Union + +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) -> Tuple[float, float]: + finite = images[np.isfinite(images)] + if not finite.size: + raise ValueError("Cannot determine display limits from an image/channel with no finite values.") + return float(finite.min()), float(finite.max()) + + +def image_norms( + images: np.ndarray, + *, + scope: str = "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] + lower, upper = _finite_range(reference) + if other is not None: + comparison = other[:, channel] if scope == "channel" else other[row, channel] + other_lower, other_upper = _finite_range(comparison) + lower, upper = min(lower, other_lower), max(upper, other_upper) + 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: str, 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..491f9e6 100644 --- a/src/virtual_stain_flow/evaluation/visualization.py +++ b/src/virtual_stain_flow/evaluation/visualization.py @@ -12,6 +12,7 @@ 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: str = "target", + target_scale_scope: str = "image", + target_limits: Optional[DisplayLimits] = None, + input_scaling: str = "independent", + input_limits: Optional[DisplayLimits] = None, + raw_scaling: str = "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), torch.from_numpy(targets)) + 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..c500cb6 100644 --- a/src/virtual_stain_flow/evaluation/visualization_utils.py +++ b/src/virtual_stain_flow/evaluation/visualization_utils.py @@ -33,7 +33,8 @@ 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}." ) - return value + # Freeze reused NumPy/tensor buffers before accessing the next sample. + return value.copy() def _stack_images(images: List[np.ndarray], name: str) -> np.ndarray: @@ -73,7 +74,10 @@ 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, including its raw image metadata. + 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. diff --git a/tests/evaluation/test_visualization_utils.py b/tests/evaluation/test_visualization_utils.py index f842151..bf2b3eb 100644 --- a/tests/evaluation/test_visualization_utils.py +++ b/tests/evaluation/test_visualization_utils.py @@ -43,3 +43,58 @@ 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]) + + +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)] + + +@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) From d8a295378e8f2c401fddf2655feaa920743da508 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci-lite[bot]" <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 22:48:40 +0000 Subject: [PATCH 3/3] [pre-commit.ci lite] apply automatic fixes --- CHANGELOG.md | 367 +++---------------- src/virtual_stain_flow/transforms/README.md | 115 ++++-- src/virtual_stain_flow/vsf_logging/README.md | 138 +------ 3 files changed, 132 insertions(+), 488 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0101f68..4416e96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,343 +1,62 @@ -# Changelog +# Datasets Subpackage Documentation -All notable changes to this project will be documented in this file. +This subpackage provides the infrastructure for managing and loading image datasets in a lazy-loaded and efficient manner, and materialized dataset classes. +It is designed to handle datasets with multiple channels and fields of view (FOVs), supporting dynamic image loading and caching to optimize memory usage and access speed. -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). +## Overview -## [0.4.10] - 2026-09-11 +The subpackage consists of three main components: -### Added - -#### Backward-compatibility config flag for ConvNeXtUNet (`virtual_stain_flow/models/unext.py`) -- Added `_pixel_shuffle_preserve_channels` to ConvNeXtUNet init/config serialization. -- `from_config` now defaults `_pixel_shuffle_preserve_channels=True` when loading older configs that do not include this key, preserving legacy behavior for previously trained models. - -### Fixed - -#### Output channel calculation in PixelShuffle2DUpBlock (`virtual_stain_flow/models/blocks/up_down_blocks.py`) -- Corrected default output-channel inference for pixel-shuffle upsampling. -- Output channels now reduce with the spatial expansion factor (for 2D: `in_channels // scale_factor^2`) instead of being preserved by default. -- Added a validation error when channel reduction would produce fewer than 1 output channel. -- Added `preserve_channels` to keep the previous behavior when needed for compatibility. - -#### Stage spatial-shape propagation for upsampling blocks (`virtual_stain_flow/models/stages.py`) -- Fixed `Stage.out_h` and `Stage.out_w` to apply shape transforms from both stage blocks, not only `Conv2DDownBlock`. -- This ensures correct output-shape reporting for stages that use upsampling blocks such as pixel shuffle and transposed convolution. - -#### Tests for PixelShuffle2DUpBlock behavior (`tests/models/test_up_down_blocks.py`) -- Updated expected output-channel assertions to match corrected pixel-shuffle defaults. -- Added tests for compatibility mode (`preserve_channels=True`). -- Added a regression test asserting that insufficient `in_channels` raises a clear `ValueError`. - -#### Added clipping to [0,1] range in max scale normalization (`virtual_stain_flow/transforms/normalizations.py`) -- Good to have for the sake of ensuring post normalization values fall within expected ranges even if normalization factor is misspecified. - ---- - -## [0.4.9] - 2026-09-11 - -### Added - - -#### Initial training resume support (`virtual_stain_flow/vsf_logging/`, `virtual_stain_flow/trainers/`) -- Added MLflow logging for optimizer state to support resuming training runs from logged artifacts and trainer state. - -### Fixed - -#### Pandas validation API update (`virtual_stain_flow/datasets/ds_engine/`) -- Replaced deprecated `DataFrame.applymap` usage with `DataFrame.map` in dataset input validation to keep the GH Actions test suite green. - - - ---- - -## [0.4.8] - 2026-08-03 - -### Fixed - -#### Best model updating/saving behavior (`virtual_stain_flow/trainers/`) -- Fixes best model updates storing only a shallow copy of statedict of models as opposed to immutable statedict. -- Fixes best model saving happens regardless of absent validation set. -In previous versions if the valiadtion set is absent from training and best model cannot possibly be determined, the trainer silently updates the `.best_model` property by the most recent model weights from training and that gets saved by the logger as the "best model". -With this fix the `trainer.best_model` will return `None` if best model cannot be determined due to missing of validation set, and correspondingly the logger will no longer write best model artifacts if not available. - -### Refactored - -#### Early stopping logic (`virtual_stain_flow/trainers/`) -- The early stopping related functionalities are now moved outside of `AbstractTrainer` to standalone utility classes to facilitate unit testing and reduce size of single module. - ---- - -## [0.4.7] - 2026-07-28 - -### Fixed - -#### Package installation discovery, expose previously unexposed subpackage - ---- - -## [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. - - ---- - -## [0.4.5] - 2026-04-30 - -### Added - -#### MLflow auto logging enhancements (`virtual_stain_flow/vsf_logging/`): - -- The logger now records model architecture tags by capturing each model config's `class_path` and setting `model..class_path` at train start. -- The loss-group auto logging routine `_log_loss_groups_config_and_tags` logs loss item names and weights as MLflow tags and persists the full loss group configuration as a JSON config artifact. - ---- - -## [0.4.4] - 2026-04-23 - -### Added - -#### Dataset wrapper and MONAI augmentation adapter (`virtual_stain_flow/datasets/`): - -Introduces a lightweight wrapper abstraction for dataset composition and a MONAI-compatible adapter for dictionary-based augmentation pipelines. This enables augmentation workflows to be layered on top of existing dataset implementations without modifying core dataset logic. - -- **`BaseWrapperDataset`** (`base_wrapper_dataset.py`): Abstract wrapper base class that forwards dataset access to an underlying dataset instance and provides recursive access to the original base dataset via the `original` property. Establishes a reusable pattern for composing dataset behaviors (e.g., augmentation, caching, preprocessing) while preserving compatibility with existing dataset APIs. -- **`MonaiAdapter`** (`monai_aug_adapter_dataset.py`): Wrapper dataset that adapts `(input, target)` tuple samples into MONAI dictionary format (`{"input": ..., "target": ...}`), applies optional MONAI `Compose` transforms, and returns transformed samples back as `(input, target)` tuples for trainer compatibility. - -#### MONAI augmentation usage example (`examples/`): - -- Added/updated **`4.data_augmentation_example.ipynb`** demonstrating: - - construction of a base dataset and crop dataset, - - application of MONAI dictionary transforms through `MonaiAdapter`, - - visualization of repeated stochastic augmentations, - - integration of the augmented dataset into a standard training dataloader/trainer workflow. - -### Refactored - -#### Visualization suite to support `BaseWrapperDataset` - ---- - -## [0.4.3] - 2025-12-16 - -### Added - -#### Crop dataset (`virtual_stain_flow/datasets/`): - -Allows the dataset to return user specified crops dynamically obtained from the full images. Supports serialization and reserialization to facilitate reproducibility. - -- **`CropImageDataset`** (`crop_dataset.py`): Dataset class for serving image crops based on a `CropManifest`. Extends `BaseImageDataset` with crop-specific state management and lazy loading via `CropFileState`. -- **`CropManifest`** (`ds_engine/crop_manifest.py`): Immutable collection of crop definitions wrapping a `DatasetManifest` for file access. Supports serialization/deserialization and factory construction from coordinate specifications. -- **`Crop`** (`ds_engine/crop_manifest.py`): Dataclass defining a single crop region with manifest index, position (x, y), and dimensions (width, height). -- **`CropIndexState`** (`ds_engine/crop_manifest.py`): Mutable state tracker for the currently active crop region. -- **`CropFileState`** (`ds_engine/crop_manifest.py`): Lazy image loading backend that wraps `FileState` to load full images and dynamically extract crop regions on demand. - -### Removed - -#### All obselete dataset classes - ---- - -## [0.4.2] - 2025-11-17 - -### Added - -#### training infrastructure (`virtual_stain_flow/engine/...`) -- Abstract away the forward pass and multiple loss accumulation from trainers -#### logging trainer (`virutal_stain_flow/trainers/logging_trainer.py`) -- New logging trainer for single generator model training using `engine` - -### Refactors - -#### abstract trainer (`virutal_stain_flow/trainers/AbstractTrainer.py`) -- Add progress bar - -### Removes - -#### Obselete trainer classes (`virutal_stain_flow/trainers/logging_trainers/...`) - ---- - -## [0.4.1] - 2025-10-16 - -### Added - -#### Loss Computation Infrastructure -- Added `LossGroup` class to abstract out the complexity of computating multiple losses surrounding a single forward pass iteration from the trainer. -Itended as prepwork for incorporating more complex wGAN training. - -##### Components: -- **`LossItem`** (`loss_group.py`): Wrapper around a `torch.nn.Module` loss to specify weights and arguments needed for computing. -- **`LossGroup`** (`loss_group.py`): Container class organizing all `LossItem`s to be computed during the same forward pass on the same set of context objects. - ---- - -## [0.4.0] - 2025-09-07 - -### Added - -#### Dataset Infrastructure Refactoring (`datasets`) -- Introduced a comprehensive refactoring of the dataset infrastructure for improved modularity, lazy loading, and memory efficiency. - -##### Core Components: -- **`DatasetManifest`** (`manifest.py`): Immutable manifest class that defines the structure of a dataset, holding a file index DataFrame where each row corresponds to a sample/FOV and columns represent channels. Validates file paths and PIL image modes during initialization. -- **`IndexState`** (`manifest.py`): Lightweight tracker for maintaining the last accessed dataset index. -- **`FileState`** (`manifest.py`): Lazy loading backend that manages image loading with configurable LRU caching. -- **`BaseImageDataset`** (`base_dataset.py`): PyTorch-compatible dataset class built on the manifest infrastructure. - -### Refactored -- Restructured dataset loading logic to use the new modular manifest-based architecture. -- Improved error handling and validation throughout the dataset pipeline. -- Enhanced type annotations and documentation for better developer experience. - ---- - -## [0.3.0] - 2025-08 - -### Added - -#### New modular backbone for `models` subpackage -##### Major changes: -- Introduced a new modular and extensible `models` subpackage for building image-to-image translation models. -The subpackage is designed around a declarative style for creating U-Net-like architectures, with a hierarchy of abstractions: - - **Blocks** (`blocks.py`, `up_down_blocks.py`): Smallest modular units, categorized into computational blocks (e.g., `Conv2DNormActBlock`, `Conv2DConvNeXtBlock`) and spatial dimension altering blocks (e.g., `Conv2DDownBlock`, `PixelShuffle2DUpBlock`). - - **Stages** (`stages.py`): Sequences of blocks for downsampling or upsampling, such as `DownStage` and `UpStage`. - - **Encoder** (`encoder.py`): Implements the downsampling path of U-Net-like architectures using `DownStage` objects. - - **Decoder** (`decoder.py`): Implements the upsampling path with skip connections using `UpStage` objects. - - **BaseModel** and **BaseGeneratorModel** (`model.py`): Added abstract base classes for models, including functionality for saving weights, configuration handling, and defining the forward pass. - - **UNet** (`unet.py`): Predefined model class supporting fully convolutional and maxpooling-based U-Net variants. - - **UNeXt** (`unext.py`): Predefined U-Net variant with a ConvNeXtV2_tiny encoder and customizable decoder. -- Added utility functions for normalization layers, activation functions, and type checking of block handles and configurations. -- Refer to the `models` README for detailed explanations of components and usage examples. - -### Refactored - -#### Repository Restructuring -- Restructured the repository from a flat layout to the conventional `/src/package_name/` structure. -This change improves module discoverability, aligns with modern Python packaging standards, and reduces potential import conflicts. -All package-related code now resides under the `src/virtual_stain_flow/` directory. -- Updated import paths throughout the codebase to reflect the new structure. -- Adjusted setup scripts and documentation to accommodate the restructuring. +1. **`DatasetManifest`**: Defines the immutable structure of a dataset, including file paths and image modes. +2. **`FileState`**: Defines the backend memory-efficient lazy loading infranstructure bridging the `DatasetManifest` and the image dataset class. +3. **`BaseImageDataset`**: A PyTorch-compatible dataset class that uses `DatasetManifest` and `FileState` to manage image loading and caching. --- -## [0.2.0] - 2025-06-19 - -### Added - -#### Overhaul Phase 1/? - Logging Framework Overhaul -A minimal rework of the logging framework as the first step to a complete overhual of the `virtual_stain_flow` software. - -This version defines a new logging subpackage that better integrates MLflow into the virtual staining model training process for a more comprehensive logging framework. +## `DatasetManifest` -#### Introduced `logging.MlflowLogger` class -Notes: This class is simiar to the old `virtual_stain_flow.callback.MlflowLogger` class, but promoted to be an independent logger class, with ability to accept logger callbacks. Key design/functionality: -- Files/metrics/parameters produced by Logger callbacks gets automatically logged to MLflow appropriately instead of being independent products untracked. -- Included some more pre-defined fine-grained run logging tags such as `experiment_type`, `model_architecture`, `target_channel_name`, `description` as logger class parameter. -- Has a `bind_trainer` and `unbind_trainer` methods to bind and unbind with the trainer instance during train step. -- User controlled mlflow run cycle, no longer autoamtically ends with the train loop, so user can perform additional logging operation before explicity ending the run. -- Has exposed `log_artifact`, `log_metric`, and `log_param` methods for manual logging of artifacts, metrics, and parameters. -- Has some access point of trainer attributes for use by logger callbacks, but subject to optimization/change. +The `DatasetManifest` class is responsible for defining the structure of a dataset. +It holds a file index (a `pandas.DataFrame`) where each row corresponds to a sample or FOV, and columns represent channels associated with that sample. +It also specifies the image mode to use when reading images. -#### Introduced `trainers.trainerAbstractLoggingTrainer` class -Notes: This class subclasses the `virtual_stain_flow.trainers.AbstractTrainer` class and preserves most of its behavior and functionalities. -Design/functionality change include: -- Binding of logger class moved from initialization to `train` method to reflect the design that logger instances should live with the training sessions. -- Early termination mode is now a parameter of the class to allow for selection of min/max optimzation mode. -- The `train` method loop invokes the logger life cycle methods: - - `logger.on_train_start()` - - `logger.on_epoch_start()` - - `logger.on_epoch_end()` - - `logger.on_train_end()` methods which in turn leads to logger's invocation of the logger callback methods. -- Requires child classes to implement the @abstract `save_model` method for unified handle for saving model weights that can be called by the logger. +## `FileState` -#### Introduced `trainers.LoggingTrainer` class -Notes: This class is nearly identical to the old Trainer class, except that: -- It is the realization of the new `AbstractLoggingTrainer` class instead of the `AbstractTrainer` class. -- It overrides the parent class `save_model` method that defines saving of the model weight. - -#### Introduced `logging.callbacks` subpackage -Notes: This is a new subpackage that is distinct from the existing `virtual_stain_flow.callbacks` -subpackage in that classes under this subpackage are passed to `logging.MlflowLogger` instances as -opposed to a `trainers.*` instances. - -##### The subpackage currently contains: -- `AbstractLoggerCallback` class: A newly introduced abstract class that defines behavior for logger -callbacks interacting with the `logging.MlflowLogger` class so product of callback gets logged appropriately as artifacts/metrics/parameters. -- `PlotPredictionCallback` class: A newly introduced class that is a realization of the `AbstractLoggerCallback` class. Serves as an example implementation of a logger callback. Similar to the `virtual_stain_flow.callbacks.intermediatePlotCallback`, plots predictions of the model on a subset of the dataset, but the additional interface with the `MlflowLogger` class ensures the -plots produced are logged as mlflow artifacts. - -### Refactored -- Internal function renames for clarity. -- Consistent attribute/property usage. -- Updated `__init__.py` directly exposing classes under subpackage. +The `FileState` class wraps a constructred `DatasetManifest` object and handles the loading of images from filepaths on-demand. +The dataset class should call the `FileState.update(idx, input_keys, target_keys)` method to request for the images corresponding to the `idx`th sample and the input/target keys to be loaded, stacked, and accessible from `FileState.input_image_raw` and `FileState.target_image_raw` attributes. --- -## [0.1.0] - 2025-03-03 - -### Added - -#### Core Framework -- Introduced a minimal yet self-contained virtual staining framework structured around modular components for model training, dataset handling, transformations, metrics, and logging. - -#### Models (`models`) -- Added `FNet`: Fully convolutional encoder-decoder for image-to-image translation. -- Added `UNet`: U-Net variant using bilinear interpolation for upsampling. -- Added GaN discriminators: - - `PatchBasedDiscriminator`: Outputs a probability map. - - `GlobalDiscriminator`: Outputs a global scalar probability. +## `BaseImageDataset` -#### Transforms (`transforms`) -- `MinMaxNormalize`: Albumentations transform for range-based normalization. -- `ZScoreNormalize`: Albumentations transform for z-score normalization. -- `PixelDepthTransform`: Converts between image bit depths (e.g., 16-bit to 8-bit). +The `BaseImageDataset` class builds on `DatasetManifest` to provide a PyTorch-compatible dataset. +It supports lazy loading of images, caching, and efficient handling of input and target channels. +- Returns paired input/target image stack as numpy arrays or torch Tensors +- Provides methods to save and load dataset configurations as JSON files for reproducibility. -#### Datasets (`datasets`) -- `ImageDataset`: Dynamically loads multi-channel microscopy images from a PE2LoadData-formatted CSV; supports input/target channel selection and Albumentations transforms. -- `PatchDataset`: Extends `ImageDataset` with configurable fixed-size cropping; supports object-centric patching and state retrieval (e.g., patch coordinates). -- `GenericImageDataset`: A simplified dataset for user-formatted directories using regex-based site/channel parsing. -- `CachedDataset`: Caches any of the above datasets in RAM to reduce I/O and speed up training. +### Usage: -#### Losses (`losses`) -- `AbstractLoss`: Base class defining standardized loss interface and trainer binding. -- `GeneratorLoss`: Combines image reconstruction and adversarial loss for training GaN generators. -- `WassersteinLoss`: Computes Wasserstein distance for GaN discriminator training. -- `GradientPenaltyLoss`: Adds gradient penalty to improve discriminator stability. +```python +import pandas as pd -#### Metrics (`metrics`) -- `AbstractMetrics`: Base class for accumulating, aggregating, and resetting batch-wise metrics. -- `MetricsWrapper`: Wraps `torch.nn.Module` metrics with accumulation and aggregation logic. -- `PSNR`: Computes Peak Signal-to-Noise Ratio (PSNR) for image quality evaluation. +from base_dataset import BaseImageDataset -#### Callbacks (`callbacks`) -- `AbstractCallback`: Base class for trainer-stage hooks (`on_train_start`, `on_epoch_end`, etc.). -- `IntermediatePlot`: Visualizes model inference during training. -- `MlflowLogger`: Logs trainer metrics and losses to an MLflow server. +# Example file index +file_index = pd.DataFrame({ + "input_channel": ["path/to/input1.tif", "path/to/input2.tif"], + "target_channel": ["path/to/target1.tif", "path/to/target2.tif"] +}) -#### Training (`trainers`) -- `AbstractTrainer`: Defines a modular training loop with support for custom models, datasets, losses, metrics, and callbacks. Exposes extensible hooks for batch and epoch-level logic. +dataset = BaseImageDataset( + file_index=file_index, + pil_image_mode="I;16", + input_channel_keys="input_channel", + target_channel_keys="target_channel", + cache_capacity=10 +) +``` +### Serialization for logging +```python +ds_config = dataset.to_config() +# or +dataset.to_json_config('loggable_artifact.json') +``` diff --git a/src/virtual_stain_flow/transforms/README.md b/src/virtual_stain_flow/transforms/README.md index 441cbce..2d5ac12 100644 --- a/src/virtual_stain_flow/transforms/README.md +++ b/src/virtual_stain_flow/transforms/README.md @@ -1,39 +1,98 @@ -# Transform Subpackage Documentation +# `virtual_stain_flow.models` -This subpackage provides image transformation utilities for virtual staining workflows, handling preprocessing, augmentation, and standardization of microscopy images. -To facilitate reproducible experiments, the transformations are also made serializable and loggable. +## Overview ---- +The `models` subpackage defines a modular and extensible framework for building image-to-image translation models. +For now most of its functionality is built around enabling a "declarative" style creation of U-Net-like architectures, +with a hierarchy of abstractions, including blocks, stages, encoders, decoders. -## Overview +The subpackage also includes some ready to use image-to-image tranlation models, +including fully convolutional or maxpooling/bilinear `UNet` and the `ConvNeXt` blocks based `ConvNeXtUNet` with tunable number of compute blocks. + +## Key Components/Building-blocks of the Framework + +![UML Diagram](assets/uml/all_models/all_models.svg) +> Figure 1. The UML diagram describes the aggregation into image-image translation models of the 3 levels of abstractions: +1. Block: the smallest modular component, performs spatial dimension preserving (computing) or changing (up/down-sampling) operations. +2. Stage: Collection of up/down sampling Blocks with repetitions compute blocks, manages the channel dimension matching between these blocks. +3. Encoder/Decoder Branch: Collection of Stages in sequence. Collects and returns skip-connection outputs if applicable. + +Diagram is generated manually. + +### Blocks +Blocks are the smallest units in the framework, responsible for directly processing feature map tensors. +They can be categorized into: + +1. **Computational Blocks**: Preserve spatial dimensions while potentially altering the number of channels. Inherits directly from `AbstractBlock`. + + Implementations include: + - `Conv2DNormActBlock`: A computational block that applies a 2D convolution followed by normalization and activation. + This block is commonly used for feature extraction and transformation while preserving spatial dimensions. + - `Conv2DConvNeXtBlock`: A computational block inspired by the ConvNeXt architecture. + It applies a series of convolutions, normalization, and activation functions designed to enhance feature extraction and representation. + This block is particularly useful for capturing complex patterns in image data while maintaining spatial dimensions. + +2. **Spatial Dimension Altering Blocks**: Change spatial dimensions, such as downsampling or upsampling. +Inherits from `AbstractUpBlock` or `AbstractDownBlock`. + + Implementations include: + - `IdentityBlock`: No-op block for first layer of UNets. + - `Conv2DDownBlock`: Down-samples image by a factor of 2. Followed by an optional Normalization layer. + - `MaxPool2DDownBlock`: Down-samples image by a factor of 2. + - `ConvTrans2DUpBlock`: Up-samples image by a factor of 2. Followed by an optional Normalization layer. + - `PixelShuffle2DUpBlock`: Up-samples image by a factor of 2. + +### Stages +Stages are collections of blocks, defining an ordered sequence of block operations, such as downsampling or upsampling followed by dimension preserving compute operation. + +A `Stage` class accepts a `in_block` handle and a `comp_block` handle, +and dynamically creates a sequence of blocks with compatible input/output sizes. +Although the `Stage` class does not enforce a specific checks for this, +it is recommended to initialize a `Stage` class with a **Spatial Dimension Altering** `in_block` and a **Computational** `comp_block`. + +The subpackage provides additional Stage classes, that still can be flexibly initialized with custom combination of `in_block` and `comp_block`, and with predefined behavior: +- `DownStage`: Handles downsampling operations from a input feature map tensor. +- `UpStage`: Handles upsampling operations with skip connections support, +processes the input feature map through a upsampling `in_block`, performs concatenation with skip connections if applicable, and finally compute with `comp_block`. + +### Encoder +The `Encoder` class implements the downsampling path of a U-Net-like architecture. +It consists of multiple `DownStage` objects, each reducing the spatial dimensions of the input tensor while capturing hierarchical features. + +### Decoder +The `Decoder` class implements the upsampling path of a U-Net-like architecture. +It consists of multiple `UpStage` objects, each increasing the spatial dimensions of the input tensor while incorporating skip connections from the encoder. + +### Model +The `model.py` module provides the foundational classes for defining models within the framework. +It includes the following key components: -This subpackage consists of the following modules: -1. **`transform_utils.py`**: Contains type definitions and validation utilities for transform objects, defining acceptable transform types and providing runtime type checking capabilities for both standard Albumentations transforms and custom LoggableTransform classes. +1. **BaseModel**: + - An abstract base class that all models inherit from. + - Provides common functionality such as saving model weights, converting configurations to/from dictionaries, and defining the forward pass. + - Enforces implementation of `forward`, `to_config`, and `from_config` methods in subclasses. -2. **`base_transform.py`**: Defines the abstract `LoggableTransform` base class that extends Albumentations' `ImageOnlyTransform`, adding serialization capabilities, naming conventions, and standardized logging interfaces for scientific reproducibility. +2. **BaseGeneratorModel**: + - Extends `BaseModel` and serves as a base class for image-to-image translation models. + - Includes properties for input/output channels, activation functions, and modular components like `Encoder` and `Decoder`. + - Implements a default forward pass that connects the input convolution, encoder, decoder, and output convolution layers. -3. **`normalizations.py`**: Implements concrete normalization transforms including `MaxScaleNormalize` for scaling images to [0,1] range and `ZScoreNormalize` for statistical standardization, both inheriting from `LoggableTransform` to ensure proper integration with the package's logging and configuration systems. +These classes establish a flexible and extensible structure for building custom models, ensuring consistency and reusability across the framework. -4. **`channelwise.py`**: Provides `ChannelwiseTransform`, a wrapper that applies a distinct transform to each channel of a channel-first image. +### Model Realizations ---- +1. **UNet**: a pre-defined model class , supporting two architectural variants: + 1. Fully convolutional U-Net with `Conv2DDownBlock` and `ConvTrans2DUpBlock`. + 2. More coventional U-Net with `MaxPool2DDownBlock` and `ConvTrans2DUpBlock`. +2. **ConvNeXtUNet**: A U-Net-like architecture that leverages the ConvNeXtV2_tiny model as its encoder. This model is designed for advanced image-to-image translation tasks and supports a modular decoder with customizable upsampling and computational blocks. + - Utilizes a pre-defined ConvNeXtV2_tiny model from the `timm` library, adapted for hierarchical feature extraction with four stages. + - Offers flexibility with two types of upsampling blocks (`PixelShuffle2DUpBlock` and `ConvTrans2DUpBlock`) and two types of computational blocks (`Conv2DConvNeXtBlock` or `Conv2DNormActBlock`). -## Usage: -See examples for in context use with datasets -```python -from normalizations import MaxScaleNormalize -from channelwise import ChannelwiseTransform +This model demonstrates the extensibility of the framework by integrating state-of-the-art architectures with the modular design principles of the subpackage. -scale_transform = MaxScaleNormalize( - normalization_factor='16bit', - name="Scale16BitImages", - p=1.0 -) +## Utilities/Handle Type-checking +Includes utility functions for retrieving normalization layers (`get_norm`), +activation functions (`get_activation`), and type checking of block handles and configurations, centralizing shared operations. -channelwise_transform = ChannelwiseTransform( - transforms=[ - MaxScaleNormalize(normalization_factor='16bit'), - MaxScaleNormalize(normalization_factor='8bit') - ] -) -``` +## Example Usage +Refer to the `examples/modular_unet_example.ipynb` notebook for detailed examples of block/stage behavior, model definition. diff --git a/src/virtual_stain_flow/vsf_logging/README.md b/src/virtual_stain_flow/vsf_logging/README.md index 653143d..74f476e 100644 --- a/src/virtual_stain_flow/vsf_logging/README.md +++ b/src/virtual_stain_flow/vsf_logging/README.md @@ -1,137 +1,3 @@ -# `virtual_stain_flow.vsf_logging` -(Prototype) Renovated Logger for Centralized MLflow Logging. +Here lives the callback classes that are meant to be fed into trainers to do stuff like saving images every epoch and logging. ---- - -## Overview - -This subpackage contains a refactored and **isolated** `MlflowLogger` logger that is its independent class. -The `MlflowLogger` works in conjunction with `virtual_stain_flow.vsf_logging.callback.LoggerCallbacks` to produce and log artifacts. - ---- - -## Subpackage Structure - -```bash -virtual_stain_flow/ -└── logging/ - ├── MlflowLogger.py # Central logging controller class - └── callbacks/ - ├── LoggerCallback.py # Abstract base class for logger-aware callbacks - └── PlotCallback.py # Example implementation to generate and prediction plots that the logger logs as artifacts -``` - ---- - -## Key Concepts - -### `MlflowLogger` - -* The **Core Logger** that: - - * Binds to a specific trainer instance only during training. - * Automatically logs metrics and losses from the trainer. - * Exposes control for MLflow run lifecycle (no longer ends with the training), thus ↵ - * allowing for manual logging by user on parameters, model weights, and artifacts. - * Accepts a list of `LoggerCallback` objects to extend automated logging behavior. - -### `AbstractLoggerCallback` - -Abstract class for logger-bound callbacks with defined lifecycle hooks at these stages of training: - * `on_train_start` - * `on_epoch_start` - * `on_epoch_end` - * `on_train_end` - * batch-wise actions are not currently supported - -and return signature: - - * Metrics (e.g., `"val_loss"`) - * Parameters (e.g., dataset config) - * Artifacts (e.g., images, `.pth` files, YAML/JSON configs) - -### `PlotPredictionCallback` - -* A concrete example of `LoggerCallback` that: - - * Accepts a dataset during initialization and dynamically retrieves the model during training to generate prediction plots - * Saves them to disk and returns paths for MLflow artifact logging - * Configurable via sampling strategy, metrics, and plotting frequency - -## Example Usage -```python -# ... import dataset/model/trainer subpackage -from virtual_stain_flow.vsf_logging import MlflowLogger -from virtual_stain_flow.vsf_logging.callbacks import PlotPredictionCallback - -""" -dataset = ... -model = ... -trainer = Trainer( - model, - dataset -) -""" - -params = { - 'foo': 1, - 'bar': 2 -} - -plot_callback = PlotPredictionCallback( - name='plot_callback1', - save_path='.', - dataset=dataset, - every_n_epochs=5, - # kwargs passed to plotter - show_plot=False - ) - -logger = MlflowLoggerV2( - name='logger1', - tracking_uri='path/to/mlruns', - experiment_name='foo1', - run_name='bar1', - experiment_type='training', - model_architecture='architecture1', - target_channel_name='foobar', - mlflow_log_params_args=params, - callbacks=[plot_callback] -) - -trainer.train(logger=logger) -``` - -Expected Logging Output: -``` -mlruns/ -└── 0/ - └── / - ├── artifacts/ - │ ├── plots/ - │ │ └── epoch/ - │ │ └── plot/ - │ │ ├── epoch_5.png - │ │ ├── epoch_10.png - │ │ └── epoch_15.png - │ └── weights/ - │ └── best/ - │ └── logger1_model_epoch_20.pth - ├── meta.yaml - ├── params/ - │ ├── foo - │ └── bar - ├── metrics/ - │ ├── val_loss - │ ├── train_loss - │ ├── train_ssim - │ └── val_ssim - └── tags/ - ├── analysis_type - ├── model - ├── channel - └── run_name - -``` - -See `examples/xx` for more detail +The callback classes must inherit the abstract class.