Upload 9 files

prostate/data.py

@@ -0,0 +1,229 @@
+# create dataloaders form csv file
+
+## ---------- imports ----------
+import os        
+import torch
+import shutil
+import numpy as np
+import pandas as pd
+from typing import Union
+from monai.utils import first
+from functools import partial
+from collections import namedtuple
+from monai.data import DataLoader as MonaiDataLoader
+      
+from . import transforms
+from .utils import num_workers
+
+
+def import_dataset(config: dict): 
+    if config.data.dataset_type == 'persistent':
+        from monai.data import PersistentDataset
+        if os.path.exists(config.data.cache_dir): 
+            shutil.rmtree(config.data.cache_dir) # rm previous cache DS
+        os.makedirs(config.data.cache_dir, exist_ok = True)
+        Dataset = partial(PersistentDataset, cache_dir = config.data.cache_dir)
+    elif config.data.dataset_type == 'cache':
+        from monai.data import CacheDataset
+        raise NotImplementedError('CacheDataset not yet implemented')
+    else:
+        from monai.data import Dataset
+    return Dataset
+
+
+class DataLoader(MonaiDataLoader): 
+    "overwrite monai DataLoader for enhanced viewing capabilities"
+    
+    def show_batch(self, 
+                   image_key: str='image', 
+                   label_key: str='label', 
+                   image_transform=lambda x: x.squeeze().transpose(0,2).flip(-2), 
+                   label_transform=lambda x: x.squeeze().transpose(0,2).flip(-2)): 
+        """Args:
+            image_key: dict key name for image to view
+            label_key: dict kex name for corresponding label. Can be a tensor or str
+            image_transform: transform input before it is passed to the viewer to ensure
+                ndim of the image is equal to 3 and image is oriented correctly
+            label_transform: transform labels before passed to the viewer, to ensure 
+                segmentations masks have same shape and orientations as images. Should be 
+                identity function of labels are str. 
+        """
+        from .viewer import ListViewer
+        
+        batch = first(self)
+        image = torch.unbind(batch[image_key], 0)
+        label = torch.unbind(batch[label_key], 0)
+        
+        ListViewer([image_transform(im) for im in image],
+                   [label_transform(im) for im in label]).show()
+
+# TODO
+## Work with 3 dataloaders
+        
+def segmentation_dataloaders(config: dict, 
+                             train: bool = None,
+                             valid: bool = None,
+                             test: bool = None,
+                            ):
+    """Create segmentation dataloaders
+    Args:
+        config: config file
+        train: whether to return a train DataLoader
+        valid: whether to return a valid DataLoader
+        test: whether to return a test DateLoader
+    Args from config: 
+        data_dir: base directory for the data
+        csv_name: path to csv file containing filenames and paths
+        image_cols: columns in csv containing path to images
+        label_cols: columns in csv containing path to label files
+        dataset_type: PersistentDataset, CacheDataset and Dataset are supported
+        cache_dir: cache directory to be used by PersistentDataset
+        batch_size: batch size for training. Valid and test are always 1
+        debug: run with reduced number of images
+    Returns:
+        list of:
+            train_loader: DataLoader (optional, if train==True)
+            valid_loader: DataLoader (optional, if valid==True)
+            test_loader: DataLoader (optional, if test==True)
+    """
+    
+    ## parse needed rguments from config
+    if train is None: train = config.data.train
+    if valid is None: valid = config.data.valid
+    if test is None: test = config.data.test
+    
+    data_dir = config.data.data_dir
+    train_csv = config.data.train_csv
+    valid_csv = config.data.valid_csv
+    test_csv = config.data.test_csv
+    image_cols = config.data.image_cols
+    label_cols = config.data.label_cols
+    dataset_type = config.data.dataset_type
+    cache_dir = config.data.cache_dir
+    batch_size = config.data.batch_size
+    debug = config.debug
+            
+    ## ---------- data dicts ----------
+
+    # first a global data dict, containing only the filepath from image_cols and label_cols is created. For this,
+    # the dataframe is reduced to only the relevant columns. Then the rows are iterated, converting each row into an
+    # individual dict, as expected by monai
+
+    if not isinstance(image_cols, (tuple, list)): image_cols = [image_cols]
+    if not isinstance(label_cols, (tuple, list)): label_cols = [label_cols]
+
+    train_df = pd.read_csv(train_csv)
+    valid_df = pd.read_csv(valid_csv)
+    test_df = pd.read_csv(test_csv)
+    if debug: 
+        train_df = train_df.sample(25)
+        valid_df = valid_df.sample(5)
+    
+    train_df['split']='train'
+    valid_df['split']='valid'
+    test_df['split']='test'
+    whole_df = []
+    if train: whole_df += [train_df]
+    if valid: whole_df += [valid_df]
+    if test: whole_df += [test_df]
+    df = pd.concat(whole_df)
+    cols = image_cols + label_cols
+    for col in cols:
+        # create absolute file name from relative fn in df and data_dir
+        df[col] = [os.path.join(data_dir, fn) for fn in df[col]]
+        if not os.path.exists(list(df[col])[0]):
+            raise FileNotFoundError(list(df[col])[0])
+
+
+    data_dict = [dict(row[1]) for row in df[cols].iterrows()]
+    # data_dict is not the correct name, list_of_data_dicts would be more accurate, but also longer.
+    # The data_dict looks like this:
+    # [
+    #  {'image_col_1': 'data_dir/path/to/image1',
+    #   'image_col_2': 'data_dir/path/to/image2'
+    #   'label_col_1': 'data_dir/path/to/label1},
+    #  {'image_col_1': 'data_dir/path/to/image1',
+    #   'image_col_2': 'data_dir/path/to/image2'
+    #   'label_col_1': 'data_dir/path/to/label1},
+    #    ...]
+    # Filename should now be absolute or relative to working directory
+
+    # now we create separate data dicts for train, valid and test data respectively
+    assert train or test or valid, 'No dataset type is specified (train/valid or test)'
+
+    if test:
+        test_files = list(map(data_dict.__getitem__, *np.where(df.split == 'test')))
+
+    if valid:
+        val_files = list(map(data_dict.__getitem__, *np.where(df.split == 'valid')))
+
+    if train:
+        train_files = list(map(data_dict.__getitem__, *np.where(df.split == 'train')))
+
+    # transforms are specified in transforms.py and are just loaded here
+    if train: train_transforms = transforms.get_train_transforms(config)
+    if valid: val_transforms = transforms.get_val_transforms(config)
+    if test: test_transforms = transforms.get_test_transforms(config)
+    
+    
+    ## ---------- construct dataloaders ----------
+    Dataset=import_dataset(config)
+    data_loaders = []
+    if train:
+        train_ds = Dataset(
+            data=train_files,
+            transform=train_transforms
+        )
+        train_loader = DataLoader(
+            train_ds,
+            batch_size=batch_size,
+            num_workers=num_workers(),
+            shuffle=True
+        )
+        data_loaders.append(train_loader)
+
+    if valid:
+        val_ds = Dataset(
+            data=val_files,
+            transform=val_transforms
+        )
+        val_loader = DataLoader(
+            val_ds,
+            batch_size=1,
+            num_workers=num_workers(),
+            shuffle=False
+        )
+        data_loaders.append(val_loader)
+
+    if test:
+        test_ds = Dataset(
+            data=test_files,
+            transform=test_transforms
+        )
+        test_loader = DataLoader(
+            test_ds,
+            batch_size=1,
+            num_workers=num_workers(),
+            shuffle=False
+        )
+        data_loaders.append(test_loader)
+
+    # if only one dataloader is constructed, return only this dataloader else return a named tuple with dataloaders,
+    # so it is clear which DataLoader is train/valid or test
+
+    if len(data_loaders) == 1:
+        return data_loaders[0]
+    else:
+        DataLoaders = namedtuple(
+            'DataLoaders',
+            # create str with specification of loader type if train and test are true but
+            # valid is false string will be 'train test'
+            ' '.join(
+                [
+                    'train' if train else '',
+                    'valid' if valid else '',
+                    'test' if test else ''
+                ]
+            ).strip()
+        )
+        return  DataLoaders(*data_loaders)

