Update normal calculation for labels isosurface rendering (napari#7431)

# References and relevant issues
From discussion on image.sc forum:

https://forum.image.sc/t/napari-labels-layers-borders-displays-with-weird-shading/105499/1

# Description
This modifies the 3D isosurface shader for the Labels layer in a few
ways. The goal is to improve appearance of the 3D rendered labels.
Specifically it should be improved at the boundaries of the volume (for
both `fast` and `smooth` gradient modes) and for 1px thin labels (for
`smooth` gradient mode).

Changes include:
* A new uniform to toggle what happens when sampling _outside_ the
volume:
  * clamping to the edge (default, current behavior)
  * returning the background value
* Update code to prevent dark surfaces in very thin (1px) labels
* Add check for adjacent background pixels to the `smooth` gradient
function
* Update the gradient calculation to include an estimate of the local
gradient instead of setting a default lambert term

The configurable clamping behavior has the effect of giving labels
"dark" (almost hollow-looking) or "filled" faces at the borders of the
label volume. Currently this is controlled by a property on the `Volume`
visual, but toggling is not exposed in the UI. I think at least @jni is
requesting to be able to toggle this in the UI. In the animation below
this is what the floating magicgui widget is toggling. I'm happy to wire
that up in this PR if there's agreement, but I could use input on a
name. We also need to choose a default.


![labels-edges-pr](https://github.com/user-attachments/assets/728505be-d212-417b-a29e-7228761ffed3)

I did a quick benchmark test on my machine and did not see significant
performance hit. I will try to do a longer test as well.

Finally, I need to look a bit into what the "background value" is. Right
now it's hard-coded in the shader, but it sounds like perhaps the
`Labels` layer may be able to have a non-zero background. I want to make
sure to support that or at least not cause a regression.

napari/_vispy/visuals/volume.py

@@ -4,7 +4,8 @@
 from napari.layers.labels._labels_constants import IsoCategoricalGradientMode
 
 FUNCTION_DEFINITIONS = """
-uniform bool u_iso_gradient;
+// switch for clamping values at volume limits
+uniform bool u_clamp_at_border;
 
 // the tolerance for testing equality of floats with floatEqual and floatNotEqual
 const float equality_tolerance = 1e-8;
@@ -25,7 +26,6 @@
     return equal;
 }
 
-
 // the background value for the iso_categorical shader
 const float categorical_bg_value = 0;
 
@@ -36,75 +36,9 @@
     adjacent_bg = adjacent_bg * int( floatEqual(val_pos, categorical_bg_value) );
     return adjacent_bg;
 }
+"""
 
-vec3 calculateGradient(vec3 loc, vec3 step, out int n_bg_borders) {
-    // calculate gradient within the volume by finite differences
-
-    n_bg_borders = 0;
-    vec3 G = vec3(0.0);
-
-    float prev;
-    float next;
-
-    prev = colorToVal($get_data(loc - vec3(step.x, 0, 0)));
-    next = colorToVal($get_data(loc + vec3(step.x, 0, 0)));
-    n_bg_borders += detectAdjacentBackground(prev, next);
-    G.x = next - prev;
-
-    prev = colorToVal($get_data(loc - vec3(0, step.y, 0)));
-    next = colorToVal($get_data(loc + vec3(0, step.y, 0)));
-    n_bg_borders += detectAdjacentBackground(prev, next);
-    G.y = next - prev;
-
-    prev = colorToVal($get_data(loc - vec3(0, 0, step.z)));
-    next = colorToVal($get_data(loc + vec3(0, 0, step.z)));
-    n_bg_borders += detectAdjacentBackground(prev, next);
-    G.z = next - prev;
-
-    return G;
-}
-
-vec3 calculateIsotropicGradient(vec3 loc, vec3 step) {
-    // calculate gradient within the volume by finite differences
-    // using a 3D sobel-feldman convolution kernel
-
-    // the kernel here is a 3x3 cube, centered on the sample at `loc`
-    // the kernel for G.z looks like this:
-
-    // [ +1 +2 +1 ]
-    // [ +2 +4 +2 ]    <-- "loc - step.z" is in the center
-    // [ +1 +2 +1 ]
-
-    // [  0  0  0 ]
-    // [  0  0  0 ]    <-- "loc" is in the center
-    // [  0  0  0 ]
-
-    // [ -1 -2 -1 ]
-    // [ -2 -4 -2 ]    <-- "loc + step.z" is in the center
-    // [ -1 -2 -1 ]
-
-    // kernels for G.x and G.y similar, but transposed
-    // see https://en.wikipedia.org/wiki/Sobel_operator#Extension_to_other_dimensions
-
-    vec3 G = vec3(0.0);
-
-    for (int i=-1; i <= 1; i++) {
-        for (int j=-1; j <= 1; j++) {
-            for (int k=-1; k <= 1; k++) {
-                float val = colorToVal($get_data(loc + vec3(i, j, k) * step));
-                G.x += val * -float(i) *
-                    (1 + float(j == 0 || k == 0) + 2 * float(j == 0 && k == 0));
-                G.y += val * -float(j) *
-                    (1 + float(i == 0 || k == 0) + 2 * float(i == 0 && k == 0));
-                G.z += val * -float(k) *
-                    (1 + float(i == 0 || j == 0) + 2 * float(i == 0 && j == 0));
-            }
-        }
-    }
-
-    return G;
-}
-
+CALCULATE_COLOR_DEFINITION = """
 vec4 calculateShadedCategoricalColor(vec4 betterColor, vec3 loc, vec3 step)
 {
     // Calculate color by incorporating ambient and diffuse lighting
@@ -114,18 +48,13 @@
     float val0 = colorToVal(color0);
     float val1 = 0;
     float val2 = 0;
-    int n_bg_borders = 0;
 
     // View direction
     vec3 V = normalize(view_ray);
 
     // Calculate normal vector from gradient
     vec3 N;
-    if (u_iso_gradient) {
-        N = calculateIsotropicGradient(loc, step);
-    } else {
-        N = calculateGradient(loc, step, n_bg_borders);
-    }
+    N = calculateGradient(loc, step, val0);
 
     // Normalize and flip normal so it points towards viewer
     N = normalize(N);
@@ -148,11 +77,6 @@
 
         // Calculate lighting properties
         float lambertTerm = clamp( dot(N,L), 0.0, 1.0 );
-        if (n_bg_borders > 0) {
-            // to fix dim pixels due to poor normal estimation,
-            // we give a default lambda to pixels surrounded by background
-            lambertTerm = 0.5;
-        }
 
         // Calculate mask
         float mask1 = lightEnabled;
@@ -171,6 +95,126 @@
 }
 """
 
+FAST_GRADIENT_DEFINITION = """
+vec3 calculateGradient(vec3 loc, vec3 step, float current_val) {
+    // calculate gradient within the volume by finite differences
+
+    vec3 G = vec3(0.0);
+
+    float prev;
+    float next;
+    int in_bounds;
+
+    for (int i=0; i<3; i++) {
+        vec3 ax_step = vec3(0.0);
+        ax_step[i] = step[i];
+
+        vec3 prev_loc = loc - ax_step;
+        if (u_clamp_at_border || (prev_loc[i] >= 0.0 && prev_loc[i] <= 1.0)) {
+            prev = colorToVal($get_data(prev_loc));
+        } else {
+            prev = categorical_bg_value;
+        }
+
+        vec3 next_loc = loc + ax_step;
+        if (u_clamp_at_border || (next_loc[i] >= 0.0 && next_loc[i] <= 1.0)) {
+            next = colorToVal($get_data(next_loc));
+        } else {
+            next = categorical_bg_value;
+        }
+
+        // add to the gradient where the adjacent voxels are both background
+        // to fix dim pixels due to poor normal estimation
+        G[i] = next - prev + (next - current_val) * 2.0 * detectAdjacentBackground(prev, next);
+    }
+
+    return G;
+}
+"""
+
+SMOOTH_GRADIENT_DEFINITION = """
+vec3 calculateGradient(vec3 loc, vec3 step, float current_val) {
+    // calculate gradient within the volume by finite differences
+    // using a 3D sobel-feldman convolution kernel
+
+    // the kernel here is a 3x3 cube, centered on the sample at `loc`
+    // the kernel for G.z looks like this:
+
+    // [ +1 +2 +1 ]
+    // [ +2 +4 +2 ]    <-- "loc - step.z" is in the center
+    // [ +1 +2 +1 ]
+
+    // [  0  0  0 ]
+    // [  0  0  0 ]    <-- "loc" is in the center
+    // [  0  0  0 ]
+
+    // [ -1 -2 -1 ]
+    // [ -2 -4 -2 ]    <-- "loc + step.z" is in the center
+    // [ -1 -2 -1 ]
+
+    // kernels for G.x and G.y similar, but transposed
+    // see https://en.wikipedia.org/wiki/Sobel_operator#Extension_to_other_dimensions
+
+    vec3 G = vec3(0.0);
+    // next and prev are the directly adjacent values along x, y, and z
+    vec3 next = vec3(0.0);
+    vec3 prev = vec3(0.0);
+
+    float val;
+    bool is_on_border = false;
+    for (int i=-1; i <= 1; i++) {
+        for (int j=-1; j <= 1; j++) {
+            for (int k=-1; k <= 1; k++) {
+                if (is_on_border && (i != 0 && j != 0 && k != 0)) {
+                    // we only care about on-axis values if we are on a border
+                    continue;
+                }
+                vec3 sample_loc = loc + vec3(i, j, k) * step;
+                bool is_in_bounds = all(greaterThanEqual(sample_loc, vec3(0.0)))
+                    && all(lessThanEqual(sample_loc, vec3(1.0)));
+
+                if (is_in_bounds || u_clamp_at_border) {
+                    val = colorToVal($get_data(sample_loc));
+                } else {
+                    val = categorical_bg_value;
+                }
+
+                G.x += val * -float(i) *
+                    (1 + float(j == 0 || k == 0) + 2 * float(j == 0 && k == 0));
+                G.y += val * -float(j) *
+                    (1 + float(i == 0 || k == 0) + 2 * float(i == 0 && k == 0));
+                G.z += val * -float(k) *
+                    (1 + float(i == 0 || j == 0) + 2 * float(i == 0 && j == 0));
+
+                next.x += int(i == 1 && j == 0 && k == 0) * val;
+                next.y += int(i == 0 && j == 1 && k == 0) * val;
+                next.z += int(i == 0 && j == 0 && k == 1) * val;
+                prev.x += int(i == -1 && j == 0 && k == 0) * val;
+                prev.y += int(i == 0 && j == -1 && k == 0) * val;
+                prev.z += int(i == 0 && j == 0 && k == -1) * val;
+
+                is_on_border = is_on_border || (!is_in_bounds && (i == 0 || j == 0 || k == 0));
+            }
+        }
+    }
+
+    if (is_on_border && u_clamp_at_border) {
+        // fallback to simple gradient calculation if we are on the border
+        // and clamping is enabled (old behavior with dark/hollow faces at the border)
+        // this makes the faces in `fast` and `smooth` look the same in both clamping modes
+        G = next - prev;
+    } else {
+        // add to the gradient where the adjacent voxels are both background
+        // to fix dim pixels due to poor normal estimation
+        G.x = G.x + (next.x - current_val) * 2.0 * detectAdjacentBackground(prev.x, next.x);
+        G.y = G.y + (next.y - current_val) * 2.0 * detectAdjacentBackground(prev.y, next.y);
+        G.z = G.z + (next.z - current_val) * 2.0 * detectAdjacentBackground(prev.z, next.z);
+    }
+
+    return G;
+}
+"""
+
 ISO_CATEGORICAL_SNIPPETS = {
     'before_loop': """
         vec4 color3 = vec4(0.0);  // final color
@@ -251,7 +295,24 @@
 
 shaders = BaseVolume._shaders.copy()
 before, after = shaders['fragment'].split('void main()')
-shaders['fragment'] = before + FUNCTION_DEFINITIONS + 'void main()' + after
+FAST_GRADIENT_SHADER = (
+    before
+    + FUNCTION_DEFINITIONS
+    + FAST_GRADIENT_DEFINITION
+    + CALCULATE_COLOR_DEFINITION
+    + 'void main()'
+    + after
+)
+SMOOTH_GRADIENT_SHADER = (
+    before
+    + FUNCTION_DEFINITIONS
+    + SMOOTH_GRADIENT_DEFINITION
+    + CALCULATE_COLOR_DEFINITION
+    + 'void main()'
+    + after
+)
+
+shaders['fragment'] = FAST_GRADIENT_SHADER
 
 rendering_methods = BaseVolume._rendering_methods.copy()
 rendering_methods['iso_categorical'] = ISO_CATEGORICAL_SNIPPETS
@@ -268,6 +329,7 @@ class Volume(TextureMixin, BaseVolume):
     def __init__(self, *args, **kwargs) -> None:  # type: ignore [no-untyped-def]
         super().__init__(*args, **kwargs)
         self.unfreeze()
+        self.clamp_at_border = False
         self.iso_gradient_mode = IsoCategoricalGradientMode.FAST.value
         self.freeze()
 
@@ -278,7 +340,29 @@ def iso_gradient_mode(self) -> str:
     @iso_gradient_mode.setter
     def iso_gradient_mode(self, value: str) -> None:
         self._iso_gradient_mode = IsoCategoricalGradientMode(value)
-        self.shared_program['u_iso_gradient'] = (
-            self._iso_gradient_mode == IsoCategoricalGradientMode.SMOOTH
+        self.shared_program.frag = (
+            SMOOTH_GRADIENT_SHADER
+            if value == IsoCategoricalGradientMode.SMOOTH
+            else FAST_GRADIENT_SHADER
         )
+        self.shared_program['u_clamp_at_border'] = self._clamp_at_border
+        self.update()
+
+    @property
+    def clamp_at_border(self) -> bool:
+        """Clamp values beyond volume limits when computing isosurface gradients.
+
+        This has an effect on the appearance of labels at the border of the volume.
+
+            True: labels will appear darker at the border. [DEFAULT]
+
+            False: labels will appear brighter at the border, as if the volume extends beyond its
+            actual limits but the labels do not.
+        """
+        return self._clamp_at_border
+
+    @clamp_at_border.setter
+    def clamp_at_border(self, value: bool) -> None:
+        self._clamp_at_border = value
+        self.shared_program['u_clamp_at_border'] = self._clamp_at_border
         self.update()