Structure tensor visualizations

docs/notebooks/structure_tensor.ipynb

@@ -2,17 +2,9 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 20,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING:root:Could not load CuPy: No module named 'cupy'\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import qim3d"
    ]
@@ -31,33 +23,36 @@
     "This notebook shows how to compute eigenvalues and eigenvectors of the **structure tensor** of a 3D volume using the `qim3d` library. The structure tensor (matrix) represents information about the local gradient directions in the volume, such that the eigenvectors represent the orientation of the structure in the volume, and the corresponding eigenvaleus indicate the magnitude.\n",
     "\n",
     "The function `qim3d.processing.structure_tensor` returns two arrays `val` and `vec` for the eigenvalues and eigenvectors, respectively.\\\n",
-    "By having the argument `visulize = True`, the function displays two figures:\n",
+    "By having the argument `visulize = True`, the function displays a figure of three subplots:\n",
     "\n",
     "* Slice of volume with vector field of the eigenvectors\n",
-    "* Orientation histogram of the eigenvectors  "
+    "* Orientation histogram of the eigenvectors\n",
+    "* Slice of volume with overlaying colors of the orientation\n",
+    "\n",
+    "For all three subplots, the colors used to visualize the orientation within the volume are from the HSV colorspace. In these visualizations, the saturation of the color corresponds to the vector component of the slicing direction (i.e. $z$-component when choosing visualization along `axis = 0`). Hence, if an orientation in the volume is orthogonal to the slicing direction, the corresponding color of the visualization will be gray."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### **Example 1**: Structure tensor of bone volume"
+    "### **Example:** Structure tensor of brain tissue volume"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "27aae46cc15f4379a890397802c53fb5",
+       "model_id": "9da567416cb64aaba65c207b09dfb2b3",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "VBox(children=(HBox(children=(IntSlider(value=64, description='Slice index', layout=Layout(width='450px'), max…"
+       "Output()"
       ]
      },
      "metadata": {},
@@ -65,29 +60,61 @@
     }
    ],
    "source": [
-    "# Import 3D volume of bone\n",
-    "bone = qim3d.examples.bone_128x128x128\n",
+    "# Import 3D volume of brain tissue\n",
+    "NT = qim3d.examples.NT_128x128x128\n",
     "\n",
+    "# Visuaize the 3D volume\n",
+    "qim3d.viz.vol(NT)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "From the visualization of the full volume, it can be seen that the circular structures of the brain tissue are aligned orthogonal to the $z$-axis (`axis = 0`). By choosing to slice the volume in this direction, the structure tensor visualizations will largely be gray, since the $z$-component of the eigenvectors are close to $0$, meaning the saturation of the coloring will be close to $0$ (i.e. gray). This can be seen below. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5e927ca8c360420aa086e285459ab3b3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HBox(children=(IntSlider(value=64, description='Slice index', layout=Layout(width='450px'), max…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
     "# Compute eigenvalues and eigenvectors of the structure tensor\n",
-    "val, vec = qim3d.processing.structure_tensor(bone, visualize = True, axis = 1)"
+    "val, vec = qim3d.processing.structure_tensor(NT, visualize = True, axis = 0) # Slicing in z-direction"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### **Example 2:** Structure tensor of brain tissue volume"
+    "By slicing the volume in the $x$-direction (`axis = 2`) instead, the orientation along the length of the structures in the brain tissue can be seen instead. Then the structure tensor visualizations will be largely blue, corresponding to eigenvectors along the $x$-direction with angles of $\\approx \\frac{\\pi}{2}$ radians ($90$ degrees)."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9648fa04a18b4f9f8a5c19749ea57d3b",
+       "model_id": "9d827725d9304f3da5d96cd5e2ae1352",
        "version_major": 2,
        "version_minor": 0
       },
@@ -100,11 +127,8 @@
     }
    ],
    "source": [
-    "# Import 3D volume of brain tissue\n",
-    "NT = qim3d.examples.NT_128x128x128\n",
-    "\n",
     "# Compute eigenvalues and eigenvectors of the structure tensor\n",
-    "val, vec = qim3d.processing.structure_tensor(NT, visualize = True, axis = 2)"
+    "val, vec = qim3d.processing.structure_tensor(NT, visualize = True, axis = 2) # Slicing in x-direction"
    ]
   }
  ],

qim3d/processing/structure_tensor_.py

@@ -9,11 +9,12 @@ def structure_tensor(
     vol: np.ndarray,
     sigma: float = 1.0,
     rho: float = 6.0,
+    base_noise: bool = True,
     full: bool = False,
     visualize = False,
     **viz_kwargs
 ) -> Tuple[np.ndarray, np.ndarray]:
-    """Wrapper for the 3D structure tensor implementation from the [structure_tensor package](https://github.com/Skielex/structure-tensor/)
+    """Wrapper for the 3D structure tensor implementation from the [structure_tensor package](https://github.com/Skielex/structure-tensor/).
 
     
     The structure tensor algorithm is a method for analyzing the orientation of fiber-like structures in 3D images. 
@@ -26,6 +27,7 @@ def structure_tensor(
         vol (np.ndarray): 3D NumPy array representing the volume.
         sigma (float, optional): A noise scale, structures smaller than sigma will be removed by smoothing.
         rho (float, optional): An integration scale giving the size over the neighborhood in which the orientation is to be analysed.
+        base_noise (bool, optional): A flag indicating whether to add a small noise to the volume. Default is True.
         full (bool, optional): A flag indicating that all three eigenvalues should be returned. Default is False.
         visualize (bool, optional): Whether to visualize the structure tensor. Default is False.
         **viz_kwargs: Additional keyword arguments for passed to `qim3d.viz.vectors`. Only used if `visualize=True`.
@@ -44,7 +46,7 @@ def structure_tensor(
         vol = qim3d.examples.NT_128x128x128
         val, vec = qim3d.processing.structure_tensor(vol, visualize = True, axis = 2)
         ```
-        ![structure tensor](assets/screenshots/structure_tensor.gif)
+        ![structure tensor](assets/screenshots/structure_tensor_visualization.gif)
 
 
     !!! info "Runtime and memory usage of the structure tensor method for different volume sizes"
@@ -82,8 +84,20 @@ def structure_tensor(
     if vol.dtype != np.float32 and vol.dtype != np.float64:
         vol = vol.astype(np.float32)
 
+    if base_noise:
+        # Add small noise to the volume
+        # FIXME: This is a temporary solution to avoid uniform regions with constant values 
+        # in the volume, which lead to numerical issues in the structure tensor computation
+        vol_noisy = vol + np.random.default_rng(seed = 0).uniform(0, 1e-10, size=vol.shape)
         
+        # Compute the structure tensor (of volume with noise)
+        s_vol = st.structure_tensor_3d(vol_noisy, sigma, rho)
+
+    else:
+        # Compute the structure tensor (of volume without noise)
         s_vol = st.structure_tensor_3d(vol, sigma, rho)
+
+    # Compute the eigenvalues and eigenvectors of the structure tensor
     val, vec = st.eig_special_3d(s_vol, full=full)
 
     if visualize:

qim3d/viz/structure_tensor.py

 import numpy as np
 from typing import Optional, Union, Tuple
 import matplotlib.pyplot as plt
+import structure_tensor as st 
 from matplotlib.gridspec import GridSpec
 import ipywidgets as widgets
 import logging as log
 
-
 def vectors(
     volume: np.ndarray,
     vec: np.ndarray,
     axis: int = 0,
     slice_idx: Optional[Union[int, float]] = None,
+    grid_size: int = 10,
     interactive: bool = True,
     figsize: Tuple[int, int] = (10, 5),
     show: bool = False,
 ) -> Union[plt.Figure, widgets.interactive]:
     """
-    Displays a grid of eigenvectors from the structure tensor to visualize the orientation of the structures in the volume.
+    Visualizes the orientation of the structures in a 3D volume using the eigenvectors of the structure tensor.
 
     Args:
         volume (np.ndarray): The 3D volume to be sliced.
         vec (np.ndarray): The eigenvectors of the structure tensor.
-        axis (int, optional): The axis along which to visualize the local thickness. Defaults to 0.
+        axis (int, optional): The axis along which to visualize the orientation. Defaults to 0.
         slice_idx (int or float, optional): The initial slice to be visualized. The slice index
             can afterwards be changed. If value is an integer, it will be the index of the slice
             to be visualized. If value is a float between 0 and 1, it will be multiplied by the
@@ -39,29 +40,43 @@ def vectors(
     Returns:
         fig (Union[plt.Figure, widgets.interactive]): If `interactive` is True, returns an interactive widget. Otherwise, returns a matplotlib figure.
 
+    Note:
+        The orientation of the vectors is visualized using an HSV color map, where the saturation corresponds to the vector component 
+        of the slicing direction (i.e. z-component when choosing visualization along `axis = 0`). Hence, if an orientation in the volume
+        is orthogonal to the slicing direction, the corresponding color of the visualization will be gray.
+
     Example:
         ```python
         import qim3d
 
         vol = qim3d.examples.NT_128x128x128
-        val, vec = qim3d.processing.structure_tensor(vol, visualize=True, axis=2)
+        val, vec = qim3d.processing.structure_tensor(vol)
 
         # Visualize the structure tensor
-        qim3d.viz.vectors(vol, vec, axis=2, slice_idx=0.5, interactive=True)
+        qim3d.viz.vectors(vol, vec, axis = 2, interactive = True)
         ```
-        ![structure tensor](assets/screenshots/structure_tensor.gif)
+        ![structure tensor](assets/screenshots/structure_tensor_visualization.gif)
 
     """
 
-    # Define Grid size limits
+    # Ensure volume is a float
+    if volume.dtype != np.float32 and volume.dtype != np.float64:
+        volume = volume.astype(np.float32)
+
+    # Normalize the volume if needed (i.e. if values are in [0, 255])
+    if volume.max() > 1.0:
+        volume = volume / 255.0
+
+    # Define grid size limits
     min_grid_size = max(1, volume.shape[axis] // 50)
     max_grid_size = max(1, volume.shape[axis] // 10)
     if max_grid_size <= min_grid_size:
         max_grid_size = min_grid_size * 5
 
-    # Testing
+    if not grid_size:
         grid_size = (min_grid_size + max_grid_size) // 2
 
+    # Testing
     if grid_size < min_grid_size or grid_size > max_grid_size:
         # Adjust grid size as little as possible to be within the limits
         grid_size = min(max(min_grid_size, grid_size), max_grid_size)
@@ -69,40 +84,62 @@ def vectors(
 
     def _structure_tensor(volume, vec, axis, slice_idx, grid_size, figsize, show):
 
-        # Create subplots
-        fig, ax = plt.subplots(1, 2, figsize=figsize, layout="constrained")
-
         # Choose the appropriate slice based on the specified dimension
         if axis == 0:
             data_slice = volume[slice_idx, :, :]
             vectors_slice_x = vec[0, slice_idx, :, :]
             vectors_slice_y = vec[1, slice_idx, :, :]
+            vectors_slice_z = vec[2, slice_idx, :, :]
+
         elif axis == 1:
             data_slice = volume[:, slice_idx, :]
             vectors_slice_x = vec[0, :, slice_idx, :]
             vectors_slice_y = vec[2, :, slice_idx, :]
+            vectors_slice_z = vec[1, :, slice_idx, :]
+
         elif axis == 2:
             data_slice = volume[:, :, slice_idx]
             vectors_slice_x = vec[1, :, :, slice_idx]
             vectors_slice_y = vec[2, :, :, slice_idx]
+            vectors_slice_z = vec[0, :, :, slice_idx]
+
         else:
             raise ValueError("Invalid dimension. Use 0 for Z, 1 for Y, or 2 for X.")
         
-        ax[0].imshow(data_slice, cmap=plt.cm.gray)
+        # Create three subplots
+        fig, ax = plt.subplots(1, 3, figsize=figsize, layout="constrained")
 
+        blend_hue_saturation = lambda hue, sat : hue * (1 - sat) + 0.5 * sat # Function for blending hue and saturation
+        blend_slice_colors = lambda slice, colors : 0.5 * (slice + colors) # Function for blending image slice with orientation colors
+        
+        # ----- Subplot 1: Image slice with orientation vectors ----- #
         # Create meshgrid with the correct dimensions
         xmesh, ymesh = np.mgrid[0 : data_slice.shape[0], 0 : data_slice.shape[1]]
 
         # Create a slice object for selecting the grid points
         g = slice(grid_size // 2, None, grid_size)
 
+        # Angles from 0 to pi
+        angles_quiver = np.mod(np.arctan2(vectors_slice_y[g, g], vectors_slice_x[g, g]), np.pi) 
+
+        # Calculate z-component (saturation)
+        saturation_quiver = (vectors_slice_z[g, g]**2)[:, :, np.newaxis] 
+
+        # Calculate hue
+        hue_quiver = plt.cm.hsv(angles_quiver / np.pi)
+
+        # Blend hue and saturation
+        rgba_quiver = blend_hue_saturation(hue_quiver, saturation_quiver)
+        rgba_quiver = np.clip(rgba_quiver, 0, 1) # Ensure rgba values are values within [0, 1]
+        rgba_quiver_flat = rgba_quiver.reshape((rgba_quiver.shape[0]*rgba_quiver.shape[1], 4)) # Flatten array for quiver plot
+
         # Plot vectors
         ax[0].quiver(
             ymesh[g, g],
             xmesh[g, g],
             vectors_slice_x[g, g],
             vectors_slice_y[g, g],
-            color="orange",
+            color=rgba_quiver_flat,
             angles="xy",
         )
         ax[0].quiver(
@@ -110,35 +147,64 @@ def vectors(
             xmesh[g, g],
             -vectors_slice_x[g, g],
             -vectors_slice_y[g, g],
-            color="orange",
+            color=rgba_quiver_flat,
             angles="xy",
         )
 
-        # Set title and turn off axis
-        ax[0].set_title(f"Slice {slice_idx}" if not interactive else None)
+        ax[0].imshow(data_slice, cmap = plt.cm.gray)
+        ax[0].set_title(f"Orientation vectors (slice {slice_idx})" if not interactive else "Orientation vectors")
         ax[0].set_axis_off()
 
-        # Orientations histogram
+        # ----- Subplot 2: Orientation histogram ----- #
         nbins = 36
-        angles = np.arctan2(vectors_slice_y, vectors_slice_x)  # angles from 0 to pi
-        distribution, bin_edges = np.histogram(angles, bins=nbins, range=(0, np.pi))
-
-        # Find the bin with the maximum count
-        peak_bin_idx = np.argmax(abs(distribution))
-        # Calculate the center of the peak bin
-        peak_angle_rad = (bin_edges[peak_bin_idx] + bin_edges[peak_bin_idx + 1]) / 2
-        # Convert the peak angle to degrees
-        peak_angle_deg = np.degrees(peak_angle_rad)
-        bin_centers = (np.arange(nbins) + 0.5) * np.pi / nbins  # half circle (180 deg)
-        colors = plt.cm.hsv(bin_centers / np.pi)
-        ax[1].bar(bin_centers, distribution, width=np.pi / nbins, color=colors)
+
+        # Angles from 0 to pi
+        angles = np.mod(np.arctan2(vectors_slice_y, vectors_slice_x), np.pi) 
+
+        # Orientation histogram over angles 
+        distribution, bin_edges = np.histogram(angles, bins=nbins, range=(0.0, np.pi))
+
+        # Half circle (180 deg)
+        bin_centers = (np.arange(nbins) + 0.5) * np.pi / nbins 
+
+        # Calculate z-component (saturation) for each bin
+        bins = np.digitize(angles.ravel(), bin_edges)
+        saturation_bin = np.array([np.mean((vectors_slice_z**2).ravel()[bins == i]) \
+                                    if np.sum(bins == i) > 0 else 0 for i in range(1, len(bin_edges))])
+
+        # Calculate hue for each bin
+        hue_bin = plt.cm.hsv(bin_centers / np.pi)
+
+        # Blend hue and saturation
+        rgba_bin = hue_bin.copy()
+        rgba_bin[:, :3] = blend_hue_saturation(hue_bin[:, :3], saturation_bin[:, np.newaxis]) 
+
+        ax[1].bar(bin_centers, distribution, width=np.pi/nbins, color=rgba_bin)
         ax[1].set_xlabel("Angle [radians]")
         ax[1].set_xlim([0, np.pi])
         ax[1].set_aspect(np.pi / ax[1].get_ylim()[1])
         ax[1].set_xticks([0, np.pi / 2, np.pi])
         ax[1].set_xticklabels(["0", "$\\frac{\\pi}{2}$", "$\\pi$"])
-        ax[1].set_ylabel("Count")
-        ax[1].set_title(f"Histogram over angles (peak at {round(peak_angle_deg)}°)")
+        ax[1].set_yticks([])
+        ax[1].set_ylabel("Frequency")
+        ax[1].set_title(f"Histogram over orientation angles")
+
+        # ----- Subplot 3: Image slice colored according to orientation ----- #
+        # Calculate z-component (saturation)
+        saturation = (vectors_slice_z**2)[:, :, np.newaxis]
+
+        # Calculate hue
+        hue = plt.cm.hsv(angles / np.pi) 
+
+        # Blend hue and saturation
+        rgba = blend_hue_saturation(hue, saturation)
+
+        # Grayscale image slice blended with orientation colors 
+        data_slice_orientation_colored = (blend_slice_colors(plt.cm.gray(data_slice), rgba) * 255).astype('uint8')
+
+        ax[2].imshow(data_slice_orientation_colored)
+        ax[2].set_title(f"Colored orientations (slice {slice_idx})" if not interactive else "Colored orientations")
+        ax[2].set_axis_off()
 
         if show:
             plt.show()
@@ -155,6 +221,7 @@ def vectors(
 
     if slice_idx is None:
         slice_idx = volume.shape[axis] // 2
+
     elif isinstance(slice_idx, float):
         if slice_idx < 0 or slice_idx > 1:
             raise ValueError(
@@ -195,8 +262,11 @@ def vectors(
         sliders_box = widgets.HBox([slide_idx_slider, grid_size_slider])
         widget_obj = widgets.VBox([sliders_box, widget_obj.children[-1]])
         widget_obj.layout.align_items = "center"
+
         if show:
             display(widget_obj)
+
         return widget_obj
+    
     else:
-        return _structure_tensor(volume, vec, axis, slice_idx, figsize, show)
+        return _structure_tensor(volume, vec, axis, slice_idx, grid_size, figsize, show)