Improve visualization and PCA

Author: fepegar

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

@@ -9,31 +9,39 @@ def _pca(
     data: torch.Tensor,
     num_components: int = 6,
     whiten: bool = True,
-    vmin: float = -2.3,
-    vmax: float = 2.3,
+    clip_range: tuple[float, float] | None = (-2.3, 2.3),
+    normalize: bool = True,
+    make_skewness_positive: bool = True,
+    **pca_kwargs,
 ) -> torch.Tensor:
     # Adapted from https://github.com/facebookresearch/capi/blob/main/eval_visualizations.py
-
+    # 2.3 is roughly 2σ for a standard-normal variable, 99% of values map inside [0,1].
     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
+    pca = PCA(n_components=num_components, whiten=whiten, **pca_kwargs)
+    projected: np.ndarray = pca.fit_transform(X).T
+    if normalize:
+        projected /= projected[0].std()
+    if make_skewness_positive:
+        for component in projected:
+            third_cumulant = np.mean(component**3)
+            second_cumulant = np.mean(component**2)
+            skewness = third_cumulant / second_cumulant ** (3 / 2)
+            if skewness < 0:
+                component *= -1
     grid: np.ndarray = rearrange(
-        projected,
+        projected.T,
         '(x y z) c -> c x y z',
         x=size_x,
         y=size_y,
         z=size_z,
     )
+    if clip_range is not None:
+        vmin, vmax = clip_range
+    else:
+        vmin, vmax = grid.min(), grid.max()
     grid = (grid - vmin) / (vmax - vmin)
     return torch.from_numpy(grid.clip(0, 1))

src/torchio/visualization.py

@@ -34,7 +34,7 @@ def import_mpl_plt():
     return mpl, plt
 
 
-def rotate(image, radiological=True, n=-1):
+def rotate(image: np.ndarray, *, radiological: bool = True, n: int = -1) -> np.ndarray:
     # Rotate for visualization purposes
     image = np.rot90(image, n, axes=(0, 1))
     if radiological:
@@ -93,13 +93,14 @@ def plot_volume(
     else:
         data = image.data[np.newaxis, channel]
     data = rearrange(data, 'c x y z -> x y z c')
+    data_numpy: np.ndarray = data.cpu().numpy()
 
     if indices is None:
-        indices = np.array(data.shape[:3]) // 2
+        indices = np.array(data_numpy.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)
+    slice_x = rotate(data_numpy[i, :, :], radiological=radiological)
+    slice_y = rotate(data_numpy[:, j, :], radiological=radiological)
+    slice_z = rotate(data_numpy[:, :, k], radiological=radiological)
 
     if isinstance(cmap, dict):
         slices = slice_x, slice_y, slice_z
@@ -118,8 +119,15 @@ def plot_volume(
     sr, sa, ss = image.spacing
     imshow_kwargs['origin'] = 'lower'
 
-    if percentiles is not None and not is_label:
-        p1, p2 = np.percentile(data, percentiles)
+    if not is_label:
+        displayed_data = np.concatenate(
+            [
+                slice_x.flatten(),
+                slice_y.flatten(),
+                slice_z.flatten(),
+            ]
+        )
+        p1, p2 = np.percentile(displayed_data, percentiles)
         imshow_kwargs['vmin'] = p1
         imshow_kwargs['vmax'] = p2