prostate/loss.py

@@ -0,0 +1,23 @@
+import monai
+from .utils import load_config
+
+
+def get_loss(config: dict): 
+    """Create a loss function of `type` with specific keyword arguments from config.
+    Example: 
+        
+        config.loss
+        >>> {'DiceCELoss': {'include_background': False, 'softmax': True, 'to_onehot_y': True}}
+
+        get_loss(config)
+        >>> DiceCELoss(
+        >>>   (dice): DiceLoss()
+        >>>   (cross_entropy): CrossEntropyLoss()
+        >>> )
+    
+    """
+    loss_type = list(config.loss.keys())[0]
+    loss_config = config.loss[loss_type]
+    loss_fun =  getattr(monai.losses, loss_type)
+    loss = loss_fun(**loss_config)
+    return loss

prostate/model.py

@@ -0,0 +1,16 @@
+# create a standard UNet
+
+from monai.networks.nets import UNet
+
+def get_model(config: dict):
+    return UNet(
+        spatial_dims=config.ndim,
+        in_channels=len(config.data.image_cols),
+        out_channels=config.model.out_channels,
+        channels=config.model.channels,
+        strides=config.model.strides,
+        num_res_units=config.model.num_res_units,
+        act=config.model.act,
+        norm=config.model.norm,
+        dropout=config.model.dropout,
+            )

prostate/optimizer.py

@@ -0,0 +1,37 @@
+import torch
+import monai
+from .utils import load_config
+
+
+def get_optimizer(model: torch.nn.Module, 
+                  config: dict): 
+    """Create an optimizer of `type` with specific keyword arguments from config.
+    Example: 
+        
+        config.optimizer
+        >>> {'Novograd': {'lr': 0.001, 'weight_decay': 0.01}}
+
+        get_optimizer(model, config)
+        >>> Novograd (
+        >>> Parameter Group 0
+        >>>     amsgrad: False
+        >>>     betas: (0.9, 0.999)
+        >>>     eps: 1e-08
+        >>>     grad_averaging: False
+        >>>     lr: 0.0001
+        >>>     weight_decay: 0.001
+        >>> )
+    
+    """
+    optimizer_type = list(config.optimizer.keys())[0]
+    opt_config = config.optimizer[optimizer_type]
+    if hasattr(torch.optim, optimizer_type): 
+        optimizer_fun = getattr(torch.optim, optimizer_type)
+    elif hasattr(monai.optimizers, optimizer_type): 
+        optimizer_fun = getattr(monai.optimizers, optimizer_type)
+    else: 
+        raise ValueError(f'Optimizer {optimizer_type} not found')
+    optimizer = optimizer_fun(model.parameters(), **opt_config)
+    return optimizer
+
+

prostate/report.py

@@ -0,0 +1,184 @@
+import os
+import io
+import cv2
+import tqdm
+import torch
+import imageio
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+class ReportGenerator():
+    "Generate markdown document, summarizing the training"
+    
+    def __init__(self, run_id, out_dir=None, log_dir=None):
+    
+        self.run_id, self.out_dir, self.log_dir = run_id, out_dir, log_dir
+
+        if log_dir: 
+            self.train_logs = pd.read_csv(os.path.join(log_dir, 'train_logs.csv'))
+            self.metric_logs = pd.read_csv(os.path.join(log_dir, 'metric_logs.csv'))
+        if out_dir: 
+            self.dice = pd.read_csv(os.path.join(out_dir, 'MeanDice_raw.csv'))
+            self.hausdorf = pd.read_csv(os.path.join(out_dir, 'HausdorffDistance_raw.csv'))
+            self.surface = pd.read_csv(os.path.join(out_dir, 'SurfaceDistance_raw.csv'))
+
+            self.mean_metrics = pd.DataFrame(
+                {"mean_dice" : [round(np.mean(self.dice[col]),3) for col in self.dice if col.startswith('class')],  
+                 "mean_hausdorf" : [round(np.mean(self.hausdorf[col]),3) for col in self.hausdorf if col.startswith('class')],   
+                 "mean_surface" : [round(np.mean(self.surface[col]),3) for col in self.surface if col.startswith('class')] 
+                }).transpose()
+    
+    def generate_report(self, loss_plot=True, metric_plot=True, boxplots=True, animation=True):
+        fn = os.path.join(self.run_id, 'report', 'SegmentationReport.md')
+        os.makedirs(os.path.join(self.run_id, 'report'), exist_ok=True)
+        with open(fn, 'w+') as f: 
+            f.write('# Segmentation Report\n\n')
+        
+        if loss_plot: 
+            fig = self.plot_loss(self.train_logs, self.metric_logs)
+            plt.savefig(os.path.join(self.run_id, 'report', 'loss_and_lr.png'), dpi = 150)
+
+            with open(fn, 'a') as f: 
+                f.write('## Loss, LR-Schedule and Key Metric\n')
+                f.write('![Loss, LR-Schedule and Key Metric](loss_and_lr.png)\n\n')
+
+        if metric_plot: 
+            fig = plt.figure("metrics", (18, 6))
+
+            ax = plt.subplot(1, 3, 1)
+            plt.ylim([0,1])
+            plt.title("Mean Dice")
+            plt.xlabel("epoch")
+            plt.plot(self.metric_logs.index, self.metric_logs.MeanDice)
+
+            ax = plt.subplot(1, 3, 2)
+            plt.title("Mean Hausdorff Distance")
+            plt.xlabel("epoch")
+            plt.plot(self.metric_logs.index, self.metric_logs.HausdorffDistance)
+
+            ax = plt.subplot(1, 3, 3)
+            plt.title("Mean Surface Distance")
+            plt.xlabel("epoch")
+            plt.plot(self.metric_logs.index, self.metric_logs.SurfaceDistance)
+
+            plt.savefig(os.path.join(self.run_id, 'report', 'metrics.png'), dpi = 150)
+            fig.clear()
+            plt.close()
+
+            with open(fn, 'a') as f: 
+                f.write('## Metrics\n')
+                f.write('![metrics](metrics.png)\n\n')
+
+        if boxplots: 
+            fig = plt.figure("boxplots", (18, 6))
+
+            ax = plt.subplot(1, 3, 1)
+            plt.title("Dice")
+            plt.xlabel("class")
+            plt.boxplot(self.dice[[col for col in self.dice if col.startswith('class')]])
+
+            ax = plt.subplot(1, 3, 2)
+            plt.title("Hausdorff Distance")
+            plt.xlabel("class")
+            plt.boxplot(self.hausdorf[[col for col in self.hausdorf if col.startswith('class')]])
+
+            ax = plt.subplot(1, 3, 3)
+            plt.title("Surface Distance")
+            plt.xlabel("class")
+            plt.boxplot(self.surface[[col for col in self.surface if col.startswith('class')]])
+
+            plt.savefig(os.path.join(self.run_id, 'report', 'boxplots.png'),dpi = 150)
+
+            fig.clear()
+            plt.close()
+
+            with open(fn, 'a') as f: 
+                f.write(f"## Individual metrics\n\n")
+                f.write(f"{self.mean_metrics.to_markdown()}\n\n")
+                f.write(f"![boxplot](boxplots.png)\n\n")
+        if animation: 
+            self.generate_gif()
+            with open(fn, 'a') as f: 
+                f.write('## Visualization of progress\n')
+                f.write('![progress](progress.gif)\n\n')
+
+    def plot_loss(self, train_logs, metric_logs): 
+        iteration = train_logs.iteration/sum(train_logs.epoch == 1)
+        fig = plt.figure("loss and lr", (12, 6))
+
+        y_max = max(metric_logs.eval_loss) + 0.5
+        if y_max > 3: y_max = 3
+
+        ax = plt.subplot(1, 2, 1)
+        plt.ylim([0,y_max])
+        plt.title("Epoch Average Loss")
+        plt.xlabel("epoch")
+        plt.plot(iteration, train_logs.loss)
+        plt.plot(metric_logs.index, metric_logs.eval_loss)
+
+        ax = plt.subplot(1, 2, 2)
+        ax.set_yscale('log')
+        plt.title("LR Schedule")
+        plt.xlabel("epoch")
+        plt.plot(iteration, train_logs.lr)
+        return fig
+
+    def get_arr_from_fig(self, fig, dpi=180):
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", dpi=dpi)
+        buf.seek(0)
+        img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
+        buf.close()
+        img = cv2.imdecode(img_arr, 1)
+        #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        return img
+    
+    def get_slices(self, im, slices): 
+        ims = torch.unbind(im[:, :, slices], -1) # extract n slices
+        ims = [i.transpose(0,1).flip(0) for i in ims] # rotate slices 90 degrees
+        if len(slices) > 4 and len(slices) % 2 == 0: 
+            n = len(slices) // 2
+            ims1 = torch.cat(ims[0:n], 1)
+            ims2 = torch.cat(ims[n:], 1)
+            return torch.cat([ims1, ims2], 0)
+        else: 
+            return torch.cat(ims, 1) # create tile
+        
+    def plot_images(self, fns, slices, cmap='Greys_r', figsize=15, **kwargs): 
+        ims = [torch.load(os.path.join(self.out_dir, 'preds', fn)).cpu().argmax(0) for fn in fns]
+        ims = [self.get_slices(im, slices) for im in ims]
+        ims = torch.cat(ims, 0)
+        plt.figure(figsize=(figsize,figsize)) 
+        plt.imshow(ims, cmap=cmap, **kwargs)
+        plt.axis('off')
+        
+    def load_segmentation_image(self, fn): 
+        im = torch.load(fn).cpu().unsqueeze(0)
+        im = torch.nn.functional.interpolate(im, (224, 224, 112))
+        im = im.argmax(1).squeeze()
+        im = self.get_slices(im, slices = (40, 48, 56, 74, 82, 90)) 
+        im = im/im.max() * 255
+        return im
+    
+    def generate_gif(self):
+        with imageio.get_writer(
+            os.path.join(self.run_id,'report','progress.gif'),
+            mode='I', 
+            fps = max(self.train_logs.epoch) // 10) as writer: # make gif 10 seconds
+            for epoch in tqdm.tqdm(list(self.train_logs.epoch.unique())):
+                seg_fn = os.path.join(self.out_dir, 'preds', f"pred_epoch_{epoch}.pt")
+                if os.path.exists(seg_fn): im = self.load_segmentation_image(seg_fn)
+
+                plt_train_logs = self.train_logs[self.train_logs.epoch <= epoch]
+                loss_plt = self.plot_loss(plt_train_logs, self.metric_logs[:epoch])
+                loss_fig = self.get_arr_from_fig(loss_plt)[:,:,0]
+
+                new_shape = im.shape[1], int(loss_fig.shape[0] / loss_fig.shape[1] * im.shape[1])
+                loss_fig = cv2.resize(loss_fig, (im.shape[1], im.shape[0]))
+
+                images = torch.cat([im, torch.tensor(loss_fig)], 0).numpy().astype(np.uint8)  
+                writer.append_data(images)
+
+                loss_plt.clear()
+                plt.close()

prostate/train.py

@@ -0,0 +1,482 @@
+import os
+import yaml
+import munch
+import torch
+import ignite
+import monai
+import shutil
+import pandas as pd
+
+from typing import Union, List, Callable
+from ignite.contrib.handlers.tqdm_logger import ProgressBar
+from monai.handlers import (
+    CheckpointSaver,
+    StatsHandler,
+    TensorBoardStatsHandler,
+    TensorBoardImageHandler,
+    ValidationHandler,
+    from_engine,
+    MeanDice,
+    EarlyStopHandler,
+    MetricLogger,
+    MetricsSaver
+)
+
+from .data import segmentation_dataloaders
+from .model import get_model
+from .optimizer import get_optimizer
+from .loss import get_loss
+from .transforms import get_val_post_transforms
+from .utils import USE_AMP
+
+def loss_logger(engine): 
+    "write loss and lr of each iteration/epoch to file"
+    iteration=engine.state.iteration
+    epoch=engine.state.epoch
+    loss=[o['loss'] for o in engine.state.output]
+    loss=sum(loss)/len(loss)
+    lr=engine.optimizer.param_groups[0]['lr']
+    log_file=os.path.join(engine.config.log_dir, 'train_logs.csv')
+    if not os.path.exists(log_file): 
+        with open(log_file, 'w+') as f: 
+            f.write('iteration,epoch,loss,lr\n')
+    with open(log_file, 'a') as f: 
+        f.write(f'{iteration},{epoch},{loss},{lr}\n')
+        
+def metric_logger(engine): 
+    "write `metrics` after each epoch to file"
+    if engine.state.epoch > 1: # only key metric is calcualted in 1st epoch, needs fix
+        metric_names=[k for k in engine.state.metrics.keys()]
+        metrics=[str(engine.state.metrics[mn]) for mn in metric_names]
+        log_file=os.path.join(engine.config.log_dir, 'metric_logs.csv')
+        if not os.path.exists(log_file): 
+            with open(log_file, 'w+') as f: 
+                f.write(','.join(metric_names) + '\n')
+        with open(log_file, 'a') as f: 
+            f.write(','.join(metrics) + '\n')
+        
+def pred_logger(engine):
+    "save `pred` each time metric improves"
+    epoch=engine.state.epoch
+    root = os.path.join(engine.config.out_dir, 'preds')
+    if not os.path.exists(root): 
+        os.makedirs(root) 
+        torch.save(
+            engine.state.output[0]['label'], 
+            os.path.join(root, f'label.pt')
+        )
+        torch.save(
+            engine.state.output[0]['image'], 
+            os.path.join(root, f'image.pt')
+        )
+        
+    if epoch==engine.state.best_metric_epoch:
+        torch.save(
+            engine.state.output[0]['pred'], 
+            os.path.join(root, f'pred_epoch_{epoch}.pt')
+        )
+
+    
+def get_val_handlers(
+    network: torch.nn.Module,
+    config: dict
+    ) -> list:
+    """Create default handlers for model validation
+    Args: 
+        network: 
+            nn.Module subclass, the model to train
+    
+    Returns:
+        a list of default handlers for validation: [
+            StatsHandler: 
+                ???
+            TensorBoardStatsHandler: 
+                Save loss from validation to `config.log_dir`, allow logging with TensorBoard
+            CheckpointSaver: 
+                Save best model to `config.model_dir`
+        ]
+    """
+    
+    val_handlers=[      
+        StatsHandler(
+            tag_name="metric_logger",
+            epoch_print_logger=metric_logger,
+            output_transform=lambda x: None
+        ),
+        StatsHandler(
+            tag_name="pred_logger",
+            epoch_print_logger=pred_logger,
+            output_transform=lambda x: None
+        ),
+        TensorBoardStatsHandler(
+            log_dir=config.log_dir, 
+            # tag_name="val_mean_dice",
+            output_transform=lambda x: None
+        ),
+        TensorBoardImageHandler(
+            log_dir=config.log_dir,
+            batch_transform=from_engine(["image", "label"]),
+            output_transform=from_engine(["pred"]),
+        ),
+        CheckpointSaver(
+            save_dir=config.model_dir, 
+            save_dict={f"network_{config.run_id}": network}, 
+            save_key_metric=True
+        ),
+
+    ]
+    
+    return val_handlers
+
+
+def get_train_handlers(
+    evaluator: monai.engines.SupervisedEvaluator,
+    config: dict
+    ) -> list: 
+    """Create default handlers for model training
+    Args: 
+        evaluator: an engine of type `monai.engines.SupervisedEvaluator` for evaluations
+        every epoch
+        
+    Returns:
+        list of default handlers for training: [
+            ValidationHandler:
+                Allows model validation every epoch
+            StatsHandler:
+                ???
+            TensorBoardStatsHandler: 
+                Save loss from validation to `config.log_dir`, allow logging with TensorBoard
+        ]
+    """
+    
+    train_handlers=[
+        ValidationHandler(
+            validator=evaluator, 
+            interval=1, 
+            epoch_level=True
+        ),
+        StatsHandler(
+            tag_name="train_loss", 
+            output_transform=from_engine(
+                ["loss"], 
+                first=True
+            )
+        ),
+        StatsHandler(
+            tag_name='loss_logger', 
+            iteration_print_logger=loss_logger
+        ), 
+        TensorBoardStatsHandler(
+            log_dir=config.log_dir,
+            tag_name="train_loss",
+            output_transform=from_engine(
+                ["loss"], 
+                first=True
+            ),
+        )
+    ]
+    
+    return train_handlers
+
+def get_evaluator(
+    config: dict,
+    device: torch.device , 
+    network: torch.nn.Module, 
+    val_data_loader: monai.data.dataloader.DataLoader, 
+    val_post_transforms: monai.transforms.compose.Compose,
+    val_handlers: Union[Callable, List]=get_val_handlers
+) -> monai.engines.SupervisedEvaluator: 
+    
+    """Create default evaluator for training of a segmentation model
+    Args: 
+        device: 
+            torch.cuda.device for model and engine
+        network: 
+            nn.Module subclass, the model to train
+        val_data_loader: 
+            Validation data loader, `monai.data.dataloader.DataLoader` subclass
+        val_post_transforms: 
+            function to create transforms OR composed transforms
+        val_handlers: 
+            function to create handerls OR List of handlers
+            
+    Returns: 
+        default evaluator for segmentation of type `monai.engines.SupervisedEvaluator`
+    """
+    
+    if callable(val_handlers): val_handlers=val_handlers()
+    
+    evaluator=monai.engines.SupervisedEvaluator(
+        device=device,
+        val_data_loader=val_data_loader,
+        network=network,
+        inferer=monai.inferers.SlidingWindowInferer(
+            roi_size=(96, 96, 96), 
+            sw_batch_size=4, 
+            overlap=0.5
+        ),
+        postprocessing=val_post_transforms,
+        key_val_metric={
+            "val_mean_dice": MeanDice(
+                include_background=False, 
+                output_transform=from_engine(
+                    ["pred", "label"]
+                )
+            )
+        },
+        val_handlers=val_handlers,
+        # if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
+        amp=USE_AMP,
+    )
+    evaluator.config=config
+    return evaluator
+
+
+class SegmentationTrainer(monai.engines.SupervisedTrainer): 
+    "Default Trainer für supervised segmentation task"
+    def __init__(self, 
+                 config: dict,
+                 progress_bar: bool=True, 
+                 early_stopping: bool=True, 
+                 metrics: list=["MeanDice", "HausdorffDistance", "SurfaceDistance"],
+                 save_latest_metrics: bool=True
+                ):
+        self.config=config
+        self._prepare_dirs()
+        self.config.device=torch.device(self.config.device)
+        
+        train_loader, val_loader=segmentation_dataloaders(
+            config=config, 
+            train=True, 
+            valid=True, 
+            test=False
+        )
+        network=get_model(config=config).to(config.device)
+        optimizer=get_optimizer(
+            network, 
+            config=config
+        )
+        loss_fn=get_loss(config=config)
+        val_post_transforms=get_val_post_transforms(config=config)
+        val_handlers=get_val_handlers(
+            network, 
+            config=config
+        )
+        
+        self.evaluator=get_evaluator(
+            config=config,
+            device=config.device, 
+            network=network, 
+            val_data_loader=val_loader, 
+            val_post_transforms=val_post_transforms,
+            val_handlers=val_handlers,
+
+        )
+        train_handlers=get_train_handlers(
+            self.evaluator, 
+            config=config
+        )
+        
+        super().__init__(
+            device=config.device, 
+            max_epochs=self.config.training.max_epochs, 
+            train_data_loader=train_loader,
+            network=network,
+            optimizer=optimizer, 
+            loss_function=loss_fn, 
+            inferer=monai.inferers.SimpleInferer(), 
+            train_handlers=train_handlers,
+            amp=USE_AMP,
+        )        
+        
+        if early_stopping: self._add_early_stopping()
+        if progress_bar: self._add_progress_bars()
+        
+        self.schedulers=[]
+        # add different metrics dynamically
+        for m in metrics: 
+            getattr(monai.handlers, m)(
+                include_background=False, 
+                reduction="mean", 
+                output_transform=from_engine(
+                    ["pred", "label"]
+                )
+            ).attach(self.evaluator, m)
+            
+        self._add_metrics_logger()
+        # add eval loss to metrics
+        self._add_eval_loss()
+        
+        if save_latest_metrics: self._add_metrics_saver()
+
+    
+    def _prepare_dirs(self)->None:
+        # create run_id, copy config file for reproducibility
+        os.makedirs(self.config.run_id, exist_ok=True)
+        with open(
+            os.path.join(
+                self.config.run_id, 
+                'config.yaml'
+            ), 'w+') as f: 
+            f.write(yaml.safe_dump(self.config))
+        
+        # delete old log_dir
+        if os.path.exists(self.config.log_dir): 
+            shutil.rmtree(self.config.log_dir) 
+            
+    def _add_early_stopping(self) -> None: 
+        early_stopping=EarlyStopHandler(
+            patience=self.config.training.early_stopping_patience, 
+            min_delta=1e-4,
+            score_function=lambda x: x.state.metrics[x.state.key_metric_name], 
+            trainer=self
+        )
+        self.evaluator.add_event_handler(
+            ignite.engine.Events.COMPLETED, 
+            early_stopping
+        )
+        
+    def _add_metrics_logger(self) -> None: 
+        self.metric_logger=MetricLogger(
+            evaluator=self.evaluator
+        )
+        self.metric_logger.attach(self)
+
+    def _add_progress_bars(self) -> None:
+        trainer_pbar=ProgressBar()
+        evaluator_pbar=ProgressBar(
+            colour='green'
+        )
+        trainer_pbar.attach(
+            self, 
+            output_transform=lambda output:{
+                'loss': torch.tensor(
+                    [x['loss'] for x in output]
+                ).mean()
+            }
+        )
+        evaluator_pbar.attach(self.evaluator)
+        
+    def _add_metrics_saver(self) -> None:
+        metric_saver=MetricsSaver(
+            save_dir=self.config.out_dir,
+            metric_details='*',
+            batch_transform=self._get_meta_dict,
+            delimiter=','
+        )
+        metric_saver.attach(self.evaluator)
+    
+    def _add_eval_loss(self)->None: 
+        # TODO improve by adding this to val handlers
+        eval_loss_handler=ignite.metrics.Loss(
+            loss_fn=self.loss_function,
+            output_transform=lambda output: (
+                output[0]['pred'].unsqueeze(0), # add batch dim
+                output[0]['label'].argmax(0, keepdim=True).unsqueeze(0) # reverse one-hot, add batch dim
+            )
+        )
+        eval_loss_handler.attach(self.evaluator, 'eval_loss')
+        
+    def _get_meta_dict(self, batch) -> list: 
+        "Get dict of metadata from engine. Needed as `batch_transform`"
+        image_cols=self.config.data.image_cols
+        image_name=image_cols[0] if isinstance(image_cols, list) else image_cols
+        key=f'{image_name}_meta_dict'
+        return [item[key] for item in batch]   
+        
+    def load_checkpoint(self, checkpoint=None): 
+        if not checkpoint: 
+            # get name of last checkpoint
+            checkpoint = os.path.join(
+                self.config.model_dir, 
+                f"network_{self.config.run_id}_key_metric={self.evaluator.state.best_metric:.4f}.pt"
+            )
+        self.network.load_state_dict(
+            torch.load(checkpoint)
+        )
+        
+    def run(self, try_resume_from_checkpoint=True) -> None:
+        """Run training, if `try_resume_from_checkpoint` tries to
+        load previous checkpoint stored at `self.config.model_dir`
+        """ 
+
+        if try_resume_from_checkpoint: 
+            checkpoints = [
+                os.path.join(
+                    self.config.model_dir, 
+                    checkpoint_name
+                ) for checkpoint_name in os.listdir(
+                    self.config.model_dir
+                ) if self.config.run_id in checkpoint_name
+            ]
+            try: 
+                checkpoint = sorted(checkpoints)[-1]
+                self.load_checkpoint(checkpoint)
+                print(f"resuming from previous checkpoint at {checkpoint}")
+            except: pass # train from scratch
+        
+        # train the model
+        super().run()
+        
+        # make metrics and losses more accessible
+        self.loss={
+            "iter": [_iter for _iter, _ in  self.metric_logger.loss],
+            "loss": [_loss for _, _loss in  self.metric_logger.loss],
+            "epoch": [_iter // self.state.epoch_length for _iter, _ in  self.metric_logger.loss]
+            }        
+        
+        self.metrics={
+            k: [item[1] for item in self.metric_logger.metrics[k]] for k in 
+                   self.evaluator.state.metric_details.keys()
+        }
+       # pd.DataFrame(self.metrics).to_csv(f"{self.config.out_dir}/metric_logs.csv")
+       # pd.DataFrame(self.loss).to_csv(f"{self.config.out_dir}/loss_logs.csv") 
+        
+    def fit_one_cycle(self, try_resume_from_checkpoint=True) -> None:
+        "Run training using one-cycle-policy"
+        assert "FitOneCycle" not in self.schedulers, "FitOneCycle already added"
+        fit_one_cycle=monai.handlers.LrScheduleHandler(
+            torch.optim.lr_scheduler.OneCycleLR(
+                optimizer=self.optimizer, 
+                max_lr=self.optimizer.param_groups[0]['lr'], 
+                steps_per_epoch=self.state.epoch_length, 
+                epochs=self.state.max_epochs
+            ),
+            epoch_level=False,
+            name="FitOneCycle"
+        )
+        fit_one_cycle.attach(self)
+        self.schedulers += ["FitOneCycle"] 
+    
+    def reduce_lr_on_plateau(self, 
+                             try_resume_from_checkpoint=True, 
+                             factor=0.1, 
+                             patience=10, 
+                             min_lr=1e-10, 
+                             verbose=True) -> None:
+        "Reduce learning rate by `factor` every `patience` epochs if kex_metric does not improve"
+        assert "ReduceLROnPlateau" not in self.schedulers, "ReduceLROnPlateau already added"
+        reduce_lr_on_plateau=monai.handlers.LrScheduleHandler(
+            torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer=self.optimizer, 
+                factor=factor, 
+                patience=patience, 
+                min_lr=min_lr, 
+                verbose=verbose
+            ), 
+            print_lr=True,
+            name='ReduceLROnPlateau', 
+            epoch_level=True, 
+            step_transform=lambda engine: engine.state.metrics[engine.state.key_metric_name],
+        )
+        reduce_lr_on_plateau.attach(self.evaluator)
+        self.schedulers += ["ReduceLROnPlateau"] 
+
+    def evaluate(self, checkpoint=None, dataloader=None): 
+        "Run evaluation with best saved checkpoint"
+        self.load_checkpoint(checkpoint)
+        if dataloader: 
+            self.evaluator.set_data(dataloader)
+            self.evaluator.state.epoch_length=len(dataloader)
+        self.evaluator.run()
+        print(f"metrics saved to {self.config.out_dir}") 

prostate/transforms.py

@@ -0,0 +1,363 @@
+# create transforms for training, validation and test dataset
+
+## TODO: Make Transforms more dynamic by directly building from config args
+## Maybe like this
+## TFM_NAME=config.transforms.keys()[0]
+## tfm_fun=getattr(monai.transforms, TFM_NAME)
+## tmfs+=[tfms_fun(keys=image+cols, **config.transforms[TFM_NAME], prob=prob, mode=mode)
+
+
+## ---------- imports ----------
+import os
+# only import of base transforms, others are imported as needed
+from monai.utils.enums import CommonKeys
+from monai.transforms import (
+    Activationsd,
+    AsDiscreted,
+    Compose,
+    ConcatItemsd,
+    KeepLargestConnectedComponentd,
+    LoadImaged,
+    EnsureChannelFirstd,
+    EnsureTyped,
+    SaveImaged,
+    ScaleIntensityd,
+    NormalizeIntensityd
+)
+# images should be interploated with `bilinear` but masks with `nearest`
+
+## ---------- base transforms ----------
+# applied everytime
+def get_base_transforms(
+    config: dict,
+    minv: int=0, 
+    maxv: int=1
+)->list:
+    
+    tfms=[]
+    tfms+=[LoadImaged(keys=config.data.image_cols+config.data.label_cols)]
+    tfms+=[EnsureChannelFirstd(keys=config.data.image_cols+config.data.label_cols)]
+    if config.transforms.spacing:
+        from monai.transforms import Spacingd
+        tfms+=[
+            Spacingd(
+                keys=config.data.image_cols+config.data.label_cols,
+                pixdim=config.transforms.spacing,
+                mode=config.transforms.mode
+            )
+        ]
+    if config.transforms.orientation:
+        from monai.transforms import Orientationd
+        tfms+=[
+            Orientationd(
+                keys=config.data.image_cols+config.data.label_cols,
+                axcodes=config.transforms.orientation
+            )
+        ]
+    tfms+=[
+        ScaleIntensityd(
+            keys=config.data.image_cols,
+            minv=minv,
+            maxv=maxv
+        )
+    ]
+    tfms+=[NormalizeIntensityd(keys=config.data.image_cols)]
+    return tfms
+
+## ---------- train transforms ----------
+
+def get_train_transforms(config: dict):
+    tfms=get_base_transforms(config=config)
+
+    # ---------- specific transforms for mri ----------
+    if 'rand_bias_field' in config.transforms.keys():
+        from monai.transforms import RandBiasFieldd
+        args=config.transforms.rand_bias_field
+        tfms+=[
+            RandBiasFieldd(
+                keys=config.data.image_cols,
+                degree=args['degree'],
+                coeff_range=args['coeff_range'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_gaussian_smooth' in config.transforms.keys():
+        from monai.transforms import RandGaussianSmoothd
+        args=config.transforms.rand_gaussian_smooth
+        tfms+=[
+            RandGaussianSmoothd(
+                keys=config.data.image_cols,
+                sigma_x=args['sigma_x'],
+                sigma_y=args['sigma_y'],
+                sigma_z=args['sigma_z'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_gibbs_nose' in config.transforms.keys():
+        from monai.transforms import RandGibbsNoised
+        args=config.transforms.rand_gibbs_nose
+        tfms+=[
+            RandGibbsNoised(
+                keys=config.data.image_cols,
+                alpha=args['alpha'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    # ---------- affine transforms ----------
+
+    if 'rand_affine' in config.transforms.keys():
+        from monai.transforms import RandAffined
+        args=config.transforms.rand_affine
+        tfms+=[
+            RandAffined(
+                keys=config.data.image_cols+config.data.label_cols,
+                rotate_range=args['rotate_range'],
+                shear_range=args['shear_range'],
+                translate_range=args['translate_range'],
+                mode=config.transforms.mode,
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_rotate90' in config.transforms.keys():
+        from monai.transforms import RandRotate90d
+        args=config.transforms.rand_rotate90
+        tfms+=[
+            RandRotate90d(
+                keys=config.data.image_cols+config.data.label_cols,
+                spatial_axes=args['spatial_axes'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_rotate' in config.transforms.keys():
+        from monai.transforms import RandRotated
+        args=config.transforms.rand_rotate
+        tfms+=[
+            RandRotated(
+                keys=config.data.image_cols+config.data.label_cols,
+                range_x=args['range_x'],
+                range_y=args['range_y'],
+                range_z=args['range_z'],
+                mode=config.transforms.mode,
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_elastic' in config.transforms.keys():
+        if config['ndim'] == 3:
+            from monai.transforms import Rand3DElasticd as RandElasticd
+        elif config['ndim'] == 2:
+            from monai.transforms import Rand2DElasticd as RandElasticd
+        args=config.transforms.rand_elastic
+        tfms+=[
+            RandElasticd(
+                keys=config.data.image_cols+config.data.label_cols,
+                sigma_range=args['sigma_range'],
+                magnitude_range=args['magnitude_range'],
+                rotate_range=args['rotate_range'],
+                shear_range=args['shear_range'],
+                translate_range=args['translate_range'],
+                mode=config.transforms.mode,
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'rand_zoom' in config.transforms.keys():
+        from monai.transforms import RandZoomd
+        args=config.transforms.rand_zoom
+        tfms+=[
+            RandZoomd(
+                keys=config.data.image_cols+config.data.label_cols,
+                min_zoom=args['min'],
+                max_zoom=args['max'],
+                mode=['area' if x == 'bilinear' else x for x in config.transforms.mode],
+                prob=config.transforms.prob
+            )
+        ]
+
+    # ---------- random cropping, very effective for large images ----------
+    # RandCropByPosNegLabeld is not advisable for data with missing lables
+    # e.g., segmentation of carcinomas which are not present on all images
+    # thus fallback to RandSpatialCropSamplesd. Completly replacing Cropping
+    # by just resizing could be discussed, but I believe it is not beneficial
+    # For the first version, this is an ungly hack. For the second version, 
+    # a better verion for transforms should be written. 
+
+    if 'rand_crop_pos_neg_label' in config.transforms.keys():
+        from monai.transforms import RandCropByPosNegLabeld
+        args=config.transforms.rand_crop_pos_neg_label
+        tfms+=[
+            RandCropByPosNegLabeld(
+                keys=config.data.image_cols+config.data.label_cols,
+                label_key=config.data.label_cols[0],
+                spatial_size=args['spatial_size'],
+                pos=args['pos'],
+                neg=args['neg'],
+                num_samples=args['num_samples'],
+                image_key=config.data.image_cols[0],
+                image_threshold=0,
+            )
+        ]
+        
+    elif 'rand_spatial_crop_samples' in config.transforms.keys():
+        from monai.transforms import RandSpatialCropSamplesd
+        args=config.transforms.rand_spatial_crop_samples
+        tfms+=[
+            RandSpatialCropSamplesd(
+                keys=config.data.image_cols+config.data.label_cols,
+                roi_size=args['roi_size'],
+                random_size=False,
+                num_samples=args['num_samples'],
+            )
+        ]
+        
+    else: 
+        raise ValueError('Either `rand_crop_pos_neg_label` or `rand_spatial_crop_samples` '\
+                         'need to be specified')
+        
+    # ---------- intensity transforms ----------
+
+    if 'gaussian_noise' in config.transforms.keys():
+        from monai.transforms import RandGaussianNoised
+        args=config.transforms.gaussian_noise
+        tfms+=[
+            RandGaussianNoised(
+                keys=config.data.image_cols,
+                mean=args['mean'],
+                std=args['std'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'shift_intensity' in config.transforms.keys():
+        from monai.transforms import RandShiftIntensityd
+        args=config.transforms.shift_intensity
+        tfms+=[
+            RandShiftIntensityd(
+                keys=config.data.image_cols,
+                offsets=args['offsets'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'gaussian_sharpen' in config.transforms.keys():
+        from monai.transforms import RandGaussianSharpend
+        args=config.transforms.gaussian_sharpen
+        tfms+=[
+            RandGaussianSharpend(
+                keys=config.data.image_cols,
+                sigma1_x=args['sigma1_x'],
+                sigma1_y=args['sigma1_y'],
+                sigma1_z=args['sigma1_z'],
+                sigma2_x=args['sigma2_x'],
+                sigma2_y=args['sigma2_y'],
+                sigma2_z=args['sigma2_z'],
+                alpha=args['alpha'],
+                prob=config.transforms.prob
+            )
+        ]
+
+    if 'adjust_contrast' in config.transforms.keys():
+        from monai.transforms import RandAdjustContrastd
+        args=config.transforms.adjust_contrast
+        tfms+=[
+            RandAdjustContrastd(
+                keys=config.data.image_cols,
+                gamma=args['gamma'],
+                prob=config.transforms.prob
+            )
+        ]
+        
+    # Concat mutlisequence data to single Tensors on the ChannelDim
+    # Rename images to `CommonKeys.IMAGE` and labels to `CommonKeys.LABELS`
+    # for more compatibility with monai.engines
+    
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.image_cols, 
+            name=CommonKeys.IMAGE, 
+            dim=0
+        )
+    ]
+
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.label_cols, 
+            name=CommonKeys.LABEL, 
+            dim=0
+        )
+    ]
+
+    return Compose(tfms)
+
+## ---------- valid transforms ----------
+
+def get_val_transforms(config: dict):
+    tfms=get_base_transforms(config=config)
+    tfms+=[EnsureTyped(keys=config.data.image_cols+config.data.label_cols)]
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.image_cols, 
+            name=CommonKeys.IMAGE, 
+            dim=0
+        )
+    ]
+
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.label_cols, 
+            name=CommonKeys.LABEL, 
+            dim=0
+        )
+    ]
+    
+    return Compose(tfms)
+
+## ---------- test transforms ----------
+# same as valid transforms
+
+def get_test_transforms(config: dict):
+    tfms=get_base_transforms(config=config)
+    tfms+=[EnsureTyped(keys=config.data.image_cols+config.data.label_cols)]
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.image_cols, 
+            name=CommonKeys.IMAGE, 
+            dim=0
+        )
+    ]
+
+    tfms+=[
+        ConcatItemsd(
+            keys=config.data.label_cols, 
+            name=CommonKeys.LABEL, 
+            dim=0
+        )
+    ]
+    
+    return Compose(tfms)
+
+
+def get_val_post_transforms(config: dict): 
+    tfms=[EnsureTyped(keys=[CommonKeys.PRED, CommonKeys.LABEL]),
+            AsDiscreted(
+                keys=CommonKeys.PRED, 
+                argmax=True, 
+                to_onehot=config.model.out_channels, 
+                num_classes=config.model.out_channels
+            ),
+            AsDiscreted(
+                keys=CommonKeys.LABEL, 
+                to_onehot=config.model.out_channels, 
+                num_classes=config.model.out_channels
+            ),
+            KeepLargestConnectedComponentd(
+                keys=CommonKeys.PRED, 
+                applied_labels=list(range(1, config.model.out_channels))
+           ),
+            ]
+    return Compose(tfms)

prostate/utils.py

@@ -0,0 +1,56 @@
+import os
+import yaml
+import torch
+import monai
+import munch 
+
+def load_config(fn: str='config.yaml'):
+    "Load config from YAML and return a serialized dictionary object"
+    with open(fn, 'r') as stream:
+        config=yaml.safe_load(stream)
+    config=munch.munchify(config)
+    
+    if not config.overwrite:
+        i=1
+        while os.path.exists(config.run_id+f'_{i}'):
+            i+=1
+        config.run_id+=f'_{i}'
+
+    config.out_dir = os.path.join(config.run_id, config.out_dir)
+    config.log_dir = os.path.join(config.run_id, config.log_dir)
+    
+    if not isinstance(config.data.image_cols, (tuple, list)): 
+        config.data.image_cols=[config.data.image_cols]
+    if not isinstance(config.data.label_cols, (tuple, list)): 
+        config.data.label_cols=[config.data.label_cols]
+    
+    config.transforms.mode=('bilinear', ) * len(config.data.image_cols) + \
+                             ('nearest', ) * len(config.data.label_cols)
+    return config
+
+    
+def num_workers():
+    "Get max supported workers -2 for multiprocessing"
+    import resource
+    import multiprocessing
+    
+    # first check for max number of open files allowed on system
+    soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE)
+    n_workers=multiprocessing.cpu_count() - 2
+    # giving each worker at least 256 open processes should allow them to run smoothly
+    max_workers = soft_limit // 256
+    if max_workers < n_workers: 
+        print(
+            "Will not use all available workers as number of allowed open files is to small"
+            "to ensure smooth multiprocessing. Current limits are:\n" 
+            f"\t soft_limit: {soft_limit}\n"
+            f"\t hard_limit: {hard_limit}\n" 
+            "try increasing the limits to at least {256*n_workers}." 
+            "See https://superuser.com/questions/1200539/cannot-increase-open-file-limit-past-4096-ubuntu"
+            "for more details"
+        )
+        return max_workers
+              
+    return n_workers
+
+USE_AMP=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False

prostate/viewer.py

@@ -0,0 +1,352 @@
+import torch
+import ipywidgets
+import numpy as np
+import matplotlib.pyplot as plt
+from IPython.display import display
+from itertools import chain, islice
+from ipywidgets import interactive, widgets
+
+def _create_label(text:str)->ipywidgets.widgets.Label:
+    "Create label widget"
+    
+    label = widgets.Label(
+            text,
+            layout=widgets.Layout(
+            width='100%', 
+            display='flex', 
+            justify_content="center"
+            )
+        )
+    return label
+        
+def _create_slider(
+    slider_min: int, 
+    slider_max: int, 
+    value: int, 
+    step: int=1, 
+    description:str ='',
+    continuous_update: bool=True, 
+    readout: bool=False, 
+    slider_type: str='IntSlider',
+    **kwargs)->ipywidgets.widgets:
+    "Create slider widget"
+    
+    slider = getattr(widgets, slider_type)(
+        min=slider_min, 
+        max=slider_max, 
+        step=step, 
+        value=value,
+        description=description, 
+        continuous_update=continuous_update,
+        readout = readout,
+        layout=widgets.Layout(width='99%', min_width='200px'),
+        style={'description_width': 'initial'}, 
+        **kwargs
+    )   
+    return slider
+    
+def _create_button(description:str)->ipywidgets.widgets.Button:
+    "Create button widget"
+    button = widgets.Button(
+            description=description,
+            layout=widgets.Layout(
+                width='95%',
+                margin='5px 5px'
+            )
+        )
+    return button
+
+def _create_togglebutton(description: str, 
+                         value: int, 
+                         **kwargs)->ipywidgets.widgets.Button:
+    "Create toggle button widget"
+    button = widgets.ToggleButton(
+            description=description,
+            value = value,
+            layout=widgets.Layout(
+                width='95%',
+                margin='5px 5px'
+            ), **kwargs
+        )
+    return button
+
+
+class BasicViewer():
+    """ Base class for viewing TensorDicom3D objects. 
+    
+    Args: 
+        x: main image object to view as rank 3 tensor
+        y: either a segmentation mask as as rank 3 tensor or a label as str. 
+        prediction: a class predicton as str
+        description: description of the whole image
+        figsize: size of image, passed as plotting argument
+        cmap: colormap for the image
+    Returns: 
+        Instance of BasicViewer
+    """
+    
+    def __init__(self, x:torch.Tensor, y=None, prediction:str=None, description: str=None, 
+                 figsize=(3, 3), cmap:str='bone'):
+        assert x.ndim == 3, f"x.ndim needs to be equal to but is {x.ndim}"
+        if isinstance(y, torch.Tensor): 
+            assert x.shape == y.shape, f"Shapes of x {x.shape} and y {y.shape} do not match"
+        self.x=x
+        self.y=y
+        self.prediction=prediction
+        self.description=description
+        self.figsize=figsize
+        self.cmap=cmap
+        self.with_mask = isinstance(y, torch.Tensor)
+        self.slice_range = (1, len(x)) # len(x) == im.shape[0]
+                            
+    def _plot_slice(self, im_slice, with_mask, px_range):
+        "Plot slice of image"
+        fig, ax = plt.subplots(1, 1, figsize=self.figsize) 
+        ax.imshow(self.x[im_slice-1, :, :].clip(*px_range), cmap=self.cmap)
+        if isinstance(self.y, (torch.Tensor)) and with_mask:
+            ax.imshow(self.y[im_slice-1, :, :], cmap='jet', alpha = 0.25)
+        plt.axis('off')
+        ax.set_xticks([])
+        ax.set_yticks([])
+        plt.show()
+    
+    def _create_image_box(self, figsize):
+        "Create widget items, order them in item_box and generate view box"
+        items = []
+        
+        if self.description: plot_description = _create_label(self.description)
+        
+        if isinstance(self.y, str):
+            label = f'{self.y} | {self.prediction}' if self.prediction else self.y
+            if self.prediction: 
+                font_color = 'green' if self.y == self.prediction else 'red'
+                y_label = _create_label(r'\(\color{' + font_color + '} {' + label  + '}\)') 
+            else: 
+                y_label = _create_label(label)
+        else: y_label = _create_label(' ')
+            
+        slice_slider = _create_slider(
+            slider_min = min(self.slice_range), 
+            slider_max = max(self.slice_range), 
+            value = max(self.slice_range)//2, 
+            readout = True)
+
+        toggle_mask_button = _create_togglebutton('Show Mask', True)
+        
+        range_slider = _create_slider(
+            slider_min = self.x.min().numpy(), 
+            slider_max = self.x.max().numpy(),  
+            value = [self.x.min().numpy(), self.x.max().numpy()], 
+            slider_type = 'FloatRangeSlider' if torch.is_floating_point(self.x) else 'IntRandSlider', 
+            step = 0.01 if torch.is_floating_point(self.x) else 1, 
+            readout=True)
+        
+        image_output = widgets.interactive_output(
+            f = self._plot_slice,
+            controls = {'im_slice': slice_slider, 
+                'with_mask': toggle_mask_button, 
+                'px_range': range_slider})
+
+        image_output.layout.height = f'{self.figsize[0]/1.2}in' # suppress flickering
+        image_output.layout.width = f'{self.figsize[1]/1.2}in' # suppress flickering
+        
+        if self.description: items.append(plot_description)
+        items.append(y_label)
+        items.append(range_slider)        
+        items.append(image_output)
+        if isinstance(self.y, torch.Tensor): 
+            slice_slider = widgets.HBox([slice_slider, toggle_mask_button])     
+        items.append(slice_slider)        
+                    
+        image_box=widgets.VBox(
+            items,
+            layout = widgets.Layout(
+                border = 'none', 
+                margin = '10px 5px 0px 0px', 
+                padding =  '5px'))
+        
+        return image_box
+    
+    def _generate_views(self):
+        image_box = self._create_image_box(self.figsize)
+        self.box = widgets.HBox(children=[image_box])
+
+    @property
+    def image_box(self):
+        return self._create_image_box(self.figsize)
+    
+    def show(self):
+        self._generate_views()
+        plt.style.use('default')
+        display(self.box)
+        
+        
+class DicomExplorer(BasicViewer):
+    """ DICOM viewer for basic image analysis inside iPython notebooks. 
+    Can display a single 3D volume together with a segmentation mask, a histogram 
+    of voxel/pixel values and some summary statistics. 
+    Allows simple windowing by clipping the pixel/voxel values to a region, which 
+    can be manually specified.    
+    
+    """
+    
+    vbox_layout = widgets.Layout(
+        margin = '10px 5px 5px 5px', 
+        padding =  '5px',
+        display='flex',
+        flex_flow='column',
+        align_items='center', 
+        min_width = '250px')
+            
+    def _plot_hist(self, px_range):
+        x = self.x.numpy().flatten()
+        fig, ax = plt.subplots(figsize=self.figsize)
+        N, bins, patches = plt.hist(x, 100, color='grey')
+        lwr = int(px_range[0] * 100/max(x))
+        upr = int(np.ceil(px_range[1] * 100/max(x)))
+        
+        for i in range(0,lwr):
+            patches[i].set_facecolor('grey' if lwr > 0 else 'darkblue')
+        for i in range(lwr, upr):
+            patches[i].set_facecolor('darkblue')
+        for i in range(upr,100):
+            patches[i].set_facecolor('grey' if upr < 100 else 'darkblue')
+        
+        plt.show()
+        
+    def _image_summary(self, px_range):
+        x = self.x.clip(*px_range)
+
+        diffs = x - x.mean()
+        var = torch.mean(torch.pow(diffs, 2.0))
+        std = torch.pow(var, 0.5)
+        zscores = diffs / std
+        skews = torch.mean(torch.pow(zscores, 3.0))
+        kurt = torch.mean(torch.pow(zscores, 4.0)) - 3.0
+
+        table = f'Statistics:\n' + \
+                f'  Mean px value:     {x.mean()} \n' + \
+                f'  Std of px values:  {x.std()} \n' + \
+                f'  Min px value:      {x.min()} \n' + \
+                f'  Max px value:      {x.max()} \n' + \
+                f'  Median px value:   {x.median()} \n' + \
+                f'  Skewness:          {skews} \n' + \
+                f'  Kurtosis:          {kurt} \n\n' + \
+                f'Tensor properties \n' + \
+                f'  Tensor shape:      {tuple(x.shape)}\n' + \
+                f'  Tensor dtype:      {x.dtype}'
+        print(table)
+        
+    def _generate_views(self):   
+        
+        slice_slider = _create_slider(
+            slider_min = min(self.slice_range), 
+            slider_max = max(self.slice_range), 
+            value = max(self.slice_range)//2, 
+            readout = True)
+
+        toggle_mask_button = _create_togglebutton('Show Mask', True)
+        
+        range_slider = _create_slider(
+            slider_min = self.x.min().numpy(), 
+            slider_max = self.x.max().numpy(),  
+            value = [self.x.min().numpy(), self.x.max().numpy()], 
+            continuous_update=False,
+            slider_type = 'FloatRangeSlider' if torch.is_floating_point(self.x) else 'IntRandSlider', 
+            step = 0.01 if torch.is_floating_point(self.x) else 1)
+        
+        image_output = widgets.interactive_output(
+            f = self._plot_slice,
+            controls = {'im_slice': slice_slider, 
+                'with_mask': toggle_mask_button, 
+                'px_range': range_slider})
+                        
+        image_output.layout.height = f'{self.figsize[0]/1.2}in' # suppress flickering
+        image_output.layout.width = f'{self.figsize[1]/1.2}in' # suppress flickering
+        
+        if isinstance(self.y, torch.Tensor): 
+            slice_slider = widgets.HBox([slice_slider, toggle_mask_button])     
+        
+        hist_output = widgets.interactive_output(
+            f = self._plot_hist,
+            controls = {'px_range': range_slider})
+        
+        hist_output.layout.height = f'{self.figsize[0]/1.2}in' # suppress flickering
+        hist_output.layout.width = f'{self.figsize[1]/1.2}in' # suppress flickering
+
+        toggle_mask_button = _create_togglebutton('Show Mask', True)
+        
+        table_output = widgets.interactive_output(
+            f = self._image_summary, 
+            controls = {'px_range': range_slider})
+        
+        table_box = widgets.VBox([table_output], layout=self.vbox_layout)
+        
+        hist_box = widgets.VBox(
+            [hist_output, range_slider],
+            layout=self.vbox_layout)
+        
+        image_box = widgets.VBox(
+            [image_output, slice_slider],
+            layout=self.vbox_layout)
+        
+        self.box = widgets.HBox(
+            [image_box, hist_box, table_box], 
+            layout = widgets.Layout(
+                border = 'solid 1px lightgrey', 
+                margin = '10px 5px 0px 0px', 
+                padding =  '5px', 
+                width = f'{self.figsize[1]*2 + 3}in'))
+        
+        
+class ListViewer(object):
+    """ Display multipple images with their masks or labels/predictions.
+    Arguments: 
+        x (tuple, list): Tensor objects to view
+        y (tuple, list): Tensor objects (in case of segmentation task) or class labels as string. 
+        predictions (str): Class predictions 
+        cmap: colormap for display of `x`
+        max_n: maximum number of items to display
+    """
+    
+    def __init__(self, x:(list, tuple), y=None, prediction:str=None, description: str=None, 
+                 figsize=(4, 4), cmap:str='bone', max_n = 9):
+        self.slice_range = (1, len(x))      
+        x = x[0:max_n]
+        if y: y = y[0:max_n]
+        self.x=x
+        self.y=y
+        self.prediction=prediction
+        self.description=description
+        self.figsize=figsize
+        self.cmap=cmap
+        self.max_n=max_n
+        
+    def _generate_views(self):
+        n_images = len(self.x)
+        image_grid, image_list = [], []
+        
+        for i in range(0, n_images):
+            image = self.x[i]
+            mask = self.y[i] if isinstance(self.y, list) else None
+            pred = self.prediction[i] if self.prediction else None
+            
+            image_list.append(
+                BasicViewer(
+                    x = image, 
+                    y = mask, 
+                    prediction = pred, 
+                    figsize = self.figsize, 
+                    cmap = self.cmap)
+                .image_box)
+            
+            if (i+1) % np.ceil(np.sqrt(n_images)) == 0 or i == n_images - 1: 
+                image_grid.append(widgets.HBox(image_list))
+                image_list = []
+        
+        self.box = widgets.VBox(children=image_grid) 
+        
+    def show(self):
+        self._generate_views()
+        plt.style.use('default')
+        display(self.box)