Use better colormap for label maps

Author: fepegar

pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
 
 [project.optional-dependencies]
 csv = ["pandas>=1"]
-plot = ["matplotlib>=3.4"]
+plot = ["distinctipy>=1.3.4", "matplotlib>=3.4"]
 
 [project.scripts]
 tiohd = "torchio.cli.print_info:app"

src/torchio/external/imports.py

@@ -20,3 +20,7 @@ def _check_and_import(package: str, extra: str) -> ModuleType:
 
 def get_pandas() -> ModuleType:
     return _check_and_import(package='pandas', extra='csv')
+
+
+def get_distinctipy() -> ModuleType:
+    return _check_and_import(package='distinctipy', extra='plot')

src/torchio/visualization.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import warnings
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -10,6 +13,9 @@
 from .transforms.preprocessing.spatial.to_canonical import ToCanonical
 from .types import TypePath
 
+if TYPE_CHECKING:
+    from matplotlib.colors import ListedColormap
+
 
 def import_mpl_plt():
     try:
@@ -28,6 +34,17 @@ def rotate(image, radiological=True, n=-1):
     return image
 
 
+def _create_categorical_colormap(data: torch.Tensor) -> ListedColormap:
+    from .external.imports import get_distinctipy
+
+    mpl, _ = import_mpl_plt()
+    distinctipy = get_distinctipy()
+    num_classes = int(data.max())
+    distinct_colors = distinctipy.get_colors(num_classes, pastel_factor=0.5, rng=0)
+    colors = [(0, 0, 0), *distinct_colors]  # prepend black
+    return mpl.colors.ListedColormap(colors)
+
+
 def plot_volume(
     image: Image,
     radiological=True,
@@ -65,7 +82,7 @@ def plot_volume(
         slice_x, slice_y, slice_z = color_labels(slices, cmap)
     else:
         if cmap is None:
-            cmap = 'cubehelix' if is_label else 'gray'
+            cmap = _create_categorical_colormap(data) if is_label else 'gray'
         imshow_kwargs['cmap'] = cmap
 
     if is_label: