Merge pull request #631 from ANTsX/fix-plot-overlay

FIX: overlay plot background is masked
ANTsX · May 15, 2024 · 9266f2f · 9266f2f
2 parents 3dcfdeb + f1a3515
commit 9266f2f
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 143 deletions.
diff --git a/ants/viz/__init__.py b/ants/viz/__init__.py
@@ -1,4 +1,3 @@
-from .create_tiled_mosaic import create_tiled_mosaic
 
 from .plot import plot
 from .movie import movie

diff --git a/ants/viz/create_tiled_mosaic.py b/ants/viz/create_tiled_mosaic.py
diff --git a/ants/viz/plot.py b/ants/viz/plot.py
@@ -328,9 +328,10 @@ def reorient_slice(x, axis):
 
             if overlay is not None:
                 ov_arr = overlay.numpy()
+                mask = ov_arr == 0
+                mask = np.ma.masked_where(mask == 0, mask)
+                ov_arr = np.ma.masked_array(ov_arr, mask)
                 ov_arr = rotate90_matrix(ov_arr)
-                if ov_arr.dtype not in ["uint8", "uint32"]:
-                    ov_arr = np.ma.masked_where(ov_arr == 0, ov_arr)
 
             fig = plt.figure()
             if title is not None:
@@ -375,8 +376,9 @@ def reorient_slice(x, axis):
                 if reorient:
                     overlay = overlay.reorient_image2("LAI")
                 ov_arr = overlay.numpy()
-                if ov_arr.dtype not in ["uint8", "uint32"]:
-                    ov_arr = np.ma.masked_where(ov_arr == 0, ov_arr)
+                mask = ov_arr == 0
+                mask = np.ma.masked_where(mask == 0, mask)
+                ov_arr = np.ma.masked_array(ov_arr, mask)
                 ov_arr = np.rollaxis(ov_arr, axis)
 
             if slices is None: