torch-em 0.9.0__py3-none-any.whl

torch_em/__init__.py

@@ -0,0 +1,11 @@
+"""
+.. include:: ../doc/start_page.md
+.. include:: ../doc/datasets_and_dataloaders.md
+"""
+from .segmentation import (
+    default_segmentation_dataset,
+    default_segmentation_loader,
+    default_segmentation_trainer,
+    get_data_loader,
+)
+from .__version__ import __version__

torch_em/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.9.0"

torch_em/classification/__init__.py

@@ -0,0 +1,6 @@
+"""Dataset, DataLoader and Trainer implementations for image classification tasks.
+"""
+from .classification import default_classification_trainer, default_classification_loader
+
+from .classification_logger import ClassificationLogger
+from .classification_trainer import ClassificationTrainer

torch_em/classification/classification.py

@@ -0,0 +1,126 @@
+from functools import partial
+from typing import Optional, Sequence, Tuple, Union
+
+import numpy as np
+import sklearn.metrics as metrics
+import torch
+import torch_em
+from numpy.typing import ArrayLike
+
+from .classification_dataset import ClassificationDataset
+from .classification_logger import ClassificationLogger
+from .classification_trainer import ClassificationTrainer
+
+
+class ClassificationMetric:
+    """Metric for classification training.
+
+    Args:
+        metric_name: The name of the metrics. The name will be looked up in `sklearn.metrics`,
+            so it must be a valid identifier in that python package.
+        metric_kwargs: Keyword arguments for the metric.
+    """
+    def __init__(self, metric_name: str = "accuracy_score", **metric_kwargs):
+        if not hasattr(metrics, metric_name):
+            raise ValueError(f"Invalid metric_name {metric_name}.")
+        self.metric = getattr(metrics, metric_name)
+        self.metric_kwargs = metric_kwargs
+
+    def __call__(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
+        """Evaluate model prediction against classification labels.
+
+        Args:
+            y_true: The classification labels.
+            y_pred: The model predictions.
+
+        Returns:
+            The metric value.
+        """
+        metric_error = 1.0 - self.metric(y_true, y_pred, **self.metric_kwargs)
+        return metric_error
+
+
+def default_classification_loader(
+    data: Sequence[ArrayLike],
+    target: Sequence[ArrayLike],
+    batch_size: int,
+    normalization: Optional[callable] = None,
+    augmentation: Optional[callable] = None,
+    image_shape: Optional[Tuple[int, ...]] = None,
+    **loader_kwargs,
+) -> torch.utils.data.DataLoader:
+    """Get a data loader for classification training.
+
+    Args:
+        data: The input data for classification. Expects a sequence of array-like data.
+            The data can be two or three dimensional.
+        target: The target data for classification. Expects a sequence of the same length as `data`.
+            Each value in the sequence must be a scalar.
+        batch_size: The batch size for the data loader.
+        normalization: The normalization function. If None, data standardization will be used.
+        augmentation: The augmentation function. If None, the default augmentations will be used.
+        image_shape: The target shape of the data. If given, each sample will be resampled to this size.
+        loader_kwargs: Additional keyword arguments for `torch.utils.data.DataLoader`.
+
+    Returns:
+        The data loader.
+    """
+    ndim = data[0].ndim - 1
+    if ndim not in (2, 3):
+        raise ValueError(f"Expect input data of dimensionality 2 or 3, got {ndim}")
+
+    if normalization is None:
+        axis = (1, 2) if ndim == 2 else (1, 2, 3)
+        normalization = partial(torch_em.transform.raw.standardize, axis=axis)
+
+    if augmentation is None:
+        augmentation = torch_em.transform.get_augmentations(ndim=ndim)
+
+    dataset = ClassificationDataset(data, target, normalization, augmentation, image_shape)
+    loader = torch_em.segmentation.get_data_loader(dataset, batch_size, **loader_kwargs)
+    return loader
+
+
+def default_classification_trainer(
+    name: str,
+    model: torch.nn.Module,
+    train_loader: torch.utils.data.DataLoader,
+    val_loader: torch.utils.data.DataLoader,
+    loss: Optional[Union[torch.nn.Module, callable]] = None,
+    metric: Optional[Union[torch.nn.Module, callable]] = None,
+    logger=ClassificationLogger,
+    trainer_class=ClassificationTrainer,
+    **kwargs,
+):
+    """Get a trainer for a classification task.
+
+    This will create an instance of `torch_em.classification.ClassificationTrainer`.
+    Check out its documentation string for details on how to configure and use the trainer.
+
+    Args:
+        name: The name for the checkpoint created by the trainer.
+        model: The classification model to train.
+        train_loader: The data loader for training.
+        val_loader: The data loader for validation.
+        loss: The loss function. If None, will use cross entropy.
+        metric: The metric function. If None, will use the accuracy error.
+        logger: The logger for keeping track of the training progress.
+        trainer_class: The trainer class.
+        kwargs: Keyword arguments for the trainer class.
+
+    Returns:
+        The classification trainer.
+    """
+    # Set the default loss and metric (if no values where passed).
+    loss = torch.nn.CrossEntropyLoss() if loss is None else loss
+    metric = ClassificationMetric() if metric is None else metric
+
+    # Metric: Note that we use lower metric = better.
+    # So we record the accuracy error instead of the accuracy..
+    trainer = torch_em.default_segmentation_trainer(
+        name, model, train_loader, val_loader,
+        loss=loss, metric=metric,
+        logger=logger, trainer_class=trainer_class,
+        **kwargs,
+    )
+    return trainer

torch_em/classification/classification_dataset.py

@@ -0,0 +1,65 @@
+from typing import Sequence, Tuple
+
+import numpy as np
+import torch
+
+from numpy.typing import ArrayLike
+from skimage.transform import resize
+
+
+class ClassificationDataset(torch.utils.data.Dataset):
+    """Dataset for classification training.
+
+    Args:
+        data: The input data for classification. Expects a sequence of array-like data.
+            The data can be two or three dimensional.
+        target: The target data for classification. Expects a sequence of the same length as `data`.
+            Each value in the sequence must be a scalar.
+        normalization: The normalization function.
+        augmentation: The augmentation function.
+        image_shape: The target shape of the data. If given, each sample will be resampled to this size.
+    """
+    def __init__(
+        self,
+        data: Sequence[ArrayLike],
+        target: Sequence[ArrayLike],
+        normalization: callable,
+        augmentation: callable,
+        image_shape: Tuple[int, ...],
+    ):
+        if len(data) != len(target):
+            raise ValueError(f"Length of data and target don't agree: {len(data)} != {len(target)}")
+        self.data = data
+        self.target = target
+        self.normalization = normalization
+        self.augmentation = augmentation
+        self.image_shape = image_shape
+
+    def __len__(self):
+        return len(self.data)
+
+    def resize(self, x):
+        """@private
+        """
+        out = [resize(channel, self.image_shape, preserve_range=True)[None] for channel in x]
+        return np.concatenate(out, axis=0)
+
+    def __getitem__(self, index):
+        x, y = self.data[index], self.target[index]
+
+        # apply normalization
+        if self.normalization is not None:
+            x = self.normalization(x)
+
+        # resize to sample shape if it was given
+        if self.image_shape is not None:
+            x = self.resize(x)
+
+        # apply augmentations (if any)
+        if self.augmentation is not None:
+            _shape = x.shape
+            # adds unwanted batch axis
+            x = self.augmentation(x)[0][0]
+            assert x.shape == _shape
+
+        return x, y

torch_em/classification/classification_logger.py

@@ -0,0 +1,137 @@
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+from sklearn.metrics import ConfusionMatrixDisplay
+from torch.utils.tensorboard import SummaryWriter
+from torch_em.trainer.logger_base import TorchEmLogger
+from torch_em.transform.raw import normalize
+
+
+def confusion_matrix(y_true, y_pred, class_labels=None, title=None, save_path=None, **plot_kwargs):
+    """@private
+    """
+    fig, ax = plt.subplots(1)
+
+    if save_path is None:
+        canvas = FigureCanvasAgg(fig)
+
+    disp = ConfusionMatrixDisplay.from_predictions(
+        y_true, y_pred, normalize="true", display_labels=class_labels
+    )
+    disp.plot(ax=ax, **plot_kwargs)
+
+    if title is not None:
+        ax.set_title(title)
+    if save_path is not None:
+        plt.savefig(save_path)
+        return
+
+    canvas.draw()
+    image = np.asarray(canvas.buffer_rgba())[..., :3]
+    image = image.transpose((2, 0, 1))
+    plt.close()
+    return image
+
+
+def make_grid(images, target=None, prediction=None, images_per_row=8, **kwargs):
+    """@private
+    """
+    assert images.ndim in (4, 5)
+    assert images.shape[1] in (1, 3), f"{images.shape}"
+
+    if images.ndim == 5:
+        is_3d = True
+        z = images.shape[2] // 2
+    else:
+        is_3d = False
+
+    n_images = images.shape[0]
+    n_rows = n_images // images_per_row
+    if n_images % images_per_row != 0:
+        n_rows += 1
+
+    images = images.detach().cpu().numpy()
+    if target is not None:
+        target = target.detach().cpu().numpy()
+    if prediction is not None:
+        prediction = prediction.max(1)[1].detach().cpu().numpy()
+
+    fig, axes = plt.subplots(n_rows, images_per_row)
+    canvas = FigureCanvasAgg(fig)
+    for r in range(n_rows):
+        for c in range(images_per_row):
+            i = r * images_per_row + c
+            if i == len(images):
+                break
+            ax = axes[r, c] if n_rows > 1 else axes[r]
+            ax.set_axis_off()
+            im = images[i, :, z] if is_3d else images[i]
+            im = im.transpose((1, 2, 0))
+            im = normalize(im, axis=(0, 1))
+            if im.shape[-1] == 3:  # rgb
+                ax.imshow(im)
+            else:
+                ax.imshow(im[..., 0], cmap="gray")
+
+            if target is None and prediction is None:
+                continue
+
+            # TODO get the class name, and if we have both target
+            # and prediction check whether they agree or not and do stuff
+            title = ""
+            if target is not None:
+                title += f"t: {target[i]} "
+            if prediction is not None:
+                title += f"p: {prediction[i]}"
+            ax.set_title(title, fontsize=8)
+
+    canvas.draw()
+    image = np.asarray(canvas.buffer_rgba())[..., :3]
+    image = image.transpose((2, 0, 1))
+    plt.close()
+    return image
+
+
+class ClassificationLogger(TorchEmLogger):
+    """Logger for classification trainer.
+
+    Args:
+        trainer: The trainer instance.
+        save_root: Root folder for saving the checkpoints and logs.
+    """
+    def __init__(self, trainer, save_root: str, **unused_kwargs):
+        super().__init__(trainer, save_root)
+        self.log_dir = f"./logs/{trainer.name}" if save_root is None else\
+            os.path.join(save_root, "logs", trainer.name)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+        self.tb = SummaryWriter(self.log_dir)
+        self.log_image_interval = trainer.log_image_interval
+
+    def add_image(self, x, y, pred, name, step):
+        """@private
+        """
+        scale_each = False
+        grid = make_grid(x, y, pred, padding=4, normalize=True, scale_each=scale_each)
+        self.tb.add_image(tag=f"{name}/images_and_predictions", img_tensor=grid, global_step=step)
+
+    def log_train(self, step, loss, lr, x, y, prediction, log_gradients=False):
+        """@private
+        """
+        self.tb.add_scalar(tag="train/loss", scalar_value=loss, global_step=step)
+        self.tb.add_scalar(tag="train/learning_rate", scalar_value=lr, global_step=step)
+        if step % self.log_image_interval == 0:
+            self.add_image(x, y, prediction, "train", step)
+
+    def log_validation(self, step, metric, loss, x, y, prediction, y_true=None, y_pred=None):
+        """@private
+        """
+        self.tb.add_scalar(tag="validation/loss", scalar_value=loss, global_step=step)
+        self.tb.add_scalar(tag="validation/metric", scalar_value=metric, global_step=step)
+        self.add_image(x, y, prediction, "validation", step)
+        if y_true is not None and y_pred is not None:
+            cm = confusion_matrix(y_true, y_pred)
+            self.tb.add_image(tag="validation/confusion_matrix", img_tensor=cm, global_step=step)

torch_em/classification/classification_trainer.py

@@ -0,0 +1,42 @@
+import warnings
+
+import numpy as np
+import torch
+import torch_em
+
+
+class ClassificationTrainer(torch_em.trainer.DefaultTrainer):
+    """Trainer for classification tasks.
+
+    This class inherits from `torch_em.trainer.DefaultTrainer` with minor changes for
+    classification instead of segmentation training. Check out the documentation of
+    the default trainer class for details on how to configure and use the trainer
+    """
+    def _validate_impl(self, forward_context):
+        self.model.eval()
+
+        loss_val = 0.0
+
+        # we use the syntax from sklearn.metrics to compute metrics
+        # over all the preditions
+        y_true, y_pred = [], []
+
+        with torch.no_grad():
+            for x, y in self.val_loader:
+                x, y = x.to(self.device), y.to(self.device)
+                with forward_context():
+                    pred, loss = self._forward_and_loss(x, y)
+                loss_val += loss.item()
+                y_true.append(y.detach().cpu().numpy())
+                y_pred.append(pred.max(1)[1].detach().cpu().numpy())
+
+        if torch.isnan(pred).any():
+            warnings.warn("Predictions are NaN")
+        loss_val /= len(self.val_loader)
+
+        y_true, y_pred = np.concatenate(y_true), np.concatenate(y_pred)
+        metric_val = self.metric(y_true, y_pred)
+
+        if self.logger is not None:
+            self.logger.log_validation(self._iteration, metric_val, loss_val, x, y, pred, y_true, y_pred)
+        return metric_val