Include visualization features.

src/visualization/__init__.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug  7 09:53:45 2024
+
+@author: jsteiglechner
+"""
+

src/visualization/color_definitions.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug  7 09:54:38 2024
+
+@author: jsteiglechner
+"""
+
+from matplotlib.colors import ListedColormap
+
+
+# https://jfly.uni-koeln.de/color/
+colorblind_friendly_colors = [
+    (0, 0, 0),  # black
+    (230/255, 159/255, 0/255),  # orange
+    (86/255, 180/255, 233/255),  # sky blue
+    (0/255, 158/255, 115/255),  # blueish green
+    (240/255, 228/255, 66/255),  # yellow
+    (0/255, 114/255, 178/255),  # blue
+    (213/255, 94/255, 0/255),  # vermillion
+    (204/255, 121/255, 167/255),  # reddish purple
+]
+colorblind_friendly_cmap = ListedColormap(
+    colors=colorblind_friendly_colors
+).with_extremes(over='0.25', under='0.75')
+colorblind_friendly_cmap_without_bg = ListedColormap(
+    colors=colorblind_friendly_colors[1:]
+).with_extremes(over='0.25', under='0.75')

src/visualization/visualize_mri_data.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+This module visualizes mask predictions of images.
+
+Created on Wed Aug  7 09:54:15 2024
+
+@author: jsteiglechner
+"""
+
+from pathlib import Path
+from typing import Union
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import Normalize, Colormap
+import nibabel as nib
+import numpy as np
+
+from src.visualization.color_definitions import colorblind_friendly_cmap
+
+
+def visualize_image_with_mask(
+    image: np.ndarray,
+    mask: np.ndarray,
+    affine: np.ndarray,
+    color_map: Colormap = colorblind_friendly_cmap,
+    target_path: Union[Path, str] = None,
+    show_plot: bool = False,
+) -> None:
+    """
+    Plot image with mask as overlay.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        DESCRIPTION.
+    mask : np.ndarray
+        DESCRIPTION.
+    affine : np.ndarray
+        DESCRIPTION.
+    target_path : Union[Path, str], optional
+        Path where to save the plot. The default is None.
+    show_plot: bool, optional
+        Flag if plot should be shown. The default is False.
+
+    Returns
+    -------
+    None
+        Plot is shown and saved optionally.
+
+    """
+    t1_min = np.percentile(image, 1)
+    t1_max = np.percentile(image, 99)
+    color_scaler = Normalize(vmin=t1_min, vmax=t1_max)
+    params = {"interpolation": "nearest"}
+
+    ref_pos_x = image.shape[0] // 2
+    ref_pos_y = image.shape[1] // 2
+    ref_pos_z = image.shape[2] // 2
+
+    min_edge_length = np.min(image.shape)
+    # consider max brain size of 24 cm (is a requirement)
+    # 1 pixel is 1mm (transformed in preprocessing)
+    zoom = max(
+        0, min_edge_length -
+        int(240 / np.min(np.linalg.norm(affine, axis=0)[:3]))
+    )
+    # adjust size lenghts, so that all lengths are the same
+    zoom_x = (image.shape[0] - min_edge_length) // 2 + zoom
+    zoom_y = (image.shape[1] - min_edge_length) // 2 + zoom
+    zoom_z = (image.shape[2] - min_edge_length) // 2 + zoom
+
+    fig = plt.figure(dpi=300)
+    plt.subplots_adjust(0, 0, 1, 1, 0, 0)
+    axes = fig.add_subplot(133)
+    axes.imshow(
+        np.fliplr(np.rot90(image[ref_pos_x, zoom_y:-(zoom_y+1),
+                                 zoom_z:-(zoom_z+1)])),
+        cmap="gray",
+        norm=color_scaler,
+    )
+    axes.imshow(
+        np.fliplr(np.rot90(mask[ref_pos_x, zoom_y:-(zoom_y+1),
+                                zoom_z:-(zoom_z+1)])),
+        color_map,
+        alpha=0.4
+        * np.fliplr(np.rot90((mask[ref_pos_x, zoom_y:-(zoom_y+1),
+                                   zoom_z:-(zoom_z+1)] != 0) * 1.)),
+        **params
+    )
+    axes.set_aspect(affine[2, 2] / affine[1, 1])
+    plt.xticks([])
+    plt.yticks([])
+
+    axes = fig.add_subplot(132)
+    axes.imshow(
+        np.fliplr(np.rot90(image[zoom_x:-(zoom_x+1), ref_pos_y,
+                                 zoom_z:-(zoom_z+1)])),
+        cmap="gray",
+        norm=color_scaler,
+    )
+    axes.imshow(
+        np.fliplr(np.rot90(mask[zoom_x:-(zoom_x+1), ref_pos_y,
+                                zoom_z:-(zoom_z+1)])),
+        color_map,
+        alpha=0.4
+        * np.fliplr(np.rot90((mask[zoom_x:-(zoom_x+1), ref_pos_y,
+                                   zoom_z:-(zoom_z+1)] != 0) * 1.)),
+        **params
+    )
+    axes.set_aspect(affine[2, 2] / affine[0, 0])
+    plt.xticks([])
+    plt.yticks([])
+
+    axes = fig.add_subplot(131)
+    axes.imshow(
+        np.fliplr(np.rot90(image[zoom_x:-(zoom_x+1), zoom_y:-(zoom_y+1),
+                                 ref_pos_z])),
+        cmap="gray",
+        norm=color_scaler,
+    )
+    axes.imshow(
+        np.fliplr(np.rot90(mask[zoom_x:-(zoom_x+1), zoom_y:-(zoom_y+1),
+                                ref_pos_z])),
+        color_map,
+        alpha=0.4
+        * np.fliplr(np.rot90((mask[zoom_x:-(zoom_x+1), zoom_y:-(zoom_y+1),
+                                   ref_pos_z] != 0) * 1.)),
+        **params
+    )
+    axes.set_aspect(affine[1, 1] / affine[0, 0])
+    plt.xticks([])
+    plt.yticks([])
+
+    if target_path:
+        plt.savefig(
+            target_path,
+            bbox_inches="tight",
+            pad_inches=0,
+            facecolor="k",
+        )
+    if show_plot:
+        plt.show()
+        return None
+    else:
+        return fig
+
+
+def visualize_prediction_from_paths(
+        image_afp, seg_afp, color_map=colorblind_friendly_cmap):
+    """Combine visualization of prediction with loading."""
+    image = nib.load(image_afp)
+    image_data = image.dataobj[:].astype(float)
+    image_affine = image.affine
+
+    mask = nib.load(seg_afp)
+    mask_data = mask.dataobj[:].astype(int)
+    mask_affine = mask.affine
+
+    if not np.allclose(image_affine, mask_affine):
+        raise ValueError(
+            "Affines do not match. Difference: "
+            f"{np.max(np.abs(image_affine - mask_affine))}")
+
+    fig = visualize_image_with_mask(
+        image_data, mask_data, image_affine,
+        color_map=color_map, show_plot=False
+    )
+
+    return fig