SAMHI/resources/CellViT/cell_segmentation/trainer/trainer_cpp_net.py

# -*- coding: utf-8 -*-
# CPP Trainer Class
#
# @ Fabian Hörst, fabian.hoerst@uk-essen.de
# Institute for Artifical Intelligence in Medicine,
# University Medicine Essen

import logging
from pathlib import Path
from typing import Union

import numpy as np
import torch
import torch.nn.functional as F

# import wandb
from matplotlib import pyplot as plt
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler
from base_ml.base_early_stopping import EarlyStopping
from models.segmentation.cell_segmentation.cellvit_cpp_net import (
    DataclassCPPStorage,
)

from cell_segmentation.trainer.trainer_stardist import CellViTStarDistTrainer
from models.segmentation.cell_segmentation.cellvit import CellViT

# import warnings
# warnings.filterwarnings("ignore", category=DeprecationWarning)


class CellViTCPPTrainer(CellViTStarDistTrainer):
    """CellViTCPP trainer class

    Args:
        model (CellViTCPP): CellViTCPP model that should be trained
        loss_fn_dict (dict): Dictionary with loss functions for each branch with a dictionary of loss functions.
            Name of branch as top-level key, followed by a dictionary with loss name, loss fn and weighting factor
            Example:
            {
                "dist_map": {"bce": {loss_fn(Callable), weight_factor(float)}, "dice": {loss_fn(Callable), weight_factor(float)}},
                "stardist_map": {"bce": {loss_fn(Callable), weight_factor(float)}, "dice": {loss_fn(Callable), weight_factor(float)}},
                "stardist_map_refined": {"bce": {loss_fn(Callable), weight_factor(float)}, "dice": {loss_fn(Callable), weight_factor(float)}},
                "nuclei_type_map": {"bce": {loss_fn(Callable), weight_factor(float)}, "dice": {loss_fn(Callable), weight_factor(float)}}
                "tissue_types": {"ce": {loss_fn(Callable), weight_factor(float)}}
            }
            Required Keys are:
                * dist_map
                * stardist_map
                * stardist_map_refined
                * nuclei_type_map
                * tissue types
        optimizer (Optimizer): Optimizer
        scheduler (_LRScheduler): Learning rate scheduler
        device (str): Cuda device to use, e.g., cuda:0.
        logger (logging.Logger): Logger module
        logdir (Union[Path, str]): Logging directory
        num_classes (int): Number of nuclei classes
        dataset_config (dict): Dataset configuration. Required Keys are:
            * "tissue_types": describing the present tissue types with corresponding integer
            * "nuclei_types": describing the present nuclei types with corresponding integer
        experiment_config (dict): Configuration of this experiment
        early_stopping (EarlyStopping, optional):  Early Stopping Class. Defaults to None.
        log_images (bool, optional): If images should be logged to WandB. Defaults to False.
        magnification (int, optional): Image magnification. Please select either 40 or 20. Defaults to 40.
        mixed_precision (bool, optional): If mixed-precision should be used. Defaults to False.
    """

    def __init__(
        self,
        model: CellViT,
        loss_fn_dict: dict,
        optimizer: Optimizer,
        scheduler: _LRScheduler,
        device: str,
        logger: logging.Logger,
        logdir: Union[Path, str],
        num_classes: int,
        dataset_config: dict,
        experiment_config: dict,
        early_stopping: EarlyStopping = None,
        log_images: bool = False,
        magnification: int = 40,
        mixed_precision: bool = False,
    ):
        super().__init__(
            model=model,
            loss_fn_dict=loss_fn_dict,
            optimizer=optimizer,
            scheduler=scheduler,
            device=device,
            logger=logger,
            logdir=logdir,
            num_classes=num_classes,
            dataset_config=dataset_config,
            experiment_config=experiment_config,
            early_stopping=early_stopping,
            log_images=log_images,
            magnification=magnification,
            mixed_precision=mixed_precision,
        )

    def unpack_predictions(
        self, predictions: dict, skip_postprocessing: bool = False
    ) -> DataclassCPPStorage:
        """Unpack the given predictions. Main focus lays on reshaping and postprocessing predictions, e.g. separating instances

        Args:
            predictions (dict): Dictionary with the following keys:
                * tissue_types: Logit tissue prediction output. Shape: (batch_size, num_tissue_classes)
                * stardist_map: Stardist output for vector prediction. Shape: (batch_size, n_rays, H, W)
                * stardist_map_refined: Stardist output for vector prediction, but refined by CPP-Net. Shape: (batch_size, n_rays, H, W)
                * dist_map: Logit output for distance map. Shape: (batch_size, 1, H, W)
                * (Optional)
                * nuclei_type_map: Logit output for nuclei instance-prediction. Shape: (batch_size, num_nuclei_classes, H, W)
            skip_postprocessing (bool, optional): If true, postprocesssing for separating nuclei and creating maps is skipped.
                Helpfull for speeding up training. Defaults to False.
        Returns:
            DataclassCPPStorage: Processed network output
        """
        predictions["nuclei_type_map"] = F.softmax(
            predictions["nuclei_type_map"], dim=1
        )  # shape: (batch_size, num_nuclei_classes, H, W)
        predictions["dist_map_sigmoid"] = F.sigmoid(predictions["dist_map"])
        # postprocessing: apply NMS and StarDist postprocessing to generate binary and multiclass cell detections

        if not skip_postprocessing:
            (
                instance_map,
                predictions["instance_types"],
                instance_types_nuclei,
            ) = self.model.calculate_instance_map(
                predictions["dist_map_sigmoid"],
                predictions["stardist_map_refined"],
                predictions["nuclei_type_map"],
            )
            instance_map = instance_map.to(self.device)
            instance_types_nuclei = instance_types_nuclei.to(self.device)
            predictions["instance_map"] = instance_map
            predictions["instance_types_nuclei"] = instance_types_nuclei

        predictions = DataclassCPPStorage(
            **predictions,
            batch_size=predictions["nuclei_type_map"].shape[0],
        )

        return predictions

    def unpack_masks(self, masks: dict, tissue_types: list) -> DataclassCPPStorage:
        """Unpack the given masks. Main focus lays on reshaping and postprocessing masks to generate one dict

        Args:
            masks (dict): Required keys are:
                * instance_map: Pixel-wise nuclear instance segmentations. Shape: (batch_size, H, W)
                * nuclei_type_map: Nuclei instance-prediction and segmentation (not binary, each instance has own integer). Shape: (batch_size, H, W)
                * nuclei_binary_map: Binary nuclei segmentations. Shape: (batch_size, H, W)
                * dist_map: Probability distance map. Shape: (batch_size, H, W)
                * stardist_map: Stardist output. Shape: (batch_size, n_rays H, W)

            tissue_types (list): List of string names of ground-truth tissue types

        Returns:
            DataclassCPPStorage: Output ground truth values

        """
        nuclei_type_maps = torch.squeeze(masks["nuclei_type_map"]).type(torch.int64)
        gt_nuclei_type_maps_onehot = F.one_hot(
            nuclei_type_maps, num_classes=self.num_classes
        ).type(
            torch.float32
        )  # background + nuclei types

        # # assemble ground truth dictionary
        gt = {
            "nuclei_type_map": gt_nuclei_type_maps_onehot.permute(0, 3, 1, 2).to(
                self.device
            ),  # shape: (batch_size, num_nuclei_classes, H, W)
            "stardist_map": masks["stardist_map"].to(
                self.device
            ),  # shape: (batch_size, nrays, H, W)
            "stardist_map_refined": masks["stardist_map"].to(
                self.device
            ),  # shape: (batch_size, nrays, H, W)
            "dist_map": masks["dist_map"].to(self.device)[
                :, None, :, :
            ],  # shape: (batch_size, 1, H, W), TODO: check if None is necessary because of shape?
            "instance_map": masks["instance_map"].to(
                self.device
            ),  # shape: (batch_size, H, W) -> each instance has one integer
            "instance_types_nuclei": (
                gt_nuclei_type_maps_onehot * masks["instance_map"][..., None]
            )
            .permute(0, 3, 1, 2)
            .to(
                self.device
            ),  # shape: (batch_size, num_nuclei_classes, H, W) -> instance has one integer, for each nuclei class
            "tissue_types": torch.Tensor([self.tissue_types[t] for t in tissue_types])
            .type(torch.LongTensor)
            .to(self.device),  # shape: batch_size
        }
        gt = DataclassCPPStorage(**gt, batch_size=gt["tissue_types"].shape[0])
        return gt

    @staticmethod
    def generate_example_image(
        imgs: Union[torch.Tensor, np.ndarray],
        predictions: dict,
        ground_truth: dict,
        num_nuclei_classes: int,
        num_images: int = 2,
    ) -> plt.Figure:
        # TODO: implement
        return None