Add colormap

src/open_dive/scripts/run.py

@@ -35,6 +35,11 @@ def main():
         nargs="+",
         help='Slice index (integer/tuple of three integers) or "m" for middle slice. Default is "m".',
     )
+    scalar_group.add_argument(
+        "--cmap",
+        default="gray",
+        help='Matplotlib or cmcrameri colormap to use for image. Default is "gray".',
+    )
     scalar_group.add_argument(
         "-o",
         "--orientation",
@@ -193,6 +198,13 @@ def main():
         default=None,
         help="Elevation angle of the view.",
     )
+    window_group.add_argument(
+        "--background_color",
+        type=float,
+        nargs=3,
+        default=(0.0, 0.0, 0.0),
+        help="Background color of the scene in RGB format (three floats between 0 and 1). Default is black (0.0, 0.0, 0.0).",
+    )
 
     args = parser.parse_args()
 
@@ -211,6 +223,7 @@ def main():
     plot_nifti(
         nifti_path=args.nifti_path,
         data_slice=args.slice,
+        nifti_cmap=args.cmap,
         orientation=args.orientation,
         size=args.size,
         zoom=args.zoom,
@@ -239,4 +252,5 @@ def main():
         azimuth=args.azimuth,
         elevation=args.elevation,
         glass_brain_path=args.glass_brain,
+        background_color=args.background_color,
     )

src/open_dive/viz.py

@@ -21,7 +21,8 @@
     slicer,
 )
 from fury.lib import Actor
-from fury.utils import apply_affine
+from fury.utils import apply_affine, apply_affine_to_actor
+from fury.transform import transform_from_matrix
 from fury.colormap import line_colors, orient2rgb, boys2rgb
 from matplotlib.colors import Colormap
 from scipy.ndimage import binary_dilation, gaussian_filter
