Add support to compute and visualize PCA

Author: fepegar

pyproject.toml

@@ -51,6 +51,9 @@ dependencies = [
 csv = ["pandas>=1"]
 plot = ["distinctipy>=1.3.4", "matplotlib>=3.4"]
 video = ["ffmpeg-python>=0.2.0"]
+sklearn = [
+    "scikit-learn>=1.6.1",
+]
 
 [project.scripts]
 tiohd = "torchio.cli.print_info:app"

src/torchio/external/imports.py

@@ -36,6 +36,10 @@ def get_ffmpeg() -> ModuleType:
     return ffmpeg
 
 
+def get_sklearn() -> ModuleType:
+    return _check_and_import(module='sklearn', extra='sklearn', package='scikit-learn')
+
+
 def _check_executable(executable: str) -> None:
     if which(executable) is None:
         message = (

src/torchio/transforms/preprocessing/intensity/pca.py

@@ -0,0 +1,39 @@
+import numpy as np
+import torch
+from einops import rearrange
+
+from ....external.imports import get_sklearn
+
+
+def _pca(
+    data: torch.Tensor,
+    num_components: int = 6,
+    whiten: bool = True,
+    vmin: float = -2.3,
+    vmax: float = 2.3,
+) -> torch.Tensor:
+    # Adapted from https://github.com/facebookresearch/capi/blob/main/eval_visualizations.py
+
+    sklearn = get_sklearn()
+    PCA = sklearn.decomposition.PCA
+
+    _, size_x, size_y, size_z = data.shape
+    X = rearrange(data, 'c x y z -> (x y z) c')
+    pca = PCA(n_components=num_components, whiten=whiten)
+    projected: np.ndarray = pca.fit_transform(X)
+    projected /= projected[:, 0].std()
+    for i in range(num_components):
+        numerator = np.mean(np.power(projected[:, i], 3))
+        denominator = np.power(np.mean(np.power(projected[:, i], 2)), 1.5)
+        skew = numerator / denominator
+        if skew < 0:
+            projected[:, i] *= -1
+    grid: np.ndarray = rearrange(
+        projected,
+        '(x y z) c -> c x y z',
+        x=size_x,
+        y=size_y,
+        z=size_z,
+    )
+    grid = (grid - vmin) / (vmax - vmin)
+    return torch.from_numpy(grid.clip(0, 1))

src/torchio/visualization.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import torch
+from einops import rearrange
 
 from .data.image import Image
 from .data.image import LabelMap
@@ -35,7 +36,7 @@ def import_mpl_plt():
 
 def rotate(image, radiological=True, n=-1):
     # Rotate for visualization purposes
-    image = np.rot90(image, n)
+    image = np.rot90(image, n, axes=(0, 1))
     if radiological:
         image = np.fliplr(image)
     return image
@@ -59,17 +60,18 @@ def _create_categorical_colormap(data: torch.Tensor) -> ListedColormap:
 def plot_volume(
     image: Image,
     radiological=True,
-    channel=-1,  # default to foreground for binary maps
+    channel=None,
     axes=None,
     cmap=None,
     output_path=None,
     show=True,
     xlabels=True,
-    percentiles: tuple[float, float] = (0.5, 99.5),
+    percentiles: tuple[float, float] = (0, 100),
     figsize=None,
     title=None,
     reorient=True,
     indices=None,
+    rgb=True,
     **imshow_kwargs,
 ):
     _, plt = import_mpl_plt()
@@ -80,14 +82,25 @@ def plot_volume(
 
     if reorient:
         image = ToCanonical()(image)  # type: ignore[assignment]
-    data = image.data[channel]
+
+    is_label = isinstance(image, LabelMap)
+    if is_label:  # probabilistic label map
+        data = image.data[np.newaxis, -1]
+    elif rgb and image.num_channels == 3:
+        data = image.data  # keep image as it is
+    elif channel is None:
+        data = image.data[0:1]  # just use the first channel
+    else:
+        data = image.data[np.newaxis, channel]
+    data = rearrange(data, 'c x y z -> x y z c')
+
     if indices is None:
-        indices = np.array(data.shape) // 2
+        indices = np.array(data.shape[:3]) // 2
     i, j, k = indices
     slice_x = rotate(data[i, :, :], radiological=radiological)
     slice_y = rotate(data[:, j, :], radiological=radiological)
     slice_z = rotate(data[:, :, k], radiological=radiological)
-    is_label = isinstance(image, LabelMap)
+
     if isinstance(cmap, dict):
         slices = slice_x, slice_y, slice_z
         slice_x, slice_y, slice_z = color_labels(slices, cmap)
@@ -98,6 +111,9 @@ def plot_volume(
 
     if is_label:
         imshow_kwargs['interpolation'] = 'none'
+    else:
+        if 'interpolation' not in imshow_kwargs:
+            imshow_kwargs['interpolation'] = 'bicubic'
 
     sr, sa, ss = image.spacing
     imshow_kwargs['origin'] = 'lower'
@@ -108,23 +124,35 @@ def plot_volume(
         imshow_kwargs['vmax'] = p2
 
     sag_aspect = ss / sa
-    sag_axis.imshow(slice_x, aspect=sag_aspect, **imshow_kwargs)
+    sag_axis.imshow(
+        slice_x,
+        aspect=sag_aspect,
+        **imshow_kwargs,
+    )
     if xlabels:
         sag_axis.set_xlabel('A')
     sag_axis.set_ylabel('S')
     sag_axis.invert_xaxis()
     sag_axis.set_title('Sagittal')
 
     cor_aspect = ss / sr
-    cor_axis.imshow(slice_y, aspect=cor_aspect, **imshow_kwargs)
+    cor_axis.imshow(
+        slice_y,
+        aspect=cor_aspect,
+        **imshow_kwargs,
+    )
     if xlabels:
         cor_axis.set_xlabel('R')
     cor_axis.set_ylabel('S')
     cor_axis.invert_xaxis()
     cor_axis.set_title('Coronal')
 
     axi_aspect = sa / sr
-    axi_axis.imshow(slice_z, aspect=axi_aspect, **imshow_kwargs)
+    axi_axis.imshow(
+        slice_z,
+        aspect=axi_aspect,
+        **imshow_kwargs,
+    )
     if xlabels:
         axi_axis.set_xlabel('R')
     axi_axis.set_ylabel('A')
@@ -223,15 +251,15 @@ def plot_histogram(x: np.ndarray, show=True, **kwargs) -> None:
 
 def color_labels(arrays, cmap_dict):
     results = []
-    for array in arrays:
-        si, sj = array.shape
+    for slice_array in arrays:
+        si, sj, _ = slice_array.shape
         rgb = np.zeros((si, sj, 3), dtype=np.uint8)
         for label, color in cmap_dict.items():
             if isinstance(color, str):
                 mpl, _ = import_mpl_plt()
                 color = mpl.colors.to_rgb(color)
                 color = [255 * n for n in color]
-            rgb[array == label] = color
+            rgb[slice_array[..., 0] == label] = color
         results.append(rgb)
     return results