diff --git a/src/virtual_stain_flow/datasets/base_dataset.py b/src/virtual_stain_flow/datasets/base_dataset.py index 6fdd20b..c3819f5 100644 --- a/src/virtual_stain_flow/datasets/base_dataset.py +++ b/src/virtual_stain_flow/datasets/base_dataset.py @@ -30,6 +30,8 @@ def __init__( input_channel_keys: Optional[Union[str, Sequence[str]]] = None, target_channel_keys: Optional[Union[str, Sequence[str]]] = None, transforms: Optional[Sequence[LoggableTransform]] = None, + input_transforms: Optional[Sequence[LoggableTransform]] = None, + target_transforms: Optional[Sequence[LoggableTransform]] = None, cache_capacity: Optional[int] = None, file_state: Optional[FileState] = None, ): @@ -49,6 +51,10 @@ def __init__( :param target_channel_keys: Keys for target channels in the file index. :param transforms: Optional sequence of LoggableTransform objects to apply to the images before returning them. + :param input_transforms: Optional sequence of LoggableTransform objects to + apply to the input image only, after `transforms`. + :param target_transforms: Optional sequence of LoggableTransform objects to + apply to the target image only, after `transforms`. :param cache_capacity: Optional capacity for caching loaded images. When set to None, default caching behavior of caching at most `file_index.shape[0]` images is used. When set to -1, unbounded @@ -86,6 +92,33 @@ def __init__( raise ValueError("All transforms must be instances of LoggableTransform.") self.transforms = transforms + self.input_transforms = self._normalize_transforms( + input_transforms, "input_transforms" + ) + self.target_transforms = self._normalize_transforms( + target_transforms, "target_transforms" + ) + + def _normalize_transforms( + self, + transforms: Optional[Sequence[LoggableTransform]], + name: str, + ) -> Sequence[LoggableTransform]: + """ + Normalize and validate a sequence of transforms. + + :param transforms: Sequence of LoggableTransform objects or None. + :param name: Name of the transform sequence for error messages. + :return: Normalized sequence of LoggableTransform objects. + """ + if not isinstance(transforms, Sequence): + transforms = [transforms] if transforms else [] + if not all(isinstance(t, LoggableTransform) for t in transforms): + raise ValueError( + f"All {name} must be instances of LoggableTransform." + ) + return transforms + def get_raw_item( self, idx: int @@ -119,28 +152,24 @@ def __len__(self) -> int: """ return len(self.manifest) - def _apply_transforms( - self, - image: np.ndarray, - ) -> np.ndarray: - """ - Applies the sequence of transforms to the input image. - - :param image: Input image as a numpy array. - :return: Transformed image as a numpy array. - """ - for transform in self.transforms: - image = transform.apply(img=image) - return image - def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: """ Overridden Dataset `__getitem__` method so class works with torch DataLoader. """ input_image_raw, target_image_raw = self.get_raw_item(idx) - return (torch.from_numpy(self._apply_transforms(input_image_raw)).float(), - torch.from_numpy(self._apply_transforms(target_image_raw)).float()) + input_image = input_image_raw + for transform in (*self.transforms, *self.input_transforms): + input_image = transform.apply(img=input_image) + + target_image = target_image_raw + for transform in (*self.transforms, *self.target_transforms): + target_image = transform.apply(img=target_image) + + return ( + torch.from_numpy(input_image).float(), + torch.from_numpy(target_image).float() + ) @property def pil_image_mode(self) -> str: diff --git a/src/virtual_stain_flow/datasets/crop_dataset.py b/src/virtual_stain_flow/datasets/crop_dataset.py index 975019f..3b1727f 100644 --- a/src/virtual_stain_flow/datasets/crop_dataset.py +++ b/src/virtual_stain_flow/datasets/crop_dataset.py @@ -2,7 +2,7 @@ crop_dataset.py """ -from typing import Any, Dict, List, Sequence, Optional, Tuple, Union, Type +from typing import Any, Dict, List, Sequence, Optional, Tuple, Union import numpy as np import pandas as pd @@ -30,6 +30,8 @@ def __init__( input_channel_keys: Optional[Union[str, Sequence[str]]] = None, target_channel_keys: Optional[Union[str, Sequence[str]]] = None, transforms: Optional[Sequence[LoggableTransform]] = None, + input_transforms: Optional[Sequence[LoggableTransform]] = None, + target_transforms: Optional[Sequence[LoggableTransform]] = None, crop_file_state: Optional[CropFileState] = None, ): """ @@ -53,6 +55,10 @@ def __init__( :param input_channel_keys: Keys for input channels in the file index. :param target_channel_keys: Keys for target channels in the file index. :param transforms: Optional sequence of transformations to apply to the images. + :param input_transforms: Optional sequence of LoggableTransform objects to + apply to the input image only, after `transforms`. + :param target_transforms: Optional sequence of LoggableTransform objects to + apply to the target image only, after `transforms`. :param crop_file_state: Optional pre-initialized CropFileState object. If provided, it takes precedence over `file_index` and `crop_specs`. Intended to be used by only .from_config class method and similar deserialization @@ -84,6 +90,13 @@ def __init__( raise ValueError("All transforms must be instances of LoggableTransform.") self.transforms = transforms + self.input_transforms = self._normalize_transforms( + input_transforms, "input_transforms" + ) + self.target_transforms = self._normalize_transforms( + target_transforms, "target_transforms" + ) + @property def pil_image_mode(self) -> str: return self.manifest.pil_image_mode @@ -157,7 +170,9 @@ def from_base_dataset( cls, base_dataset: BaseImageDataset, transforms: Optional[Sequence[LoggableTransform]] = None, - how: Type[CropGenerator] = generate_center_crops, + input_transforms: Optional[Sequence[LoggableTransform]] = None, + target_transforms: Optional[Sequence[LoggableTransform]] = None, + how: CropGenerator = generate_center_crops, **kwargs: Any ) -> 'CropImageDataset': """ @@ -166,6 +181,10 @@ def from_base_dataset( :param base_dataset: The BaseImageDataset to convert. :param how: A function that generates crop specifications from the base dataset. Default is `generate_center_crops`. + :param input_transforms: Optional sequence of LoggableTransform objects to + apply to the input image only, after `transforms`. + :param target_transforms: Optional sequence of LoggableTransform objects to + apply to the target image only, after `transforms`. :param kwargs: Additional keyword arguments for the `how` function. """ @@ -177,6 +196,8 @@ def from_base_dataset( return cls( file_index=base_dataset.file_index, transforms=transforms, + input_transforms=input_transforms, + target_transforms=target_transforms, crop_specs=crop_specs, pil_image_mode=base_dataset.pil_image_mode, input_channel_keys=base_dataset.input_channel_keys, diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generator.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generator.py index ad5cad1..406aca8 100644 --- a/src/virtual_stain_flow/datasets/ds_engine/crop_generator.py +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generator.py @@ -3,98 +3,22 @@ Utilities for generating crop coordinates from BaseImageDataset objects. Designed for easy creation of CropImageDataset instances. +Made Facade to account for increased complexity and future expansion. """ -from typing import Dict, List, Tuple, Any, Protocol - -from ..base_dataset import BaseImageDataset -from .ds_utils import ( - _get_active_channels, - _validate_same_dimensions_across_channels +from .crop_generators.protocol import CropSpec, CropMap, CropGenerator +from .crop_generators.center import ( + generate_center_crops, + _compute_center_crop, ) - -CropSpec = Tuple[Tuple[int, int], int, int] -CropMap = Dict[int, List[CropSpec]] - - -class CropGenerator(Protocol): - """ - Protocol for crop generator functions. - """ - def __call__( - self, - dataset: BaseImageDataset, - **kwargs: Any - ) -> CropMap: - pass - - -def _compute_center_crop( - image_width: int, - image_height: int, - crop_size: int -) -> Tuple[int, int]: - """ - Compute top-left (x, y) coordinates for a center crop. - - :param image_width: Width of the source image. - :param image_height: Height of the source image. - :param crop_size: Size of the square crop (width and height). - :return: Tuple of (x, y) for top-left corner of center crop. - :raises ValueError: If crop_size exceeds image dimensions. - """ - if crop_size > image_width or crop_size > image_height: - raise ValueError( - f"crop_size ({crop_size}) exceeds image dimensions " - f"({image_width}x{image_height})." - ) - - x = (image_width - crop_size) // 2 - y = (image_height - crop_size) // 2 - return x, y - - -def generate_center_crops( - dataset: BaseImageDataset, - crop_size: int, -) -> CropMap: - """ - Generate center crop coordinates for each sample in a BaseImageDataset. - - :param dataset: A BaseImageDataset instance (or compatible object with - `file_state.manifest` attribute supporting `get_image_dimensions()`). - :param crop_size: Size of the square crop (same width and height). - :return: Dictionary mapping manifest indices to lists of crop specs. - Format: {manifest_idx: [((x, y), width, height), ...]} - :raises ValueError: If crop_size is non-positive, if no active channels - are configured, or if channel dimensions don't match for any sample. - """ - if crop_size <= 0: - raise ValueError(f"crop_size must be positive, got {crop_size}.") - - active_channels = _get_active_channels(dataset) - if not active_channels: - raise ValueError( - "No active channels configured. Set input_channel_keys and/or " - "target_channel_keys on the dataset before generating crops." - ) - - manifest = dataset.file_state.manifest - crop_specs: Dict[int, List[Tuple[Tuple[int, int], int, int]]] = {} - - for idx in range(len(dataset)): - # Get dimensions for all active channels - dims = manifest.get_image_dimensions(idx, channels=active_channels) - - # Validate all channels have matching dimensions - width, height = _validate_same_dimensions_across_channels( - dims, active_channels, idx - ) - - # Compute center crop coordinates - x, y = _compute_center_crop(width, height, crop_size) - - # Store as crop_specs format: {idx: [((x, y), w, h), ...]} - crop_specs[idx] = [((x, y), crop_size, crop_size)] - - return crop_specs +from .crop_generators.point_centered import generate_point_centered_crops +from .crop_generators.tile import generate_tile_crops +__all__ = [ + "CropSpec", + "CropMap", + "CropGenerator", + "generate_center_crops", + "generate_point_centered_crops", + "generate_tile_crops", + "_compute_center_crop", +] diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/__init__.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/__init__.py new file mode 100644 index 0000000..5ff9463 --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/__init__.py @@ -0,0 +1,14 @@ +from .protocol import CropMap, CropSpec, CropGenerator +from .center import generate_center_crops, _compute_center_crop +from .point_centered import generate_point_centered_crops +from .tile import generate_tile_crops + +__all__ = [ + "CropMap", + "CropSpec", + "CropGenerator", + "generate_center_crops", + "generate_point_centered_crops", + "generate_tile_crops", + "_compute_center_crop" +] diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/center.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/center.py new file mode 100644 index 0000000..9ffb8dc --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/center.py @@ -0,0 +1,104 @@ +""" +center.py + +Helper module for generating crops centered at FOV +""" + +from typing import Dict, List, Tuple + +from .crop_summary import warn_formatted_crop_summary +from .protocol import CropMap +from ...base_dataset import BaseImageDataset +from ..ds_utils import ( + _get_active_channels, + _validate_same_dimensions_across_channels +) + + +def _compute_center_crop( + image_width: int, + image_height: int, + crop_size: int +) -> Tuple[int, int]: + """ + Compute top-left (x, y) coordinates for a center crop. + + :param image_width: Width of the source image. + :param image_height: Height of the source image. + :param crop_size: Size of the square crop (width and height). + :return: Tuple of (x, y) for top-left corner of center crop. + :raises ValueError: If crop_size exceeds image dimensions. + """ + if crop_size > image_width or crop_size > image_height: + raise ValueError( + f"crop_size ({crop_size}) exceeds image dimensions " + f"({image_width}x{image_height})." + ) + + x = (image_width - crop_size) // 2 + y = (image_height - crop_size) // 2 + return x, y + + +def generate_center_crops( + dataset: BaseImageDataset, + crop_size: int, + verbose: bool = False, +) -> CropMap: + """ + Generate center crop coordinates for each sample in a BaseImageDataset. + + :param dataset: A BaseImageDataset instance (or compatible object with + `file_state.manifest` attribute supporting `get_image_dimensions()`). + :param crop_size: Size of the square crop (same width and height). + :param verbose: If True, emit summary statistics for crop generation. + :return: Dictionary mapping manifest indices to lists of crop specs. + Format: {manifest_idx: [((x, y), width, height), ...]} + :raises ValueError: If crop_size is non-positive, if no active channels + are configured, or if channel dimensions don't match for any sample. + """ + if crop_size <= 0: + raise ValueError(f"crop_size must be positive, got {crop_size}.") + + active_channels = _get_active_channels(dataset) + if not active_channels: + raise ValueError( + "No active channels configured. Set input_channel_keys and/or " + "target_channel_keys on the dataset before generating crops." + ) + + manifest = dataset.file_state.manifest + crop_specs: Dict[int, List[Tuple[Tuple[int, int], int, int]]] = {} + total_samples = len(dataset) + + for idx in range(total_samples): + # Get dimensions for all active channels + dims = manifest.get_image_dimensions(idx, channels=active_channels) + + # Validate all channels have matching dimensions + width, height = _validate_same_dimensions_across_channels( + dims, active_channels, idx + ) + + # Compute center crop coordinates + x, y = _compute_center_crop(width, height, crop_size) + + # Store as crop_specs format: {idx: [((x, y), w, h), ...]} + crop_specs[idx] = [((x, y), crop_size, crop_size)] + + if verbose: + successful_center_crops = len(crop_specs) + warn_formatted_crop_summary( + title="Center crop generation statistics:", + detail_line=( + "Generation criterion: exactly one centered crop is generated " + "per sample when crop size fits within image bounds." + ), + metrics={ + "Success count / total dataset count": ( + f"{successful_center_crops}/{total_samples}" + ), + }, + ) + + return crop_specs diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/crop_summary.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/crop_summary.py new file mode 100644 index 0000000..074d4f8 --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/crop_summary.py @@ -0,0 +1,91 @@ +""" +crop_summary.py + +Helper module for formatting crop generation summaries, +shared by crop generating modules +""" + +from collections.abc import Mapping +import warnings + + +class CropSummaryWarning(UserWarning): + """Warning category for crop generation summary reports.""" + + +def _stringify_metric_value(value: object) -> str: + """Convert metric values to compact, human-readable strings.""" + if isinstance(value, float): + return f"{value:.2f}" + return str(value) + + +def build_stats_table(metrics: Mapping[str, object]) -> str: + """ + Build a simple ASCII table from metric-value key-value pairs. + + :param metrics: Ordered metric mapping of label -> value. + :return: Multi-line ASCII table. + """ + metric_header = "Metric" + value_header = "Value" + + rows = [(metric, _stringify_metric_value(value)) for metric, value in metrics.items()] + if not rows: + rows = [("No metrics", "n/a")] + + metric_width = max(len(metric_header), *(len(metric) for metric, _ in rows)) + value_width = max(len(value_header), *(len(value) for _, value in rows)) + + separator = f"+-{'-' * metric_width}-+-{'-' * value_width}-+" + header = f"| {metric_header:<{metric_width}} | {value_header:<{value_width}} |" + body = [ + f"| {metric:<{metric_width}} | {value:>{value_width}} |" + for metric, value in rows + ] + + return "\n".join([separator, header, separator, *body, separator]) + + +def format_crop_summary( + title: str, + metrics: Mapping[str, object], + detail_line: str | None = None, +) -> str: + """ + Format a crop summary title, optional detail line, and metric table. + + :param title: Summary title line. + :param metrics: Ordered metric mapping of label -> value. + :param detail_line: Optional descriptive one-liner for context. + :return: Formatted summary text. + """ + stats_table = build_stats_table(metrics) + if detail_line: + return f"{title}\n{detail_line}\n{stats_table}" + return f"{title}\n{stats_table}" + + +def warn_formatted_crop_summary( + title: str, + metrics: Mapping[str, object], + detail_line: str | None = None, +) -> None: + """ + Emit a module-scoped crop summary warning without file/line prefixes. + + :param title: Summary title line. + :param metrics: Ordered metric mapping of label -> value. + :param detail_line: Optional descriptive one-liner for context. + """ + message = format_crop_summary(title=title, metrics=metrics, detail_line=detail_line) + + original_formatwarning = warnings.formatwarning + try: + warnings.formatwarning = ( + lambda msg, category, filename, lineno, line=None: + f"{category.__name__}: {msg}\n" + ) + warnings.warn(message, category=CropSummaryWarning, stacklevel=2) + finally: + warnings.formatwarning = original_formatwarning diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/point_centered.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/point_centered.py new file mode 100644 index 0000000..a7769ab --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/point_centered.py @@ -0,0 +1,183 @@ +""" +point_centered.py + +Helper module for generating crops centered around specified points within an image. +""" + +import numpy as np + +from .crop_summary import warn_formatted_crop_summary +from .protocol import CropMap +from ...base_dataset import BaseImageDataset +from ..ds_utils import ( + _get_active_channels, + _validate_same_dimensions_across_channels +) + + +def _compute_point_centered_crops( + image_width: int, + image_height: int, + crop_size: int, + centers: dict[str, np.ndarray], + track_skipped: bool = False, +) -> tuple[list[tuple[int, int]], int, int]: + """ + Compute the top-left coordinates of crops centered around specified points, + ensuring that the crops fit fully within the image boundaries. + + :param image_width: Width of the image. + :param image_height: Height of the image. + :param crop_size: Size of the crops to generate. + :param centers: Dictionary containing 'X' and 'Y' coordinates of the centers. + :param track_skipped: Whether to count points skipped due to boundary checks. + :return: Tuple of: + - List of (x, y) tuples representing the top-left coordinates of the crops. + - Number of skipped points. + - Number of accepted points used for crops. + """ + + crops = [] + skipped_count = 0 + accepted_count = 0 + + for x, y in zip(centers['X'], centers['Y']): + + if x <= crop_size // 2 or x >= image_width - crop_size // 2 or \ + y <= crop_size // 2 or y >= image_height - crop_size // 2: + if track_skipped: + skipped_count += 1 + continue # Skip points too close to the edge for a full crop + + crop_x = int(x - crop_size // 2) + crop_y = int(y - crop_size // 2) + crops.append((crop_x, crop_y)) + if track_skipped: + accepted_count += 1 + + return crops, skipped_count, accepted_count + + +def generate_point_centered_crops( + dataset: BaseImageDataset, + crop_size: int | None = None, + mapping: list[dict[str, np.ndarray]] | None = None, + verbose: bool = False, +) -> CropMap: + """ + Generate crop specifications centered around specified points + for each image in the dataset. + + :param dataset: BaseImageDataset containing the images and metadata. + :param crop_size: Size of the crops to generate. Made optional + for better error handling when called from the CropImageDataset + .from_base_dataset class method. + :param mapping: List of dictionaries containing the centers for each image. + Each dictionary should have keys 'X' and 'Y' with corresponding numpy + arrays of coordinates. + Made optional for the same reason as crop_size. + :param verbose: If True, compute and report skipped-point statistics. + :return: CropMap containing the generated crop specifications. + """ + + if crop_size is None: + raise ValueError( + "crop_size must be provided for point-centered crop generation." + ) + if mapping is None: + raise ValueError( + "mapping must be provided for point-centered crop generation." + ) + + if not isinstance(crop_size, int): + raise ValueError(f"crop_size must be an integer, got {type(crop_size).__name__}.") + if crop_size <= 0: + raise ValueError(f"crop_size must be positive, got {crop_size}.") + + active_channels = _get_active_channels(dataset) + if not active_channels: + raise ValueError( + "No active channels configured. Set input_channel_keys and/or " + "target_channel_keys on the dataset before generating crops." + ) + + manifest = dataset.file_state.manifest + crop_specs: dict[int, list[tuple[tuple[int, int], int, int]]] = {} + skipped_points_per_image: list[int] = [] + accepted_points_per_image: list[int] = [] + + for idx in range(len(dataset)): + + dims = manifest.get_image_dimensions(idx, channels=active_channels) + + width, height = _validate_same_dimensions_across_channels( + dims, active_channels, idx + ) + + centers = mapping[idx] + + crop_lists, skipped_count, accepted_count = _compute_point_centered_crops( + width, + height, + crop_size, + centers, + track_skipped=verbose, + ) + crop_specs[idx] = [ + (crop_coords, crop_size, crop_size) for crop_coords in crop_lists + ] + + if verbose: + skipped_points_per_image.append(skipped_count) + accepted_points_per_image.append(accepted_count) + + if verbose: + total_points_skipped = sum(skipped_points_per_image) + total_points_accepted = sum(accepted_points_per_image) + num_images = len(skipped_points_per_image) + + avg_accepted_per_image = ( + total_points_accepted / num_images if num_images > 0 else 0.0 + ) + images_with_accepted = sum(1 for count in accepted_points_per_image if count > 0) + avg_accepted_per_affected_image = ( + total_points_accepted / images_with_accepted + if images_with_accepted > 0 + else 0.0 + ) + max_accepted_single_image = max(accepted_points_per_image, default=0) + + if total_points_skipped > 0: + avg_skipped_per_image = ( + total_points_skipped / num_images if num_images > 0 else 0.0 + ) + images_with_skips = sum(1 for count in skipped_points_per_image if count > 0) + avg_skipped_per_affected_image = ( + total_points_skipped / images_with_skips + if images_with_skips > 0 + else 0.0 + ) + max_skipped_single_image = max(skipped_points_per_image, default=0) + + warn_formatted_crop_summary( + title="Point-centered crop generation statistics:", + detail_line=( + "Exclusion criterion: a point is rejected when its center " + "is within or on half a crop-size from any image border, " + "which prevents full in-bounds crop extraction." + ), + metrics={ + "Total accepted": total_points_accepted, + "Total rejected": total_points_skipped, + "Mean accepted / image": avg_accepted_per_image, + "Mean rejected / image": avg_skipped_per_image, + "Mean accepted / affected image": avg_accepted_per_affected_image, + "Mean rejected / affected image": avg_skipped_per_affected_image, + "Images with accepted": f"{images_with_accepted}/{num_images}", + "Images with rejected": f"{images_with_skips}/{num_images}", + "Max accepted in one image": max_accepted_single_image, + "Max rejected in one image": max_skipped_single_image, + }, + ) + + return crop_specs diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/protocol.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/protocol.py new file mode 100644 index 0000000..374dc91 --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/protocol.py @@ -0,0 +1,25 @@ +""" +protocol.py + +Defines the CropGenerator protocol for crop generator functions. +""" + +from typing import Dict, List, Tuple, Any, Protocol + +from ...base_dataset import BaseImageDataset + + +CropSpec = Tuple[Tuple[int, int], int, int] +CropMap = Dict[int, List[CropSpec]] + + +class CropGenerator(Protocol): + """ + Protocol for crop generator functions. + """ + def __call__( + self, + dataset: BaseImageDataset, + **kwargs: Any + ) -> CropMap: + pass diff --git a/src/virtual_stain_flow/datasets/ds_engine/crop_generators/tile.py b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/tile.py new file mode 100644 index 0000000..1e20703 --- /dev/null +++ b/src/virtual_stain_flow/datasets/ds_engine/crop_generators/tile.py @@ -0,0 +1,153 @@ +""" +tile.py + +Crop generator module for creating non-overlapping tile crops centered within +images in a BaseImageDataset. Provides a best-effort tiling approach that +maximizes the number of full tiles while centering the grid within the image +boundaries. + +The non-overlapping decision is motivated by the property of convolutional +neural networks being translation-invariant, that it learns the same +information from a blob of pixels whether it is located at the center, the +left/right/top/down edge/corner. Having non-overlapping tiles ensures +minimal redundancy in the information captured across tiles. +""" + +from typing import List + +from .crop_summary import warn_formatted_crop_summary +from .protocol import CropMap, CropSpec +from ...base_dataset import BaseImageDataset +from ..ds_utils import ( + _get_active_channels, + _validate_same_dimensions_across_channels +) + + +def _compute_centered_tile_positions( + image_extent: int, + crop_size: int +) -> List[int]: + """ + Compute 1D tile start positions for non-overlapping tiles centered + within an image extent. + + :param image_extent: Size of image extent along one axis (width or height). + :param crop_size: Size of one tile along the same axis. + :return: List of start positions for each tile. + :raises ValueError: If crop_size exceeds image extent. + """ + n_tiles = image_extent // crop_size + if n_tiles == 0: + raise ValueError( + f"crop_size ({crop_size}) exceeds image dimensions " + f"along one axis ({image_extent})." + ) + + covered_extent = n_tiles * crop_size + margin = image_extent - covered_extent + start = margin // 2 + + return [start + (tile_idx * crop_size) for tile_idx in range(n_tiles)] + + +def _compute_centered_tile_crops( + image_width: int, + image_height: int, + crop_size: int +) -> List[CropSpec]: + """ + Compute non-overlapping square tiles arranged on a centered grid. + + :param image_width: Width of the source image. + :param image_height: Height of the source image. + :param crop_size: Size of square tile crop. + :return: List of crop specs in format ((x, y), width, height). + :raises ValueError: If crop_size exceeds image width or height. + """ + if crop_size > image_width or crop_size > image_height: + raise ValueError( + f"crop_size ({crop_size}) exceeds image dimensions " + f"({image_width}x{image_height})." + ) + + x_positions = _compute_centered_tile_positions(image_width, crop_size) + y_positions = _compute_centered_tile_positions(image_height, crop_size) + + return [ + ((x, y), crop_size, crop_size) + for y in y_positions + for x in x_positions + ] + + +def generate_tile_crops( + dataset: BaseImageDataset, + crop_size: int, + verbose: bool = False, +) -> CropMap: + """ + Generate best-effort centered, non-overlapping tiling crops + for each sample in a BaseImageDataset. + + Tiling is "best-effort" in that it uses as many full, non-overlapping + tiles of size `crop_size` as possible along each axis, then centers + the tile grid within the remaining field of view. + + :param dataset: A BaseImageDataset instance (or compatible object with + `file_state.manifest` attribute supporting `get_image_dimensions()`). + :param crop_size: Size of the square crop tile (same width and height). + :param verbose: If True, emit summary statistics for crop generation. + :return: Dictionary mapping manifest indices to lists of crop specs. + Format: {manifest_idx: [((x, y), width, height), ...]} + :raises ValueError: If crop_size is non-positive, if no active channels + are configured, or if channel dimensions don't match for any sample. + """ + if crop_size <= 0: + raise ValueError(f"crop_size must be positive, got {crop_size}.") + + active_channels = _get_active_channels(dataset) + if not active_channels: + raise ValueError( + "No active channels configured. Set input_channel_keys and/or " + "target_channel_keys on the dataset before generating crops." + ) + + manifest = dataset.file_state.manifest + crop_specs: CropMap = {} + tile_count_distribution: dict[int, int] = {} + total_samples = len(dataset) + + for idx in range(total_samples): + dims = manifest.get_image_dimensions(idx, channels=active_channels) + + width, height = _validate_same_dimensions_across_channels( + dims, active_channels, idx + ) + + crop_specs[idx] = _compute_centered_tile_crops( + width, height, crop_size + ) + tile_count = len(crop_specs[idx]) + tile_count_distribution[tile_count] = ( + tile_count_distribution.get(tile_count, 0) + 1 + ) + + if verbose: + distribution_metrics = { + f"Tiles per FOV = {tile_count}": full_fov_count + for tile_count, full_fov_count in sorted(tile_count_distribution.items()) + } + warn_formatted_crop_summary( + title="Tile crop generation statistics:", + detail_line=( + "Tiling criterion: each FOV receives as many full, non-overlapping " + "tiles as fit along width and height, centered within remaining margins." + ), + metrics={ + "Total dataset count": total_samples, + **distribution_metrics, + }, + ) + + return crop_specs diff --git a/src/virtual_stain_flow/transforms/README.md b/src/virtual_stain_flow/transforms/README.md index df63f20..441cbce 100644 --- a/src/virtual_stain_flow/transforms/README.md +++ b/src/virtual_stain_flow/transforms/README.md @@ -7,23 +7,33 @@ To facilitate reproducible experiments, the transformations are also made serial ## Overview -This subpackage consists of three modules: +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. 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. 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. +4. **`channelwise.py`**: Provides `ChannelwiseTransform`, a wrapper that applies a distinct transform to each channel of a channel-first image. + --- ## Usage: See examples for in context use with datasets ```python from normalizations import MaxScaleNormalize +from channelwise import ChannelwiseTransform scale_transform = MaxScaleNormalize( normalization_factor='16bit', name="Scale16BitImages", p=1.0 ) + +channelwise_transform = ChannelwiseTransform( + transforms=[ + MaxScaleNormalize(normalization_factor='16bit'), + MaxScaleNormalize(normalization_factor='8bit') + ] +) ``` diff --git a/src/virtual_stain_flow/transforms/__init__.py b/src/virtual_stain_flow/transforms/__init__.py index 4c6b16b..b5bfa14 100644 --- a/src/virtual_stain_flow/transforms/__init__.py +++ b/src/virtual_stain_flow/transforms/__init__.py @@ -1,3 +1,15 @@ """ /transforms/__init__.py """ + +from .channelwise import ChannelwiseTransform +from .normalizations import ( + MaxScaleNormalize, + ZScoreNormalize, +) + +__all__ = [ + "ChannelwiseTransform", + "MaxScaleNormalize", + "ZScoreNormalize", +] diff --git a/src/virtual_stain_flow/transforms/channelwise.py b/src/virtual_stain_flow/transforms/channelwise.py new file mode 100644 index 0000000..5114b16 --- /dev/null +++ b/src/virtual_stain_flow/transforms/channelwise.py @@ -0,0 +1,108 @@ +""" +channelwise.py + +Defines channel-specific transform wrappers. +Useful if the values of different channels take different ranges or distributions +and needs separate transformations/normalizations for each channel. +""" + +from typing import List, Optional, Sequence + +import numpy as np + +from .base_transform import LoggableTransform + + +class ChannelwiseTransform(LoggableTransform): + """ + Apply a list of transforms to a channel-first image, one transform per channel. + + A potential use case of this transform would be two modal input combining + two channels of different images (image augmentations). + An example would be brightfield image + phase retrieved image from brightfield z-stacks. + The former would take the range of a grayscale image, while the latter + would be in unit of radians for phase delay, centered around zero. + A reasonable channelwise transform would be a bitdepth max normalization + applied to the brightfield image channel and a z-score/radian normalization + applied to the phase retrieved image channel. + """ + + def __init__( + self, + transforms: Sequence[Optional[LoggableTransform]], + name: str = "ChannelwiseTransform", + p: float = 1.0, + channel_axis: int = 0, + ): + super().__init__(name=name, p=p) + + if channel_axis != 0: + raise ValueError("Only channel-first images with channel_axis=0 are supported.") + + if not isinstance(transforms, Sequence) or len(transforms) == 0: + raise ValueError("Expected a non-empty sequence of LoggableTransform or None.") + if not all((t is None) or isinstance(t, LoggableTransform) for t in transforms): + raise ValueError("All transforms must be instances of LoggableTransform or None.") + + self._transforms: List[Optional[LoggableTransform]] = list(transforms) + self._channel_axis = channel_axis + + @property + def transforms(self) -> List[Optional[LoggableTransform]]: + return self._transforms + + @property + def channel_axis(self) -> int: + return self._channel_axis + + def apply(self, img: np.ndarray, **params) -> np.ndarray: + if not isinstance(img, np.ndarray): + raise TypeError( + "Expected input image to be a NumPy array, " + f"got {type(img).__name__} instead." + ) + + if img.ndim != 3: + raise ValueError( + "Expected a channel-first image with shape (C, H, W)." + ) + + channel_count = img.shape[self._channel_axis] + if channel_count != len(self._transforms): + raise ValueError( + "Number of transforms must match number of channels. " + f"Got {len(self._transforms)} transforms for {channel_count} channels." + ) + + transformed_channels = [] + for channel_idx, transform in enumerate(self._transforms): + channel = img[channel_idx:channel_idx + 1, ...] + if transform is None: + transformed_channels.append(channel) + else: + transformed_channels.append(transform.apply(img=channel)) + + return np.concatenate(transformed_channels, axis=0) + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(name={self._name}, " + f"channels={len(self._transforms)}, p={self.p})" + ) + + def to_config(self) -> dict: + """ + Returns a dictionary containing the configuration of the transform. + + Helps with training reproducibility by allowing the logger to export + the configuration of the transform for later use. + """ + return { + "class": self.__class__.__name__, + "name": self._name, + "params": { + "channel_axis": self._channel_axis, + "p": self.p, + "transforms": [t.to_config() if t is not None else None for t in self._transforms], + }, + } diff --git a/tests/datasets/ds_engine/test_crop_generator.py b/tests/datasets/ds_engine/test_crop_generator.py index 3f23327..e9e87f0 100644 --- a/tests/datasets/ds_engine/test_crop_generator.py +++ b/tests/datasets/ds_engine/test_crop_generator.py @@ -9,6 +9,9 @@ _compute_center_crop, generate_center_crops, ) +from virtual_stain_flow.datasets.ds_engine.crop_generators.crop_summary import ( + CropSummaryWarning, +) class TestComputeCenterCrop: @@ -87,3 +90,11 @@ def test_raises_error_when_crop_too_large(self, basic_dataset): # Images are 10x10, so crop_size=20 should fail with pytest.raises(ValueError, match="exceeds image dimensions"): generate_center_crops(basic_dataset, crop_size=20) + + def test_verbose_emits_summary_warning(self, basic_dataset): + """Should emit one summary warning in verbose mode.""" + with pytest.warns(CropSummaryWarning, match="Center crop generation statistics") as record: + generate_center_crops(basic_dataset, crop_size=4, verbose=True) + + message = str(record[0].message) + assert "Success count / total dataset count" in message diff --git a/tests/datasets/ds_engine/test_crop_summary.py b/tests/datasets/ds_engine/test_crop_summary.py new file mode 100644 index 0000000..244f195 --- /dev/null +++ b/tests/datasets/ds_engine/test_crop_summary.py @@ -0,0 +1,48 @@ +""" +Minimal tests for crop summary formatting helpers. +""" + +import pytest + +from virtual_stain_flow.datasets.ds_engine.crop_generators.crop_summary import ( + CropSummaryWarning, + build_stats_table, + format_crop_summary, + warn_formatted_crop_summary, +) + + +def test_build_stats_table_renders_headers_and_rows(): + """Should render a compact ASCII table from metric key-value pairs.""" + table = build_stats_table({"Total accepted": 12, "Total rejected": 3}) + + assert "Metric" in table + assert "Value" in table + assert "Total accepted" in table + assert "Total rejected" in table + + +def test_warn_formatted_crop_summary_emits_warning_with_formatted_body(): + """Should emit CropSummaryWarning containing title, detail, and table rows.""" + with pytest.warns(CropSummaryWarning, match="Example summary") as record: + warn_formatted_crop_summary( + title="Example summary", + detail_line="Example criterion line.", + metrics={"Total": 2, "Mean": 1.0}, + ) + + message = str(record[0].message) + assert "Example criterion line." in message + assert "Total" in message + assert "Mean" in message + + +def test_format_crop_summary_without_detail_still_includes_table(): + """Should format title and table when detail line is omitted.""" + summary = format_crop_summary( + title="No detail summary", + metrics={"Success": "3/3"}, + ) + + assert "No detail summary" in summary + assert "Success" in summary diff --git a/tests/datasets/ds_engine/test_point_centered_crop_generator.py b/tests/datasets/ds_engine/test_point_centered_crop_generator.py new file mode 100644 index 0000000..95eee01 --- /dev/null +++ b/tests/datasets/ds_engine/test_point_centered_crop_generator.py @@ -0,0 +1,146 @@ +""" +Tests for point-centered crop generator core functionality. +""" + +import warnings + +import numpy as np +import pytest + +from virtual_stain_flow.datasets.base_dataset import BaseImageDataset +from virtual_stain_flow.datasets.ds_engine.crop_generator import ( + generate_point_centered_crops, +) +from virtual_stain_flow.datasets.ds_engine.crop_generators.crop_summary import ( + CropSummaryWarning, +) +from virtual_stain_flow.datasets.ds_engine.crop_generators.point_centered import ( + _compute_point_centered_crops, +) + + +class TestComputePointCenteredCrops: + """Core tests for point-centered crop coordinate computation.""" + + def test_computes_crops_and_tracks_counts(self): + """Should return valid crop coordinates with accepted/rejected counts.""" + centers = { + "X": np.array([5, 2, 8, 7]), + "Y": np.array([5, 5, 8, 7]), + } + + crops, skipped, accepted = _compute_point_centered_crops( + image_width=10, + image_height=10, + crop_size=4, + centers=centers, + track_skipped=True, + ) + + assert crops == [(3, 3), (5, 5)] + assert skipped == 2 + assert accepted == 2 + + +class TestGeneratePointCenteredCrops: + """Core tests for generate_point_centered_crops.""" + + def test_generates_expected_crop_specs(self, basic_dataset): + """Should generate crop specs in expected format for each sample.""" + mapping = [ + {"X": np.array([5, 7]), "Y": np.array([5, 7])}, + {"X": np.array([6]), "Y": np.array([6])}, + {"X": np.array([5, 2]), "Y": np.array([5, 5])}, + ] + + crop_specs = generate_point_centered_crops( + dataset=basic_dataset, + crop_size=4, + mapping=mapping, + ) + + assert set(crop_specs.keys()) == {0, 1, 2} + assert crop_specs[0] == [((3, 3), 4, 4), ((5, 5), 4, 4)] + assert crop_specs[1] == [((4, 4), 4, 4)] + assert crop_specs[2] == [((3, 3), 4, 4)] + + def test_verbose_emits_summary_warning_when_points_rejected(self, basic_dataset): + """Should emit one summary warning when any points are rejected.""" + mapping = [ + {"X": np.array([5, 2]), "Y": np.array([5, 5])}, + {"X": np.array([6]), "Y": np.array([6])}, + {"X": np.array([7]), "Y": np.array([7])}, + ] + + with pytest.warns( + CropSummaryWarning, + match="Point-centered crop generation statistics", + ) as record: + generate_point_centered_crops( + dataset=basic_dataset, + crop_size=4, + mapping=mapping, + verbose=True, + ) + + message = str(record[0].message) + assert "Total accepted" in message + assert "Total rejected" in message + + def test_verbose_no_warning_when_no_points_rejected(self, basic_dataset): + """Should not emit summary warning when no points are rejected.""" + mapping = [ + {"X": np.array([5, 7]), "Y": np.array([5, 7])}, + {"X": np.array([6]), "Y": np.array([6])}, + {"X": np.array([7]), "Y": np.array([7])}, + ] + + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + generate_point_centered_crops( + dataset=basic_dataset, + crop_size=4, + mapping=mapping, + verbose=True, + ) + + assert len(record) == 0 + + @pytest.mark.parametrize( + "crop_size,mapping,match", + [ + (None, [{"X": np.array([5]), "Y": np.array([5])}], "crop_size must be provided"), + (4, None, "mapping must be provided"), + (1.5, [{"X": np.array([5]), "Y": np.array([5])}], "crop_size must be an integer"), + (0, [{"X": np.array([5]), "Y": np.array([5])}], "crop_size must be positive"), + ], + ) + def test_validates_required_arguments(self, basic_dataset, crop_size, mapping, match): + """Should validate required inputs and crop_size constraints.""" + with pytest.raises(ValueError, match=match): + generate_point_centered_crops( + dataset=basic_dataset, + crop_size=crop_size, + mapping=mapping, + ) + + def test_raises_error_when_no_active_channels(self, file_index): + """Should raise ValueError when no channels are configured.""" + dataset = BaseImageDataset( + file_index=file_index, + pil_image_mode="I;16", + input_channel_keys=None, + target_channel_keys=None, + ) + mapping = [ + {"X": np.array([5]), "Y": np.array([5])}, + {"X": np.array([5]), "Y": np.array([5])}, + {"X": np.array([5]), "Y": np.array([5])}, + ] + + with pytest.raises(ValueError, match="No active channels"): + generate_point_centered_crops( + dataset=dataset, + crop_size=4, + mapping=mapping, + ) diff --git a/tests/datasets/ds_engine/test_tile_crop_generator.py b/tests/datasets/ds_engine/test_tile_crop_generator.py new file mode 100644 index 0000000..2785534 --- /dev/null +++ b/tests/datasets/ds_engine/test_tile_crop_generator.py @@ -0,0 +1,116 @@ +""" +Tests for tile crop generator functions. +""" + +import pytest + +from virtual_stain_flow.datasets.base_dataset import BaseImageDataset +from virtual_stain_flow.datasets.ds_engine.crop_generator import generate_tile_crops +from virtual_stain_flow.datasets.ds_engine.crop_generators.crop_summary import ( + CropSummaryWarning, +) +from virtual_stain_flow.datasets.ds_engine.crop_generators.tile import ( + _compute_centered_tile_positions, + _compute_centered_tile_crops, +) + + +class TestComputeCenteredTilePositions: + """Tests for 1D centered tile position computation.""" + + def test_returns_single_centered_position_when_one_tile_fits(self): + """Should center one tile when only one fits in the extent.""" + positions = _compute_centered_tile_positions(10, 8) + assert positions == [1] + + def test_returns_evenly_spaced_non_overlapping_positions(self): + """Should produce step=crop_size positions with centered margin.""" + positions = _compute_centered_tile_positions(10, 4) + assert positions == [1, 5] + + def test_returns_zero_margin_positions_when_exact_multiple(self): + """Should start at zero when extent is an exact multiple.""" + positions = _compute_centered_tile_positions(12, 4) + assert positions == [0, 4, 8] + + def test_raises_error_when_crop_larger_than_extent(self): + """Should raise ValueError when no tile can fit.""" + with pytest.raises(ValueError, match="exceeds image dimensions"): + _compute_centered_tile_positions(3, 4) + + +class TestComputeCenteredTileCrops: + """Tests for 2D centered tile crop computation.""" + + def test_computes_centered_non_overlapping_grid(self): + """Should produce Cartesian product of centered x/y positions.""" + crops = _compute_centered_tile_crops(10, 10, 4) + assert crops == [ + ((1, 1), 4, 4), + ((5, 1), 4, 4), + ((1, 5), 4, 4), + ((5, 5), 4, 4), + ] + + def test_raises_error_when_crop_larger_than_image(self): + """Should fail when crop_size exceeds width or height.""" + with pytest.raises(ValueError, match="exceeds image dimensions"): + _compute_centered_tile_crops(3, 10, 4) + + +class TestGenerateTileCrops: + """Tests for generate_tile_crops function.""" + + def test_generates_tiled_crops_for_all_samples(self, basic_dataset): + """Should generate tile crops for each sample index.""" + crop_specs = generate_tile_crops(basic_dataset, crop_size=4) + + assert len(crop_specs) == 3 + assert set(crop_specs.keys()) == {0, 1, 2} + + def test_crop_specs_have_centered_non_overlapping_layout(self, basic_dataset): + """Should return centered non-overlapping 4x4 tiles for 10x10 images.""" + crop_specs = generate_tile_crops(basic_dataset, crop_size=4) + + expected = [ + ((1, 1), 4, 4), + ((5, 1), 4, 4), + ((1, 5), 4, 4), + ((5, 5), 4, 4), + ] + + for _, crops in crop_specs.items(): + assert crops == expected + + def test_raises_error_for_non_positive_crop_size(self, basic_dataset): + """Should raise ValueError for crop_size <= 0.""" + with pytest.raises(ValueError, match="crop_size must be positive"): + generate_tile_crops(basic_dataset, crop_size=0) + + with pytest.raises(ValueError, match="crop_size must be positive"): + generate_tile_crops(basic_dataset, crop_size=-1) + + def test_raises_error_when_no_active_channels(self, file_index): + """Should raise ValueError when no channels are configured.""" + dataset = BaseImageDataset( + file_index=file_index, + pil_image_mode="I;16", + input_channel_keys=None, + target_channel_keys=None, + ) + + with pytest.raises(ValueError, match="No active channels"): + generate_tile_crops(dataset, crop_size=4) + + def test_raises_error_when_crop_too_large(self, basic_dataset): + """Should raise ValueError when crop_size exceeds image dimensions.""" + with pytest.raises(ValueError, match="exceeds image dimensions"): + generate_tile_crops(basic_dataset, crop_size=20) + + def test_verbose_emits_tile_distribution_summary(self, basic_dataset): + """Should emit one summary warning with tile-count distribution.""" + with pytest.warns(CropSummaryWarning, match="Tile crop generation statistics") as record: + generate_tile_crops(basic_dataset, crop_size=4, verbose=True) + + message = str(record[0].message) + assert "Tiles per FOV = 4" in message