Colormap for qim3d.viz.vol

See https://www.notion.so/qim-dtu/Colormap-for-qim3d-viz-vol-123bab83b2f980369353d0e473c961a0

qim3d/viz/k3d.py

@@ -8,6 +8,8 @@ Volumetric visualization using K3D
 """
 
 import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import Colormap
 from qim3d.utils.logger import log
 from qim3d.utils.misc import downscale_img, scale_to_float16
 
@@ -19,6 +21,7 @@ def vol(
     save=False,
     grid_visible=False,
     cmap=None,
+    objects=False,
     vmin=None,
     vmax=None,
     samples="auto",
@@ -40,7 +43,8 @@ def vol(
             If a string is provided, it's interpreted as the file path where the HTML
             file will be saved. Defaults to False.
         grid_visible (bool, optional): If True, the grid is visible in the plot. Defaults to False.
-        cmap (list, optional): The color map to be used for the volume rendering. Defaults to None.
+        cmap (str or matplotlib.colors.Colormap or list, optional): The color map to be used for the volume rendering. If a string is passed, it should be a matplotlib colormap name. Defaults to None.
+        objects (bool): Set to True if doing an object label visualization with a corresponding cmap; otherwise, the plot may appear poorly. Defaults to False.
         vmin (float, optional): Together with vmax defines the data range the colormap covers. By default colormap covers the full range. Defaults to None.
         vmax (float, optional): Together with vmin defines the data range the colormap covers. By default colormap covers the full range. Defaults to None
         samples (int, optional): The number of samples to be used for the volume rendering in k3d. Defaults to 512.
@@ -55,6 +59,9 @@ def vol(
     Raises:
         ValueError: If `aspectmode` is not `'data'` or `'cube'`.
 
+    Tip:
+        The function can be used for object label visualization using a `cmap` created with `qim3d.viz.colormaps.objects` along with setting `objects=True`. The latter ensures appropriate rendering.
+    
     Example:
         Display a volume inline:
 
@@ -122,6 +129,27 @@ def vol(
     if vmax:
         color_range[1] = vmax
 
+    # Handle the different formats that cmap can take
+    if cmap:
+        if isinstance(cmap, str):
+            cmap = plt.get_cmap(cmap) # Convert to Colormap object
+        if isinstance(cmap, Colormap):
+            # Convert to the format of cmap required by k3d.volume
+            attr_vals = np.linspace(0.0, 1.0, num=cmap.N)
+            RGB_vals = cmap(np.arange(0, cmap.N))[:, :3]
+            cmap = np.column_stack((attr_vals, RGB_vals)).tolist()
+    
+    # Default k3d.volume settings
+    opacity_function = []
+    interpolation = True
+    if objects:
+        # without these settings, the plot will look bad when cmap is created with qim3d.viz.colormaps.objects
+        opacity_function = [
+            (0.0, 0.95),
+            (1.0, 0.95)
+        ]
+        interpolation = False
+    
     # Create the volume plot
     plt_volume = k3d.volume(
         img,
@@ -133,6 +161,8 @@ def vol(
         color_map=cmap,
         samples=samples,
         color_range=color_range,
+        opacity_function=opacity_function,
+        interpolation=interpolation
     )
     plot = k3d.plot(grid_visible=grid_visible, **kwargs)
     plot += plt_volume