@@ -32,6 +33,7 @@ def plot_nifti(
     data_slice: str | tuple[int, int, int] | int = "m",
     orientation: str = "axial",
     size: tuple[int, int] = (600, 400),
+    nifti_cmap: str | None = "gray",
     zoom: float = 1.0,
     azimuth: float | None = None,
     elevation: float | None = None,
@@ -57,6 +59,7 @@ def plot_nifti(
     sh_basis: str = "descoteaux07",
     scale: int = 1,
     glass_brain_path: os.PathLike | None = None,
+    background_color: tuple[float, float, float] = (0.0, 0.0, 0.0),
     **kwargs,
 ) -> None:
     """Create a 2D rendering of a NIFTI slice.
@@ -69,6 +72,8 @@ def plot_nifti(
         Slice to plot or "m" for middle slice
     orientation : str, default "axial"
         Can be "axial", "sagittal" or "coronal"
+    nifti_cmap : str, default "gray"
+        Colormap to use for the NIFTI image
     size : tuple, default (600, 400)
         Size of window
     zoom : float, default 1.0
@@ -121,6 +126,8 @@ def plot_nifti(
         Scale of the tensor glyphs or ODF glyphs
     glass_brain_path : os.PathLike, optional
         Optional glass brain mask to overlay
+    background_color : tuple of float, default (0.0, 0.0, 0.0)
+        Background color of the scene, in RGB format from 0 to 1
 
     **kwargs
         Additional keyword arguments to pass to fury.actor.slicer
@@ -175,7 +182,7 @@ def plot_nifti(
         scene_bound_nifti = nib.load(scene_bound_nifti_path)
         # scene_bound_data, scene_bound_affine = load_nifti(scene_bound_nifti_path)
         scene_bound_nifti = nib.as_closest_canonical(scene_bound_nifti)
-        scene_bound_affine = scene_bound_nifti.affine
+        scene_bound_affine = np.eye(4)  # scene_bound_nifti.affine
         scene_bound_data = scene_bound_nifti.get_fdata()
         scene_bound_data_shape = scene_bound_data.shape
 
@@ -211,7 +218,7 @@ def plot_nifti(
             volume_idx=volume_idx,
             value_range=value_range,
             opacity=opacity,
-            cmap="gray",
+            cmap=nifti_cmap,
             **kwargs,
         )
         scene.add(slice_actor)
@@ -237,6 +244,7 @@ def plot_nifti(
             colorbar_position=(0.8, 0.1),
             colorbar_height=0.5,
             colorbar_width=0.1,
+            cmap=plt.get_cmap(nifti_cmap),
         )
         scene.add(scalar_bar)
 
@@ -276,11 +284,17 @@ def plot_nifti(
             colors = tractography_cmap
 
         # Add each tractography with its corresponding color
+        affine = (
+            np.linalg.inv(scene_bound_affine)
+            if scene_bound_affine is not None
+            else None
+        )
         stream_actors = _create_tractography_actor(
             tractography_path,
             colors=colors,
             tractography_opacity=tractography_opacity,
             tractography_color_by_endpoints=tractography_color_by_endpoints,
+            affine=affine,
         )
         for stream_actor in stream_actors:
             scene.add(stream_actor)
@@ -317,6 +331,9 @@ def plot_nifti(
             scene_bound_data_shape = glass_brain_data.shape
             scene_bound_affine = glass_brain_affine
 
+    # Set background color
+    scene.background(background_color)
+
     _set_camera(
         scene=scene,
         focus=focus,
@@ -405,6 +422,7 @@ def _create_nifti_actor(
     # nifti = nib.load(nifti_path)
     # nifti = nib.as_closest_canonical(nifti)
     nifti = nib.load(nifti_path)
+    nifti = nib.as_closest_canonical(nifti)
 
     if len(nifti.shape) == 4:
         if volume_idx is None:
@@ -425,7 +443,7 @@ def _create_nifti_actor(
         data = nifti.get_fdata()
 
     # Get the data and affine
-    affine = nifti.affine
+    affine = np.eye(4)  # nifti.affine
 
     # value range
     if cmap == "slant":
@@ -506,6 +524,9 @@ def _create_colorbar_actor(
     colorbar.SetPosition(*colorbar_position)  # Position of the colorbar
     colorbar.SetHeight(colorbar_height)  # Adjust height (increase size)
     colorbar.SetWidth(colorbar_width)  # Adjust width (increase size)
+    label_prop = colorbar.GetLabelTextProperty()
+    label_prop.ItalicOff()
+    label_prop.SetFontFamilyToArial()
 
     if not labels:
         colorbar.SetLabelFormat("")
@@ -518,6 +539,7 @@ def _create_tractography_actor(
     colors: list[tuple[float, float, float]] | str,
     tractography_opacity: list[float] = [0.6],
     tractography_color_by_endpoints: bool = False,
+    affine: np.ndarray | None = None,
 ) -> list[Actor]:
     """Create tractography actors from a list of files."""
 
@@ -563,17 +585,35 @@ def _create_tractography_actor(
                 linewidth=0.2,
                 opacity=tractography_opacity[i],
             )
+            stream_actor = (
+                apply_affine_to_actor(stream_actor, affine)
+                if affine is not None
+                else stream_actor
+            )
             stream_actors.append(stream_actor)
     else:
-        for i, (tract_file, color) in enumerate(zip(tractography_path, colors)):
+        # If colors is a string, we need to sample the colormap
+        if isinstance(colors, str):
+            cmap = plt.get_cmap(colors)
+            colors = [
+                cmap(i / (len(tractography_path)))
+                for i in range(len(tractography_path))
+            ]
+        for i, (tract_file) in enumerate(tractography_path):
             streamlines_nifti = nib.streamlines.load(tract_file)
             streamlines = streamlines_nifti.streamlines
+            color = colors[i] if isinstance(colors, list) else colors
             stream_actor = actor.line(
                 streamlines,
                 colors=color,
                 linewidth=0.2,
                 opacity=tractography_opacity[i],
             )
+            stream_actor = (
+                apply_affine_to_actor(stream_actor, affine)
+                if affine is not None
+                else stream_actor
+            )
             stream_actors.append(stream_actor)
     return stream_actors
 
@@ -591,7 +631,7 @@ def _create_tensor_actor(
     tensor_nifti = nib.load(tensor_path)
     tensor_nifti = nib.as_closest_canonical(tensor_nifti)
     tensor_data = tensor_nifti.get_fdata()
-    tensor_affine = tensor_nifti.affine
+    tensor_affine = np.eye(4)  # tensor_nifti.affine
     tensor_matrix = from_lower_triangular(tensor_data)
     eigvals, eigvecs = decompose_tensor(tensor_matrix)
     mask = np.ones(tensor_data.shape[:3])
@@ -652,7 +692,7 @@ def _create_odf_actor(
     odf_nifti = nib.load(odf_path)
     odf_nifti = nib.as_closest_canonical(odf_nifti)
     odf_data = odf_nifti.get_fdata()
-    odf_affine = odf_nifti.affine
+    odf_affine = np.eye(4)  # odf_nifti.affine
     sphere = get_sphere(name="repulsion724")  # Use a precomputed sphere
     sh_order_max = calculate_max_order(odf_data.shape[-1])
     B, _ = sh_to_sf_matrix(
@@ -663,7 +703,7 @@ def _create_odf_actor(
         sphere=sphere,
         B_matrix=B,
         scale=scale,
-        norm=False,
+        norm=None,
         affine=odf_affine,
     )