Datasets

mtlearn.datasets contains PyTorch helpers for paired image datasets and datasets that generate targets from input images with a callable.

Public PyTorch dataset helpers for mtlearn.

class mtlearn.datasets.GeneratedTargetImageDataset(root_dir, target_fn, *, num_rows=None, num_cols=None, grayscale=True, invert_in=False, extensions=('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'), dtype=torch.float32, scale_in=True, scale_out=True, suffix_in='')

Bases: Dataset

Read source images and generate each target with target_fn.

Samples are returned as (input_tensor, target_tensor, filename) with channel-first tensors. Resize is optional; when configured, the source image is resized before calling target_fn. The generated target must match the post-resize source spatial shape.

Parameters:

  • root_dir (str)

  • target_fn (Callable[[np.ndarray], np.ndarray])

  • num_rows (int | None)

  • num_cols (int | None)

  • grayscale (bool)

  • invert_in (bool)

  • extensions (tuple[str, ...])

  • dtype (torch.dtype)

  • scale_in (bool)

  • scale_out (bool)

  • suffix_in (str)

train_test_split(test_size=0.25, shuffle=True, random_state=42)[source]

Return (train_subset, test_subset) using stable dataset indices.

class mtlearn.datasets.PairedImageDataset(root_dir, num_rows=None, num_cols=None, *, grayscale_in=True, grayscale_target=True, invert_in=False, invert_target=False, extensions=('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'), dtype=torch.float32, scale_in=True, scale_out=True, prefix_in='', prefix_target='', suffix_in='_in', suffix_target='_target')

Bases: Dataset

Read matched input/target image pairs from one directory.

Samples are returned as (input_tensor, target_tensor, filename) with channel-first tensors. Resize is optional; when configured, both input and target images are resized to (num_rows, num_cols). Without resize, the input and target spatial shapes must already match.

Parameters:

  • root_dir (str)

  • num_rows (int | None)

  • num_cols (int | None)

  • grayscale_in (bool)

  • grayscale_target (bool)

  • invert_in (bool)

  • invert_target (bool)

  • extensions (tuple[str, ...])

  • dtype (torch.dtype)

  • scale_in (bool)

  • scale_out (bool)

  • prefix_in (str)

  • prefix_target (str)

  • suffix_in (str)

  • suffix_target (str)

train_test_split(test_size=0.25, shuffle=True, random_state=42)[source]

Return (train_subset, test_subset) using stable pair indices.