Repository: napari/napari

Issue:
High memory usage - labels layers with high max values
## 🐛 Bug
When viewing a labels layer with very high values (e.g. >10e7), there's a spike in memory usage (in `0.4.8rc6`). This often crashes napari.

Presumably linked to #2723, this only happens when (I assume) the lookup isn't computed when the layer is loaded. If `array[0]` contains these high values, it works fine, but if not, the memory spike occurs when trying to view them (see example).

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

**This works fine**
```python
label_array = (5 * 10e7 * np.ones((2, 2, 2))).astype(np.uint32)
viewer = napari.view_labels(label_array)
napari.run()
```

**This causes a 10GB spike in memory usage**
```python
label_array = (5 * 10e7 * np.ones((2, 2, 2))).astype(np.uint32)
label_array[0, :, :] = 0
viewer = napari.view_labels(label_array)
napari.run()
```

## Environment

```
napari: 0.4.8rc6
Platform: Linux-4.18.2-041802-generic-x86_64-with-glibc2.10
System: Ubuntu 18.04.5 LTS
Python: 3.8.5 (default, Sep 4 2020, 07:30:14) [GCC 7.3.0]
Qt: 5.15.2
PySide2: 5.15.2
NumPy: 1.19.5
SciPy: 1.6.0
Dask: 2021.01.0
VisPy: 0.6.6

OpenGL:
- GL version: 4.6.0 NVIDIA 460.56
- MAX_TEXTURE_SIZE: 32768

Screens:
- screen 1: resolution 1920x1080, scale 1.0

Plugins:
- brainglobe-io: 0.0.2
- brainreg-segment: 0.2.7
- brainreg-standard: 0.0.2
- brainreg: 0.2.5.dev0+gd4960af.d20210310
- brainreg_standard: 0.2.5.dev0+gd4960af.d20210310
- cellfinder: 0.2.1
- console: 0.0.3
- ndtiffs: 0.1.1
- scikit-image: 0.4.8rc6
- svg: 0.1.4
- ```



Discussion:
This issue has been mentioned on **Image.sc Forum**. There might be relevant details there:

https://forum.image.sc/t/brainreg-tutorial-error-in-napari/52702/21

*sigh* thanks @adamltyson! @DragaDoncila and I tried to account for this case, but it must have slipped through somehow.

The code logic is here

https://github.com/napari/napari/blob/5116021e1a1ea96176560f2880f1369134d4ba5b/napari/layers/labels/labels.py#L763-L776

basically, if we encounter a new value that is bigger than the lookup array, we set the lookup function again, and *in theory* it should go back to the old mode when it finds that value. I don't know why it's not doing so, but will investigate. :facepalm: 

btw @DragaDoncila, *just above* that code, we have:

https://github.com/napari/napari/blob/5116021e1a1ea96176560f2880f1369134d4ba5b/napari/layers/labels/labels.py#L753-L761

ie when the view mode is "SELECTED" we don't do the switch. You don't happen to be setting the view mode, @adamltyson, do you?
sorry, the setting on the UI is "show selected", on the API it's `labels_layer.show_selected_label`:

https://github.com/napari/napari/blob/5116021e1a1ea96176560f2880f1369134d4ba5b/napari/layers/labels/labels.py#L849-L858

It also looks like using direct color mode will also explode the memory. :facepalm: 
Hi @jni, nope I'm not touching `labels_layer.show_selected_label`. The example above reproduces the issue, but our normal use case is [here](https://github.com/brainglobe/brainglobe-napari-io/blob/1eb4b2038dc79886fda507ae38c415af93f41c45/brainglobe_napari_io/brainreg/reader_dir.py#L85). 
In my case to reproduce bug I need to change the layer using a slider. 
Sorry, yes I forgot that bit. The memory spike happens when trying to view the part of the layer with the high values.

Resolving patch:
diff --git a/napari/layers/labels/labels.py b/napari/layers/labels/labels.py
--- a/napari/layers/labels/labels.py
+++ b/napari/layers/labels/labels.py
@@ -754,34 +754,36 @@ def _lookup_with_index(self, im, selected_label=None):
             Value of selected label to color, by default None
         """
         if selected_label:
-            selected_color = low_discrepancy_image(selected_label, self._seed)
             if selected_label > len(self._all_vals):
-                self._all_vals = low_discrepancy_image(
-                    np.arange(selected_label + 1), self._seed
+                self._color_lookup_func = self._get_color_lookup_func(
+                    im, max(np.max(im), selected_label)
+                )
+            if (
+                self._color_lookup_func
+                == self._lookup_with_low_discrepancy_image
+            ):
+                image = self._color_lookup_func(im, selected_label)
+            else:
+                colors = np.zeros_like(self._all_vals)
+                colors[selected_label] = low_discrepancy_image(
+                    selected_label, self._seed
                 )
-            colors = np.zeros(len(self._all_vals))
-            colors[selected_label] = selected_color
-            image = colors[im]
+                image = colors[im]
         else:
             try:
                 image = self._all_vals[im]
             except IndexError:
-                max_val = np.max(im)
                 self._color_lookup_func = self._get_color_lookup_func(
-                    im, max_val
+                    im, np.max(im)
                 )
                 if (
                     self._color_lookup_func
-                    is self._lookup_with_low_discrepancy_image
+                    == self._lookup_with_low_discrepancy_image
                 ):
                     # revert to "classic" mode converting all pixels since we
                     # encountered a large value in the raw labels image
                     image = self._color_lookup_func(im, selected_label)
                 else:
-                    self._all_vals = low_discrepancy_image(
-                        np.arange(max_val + 1), self._seed
-                    )
-                    self._all_vals[0] = 0
                     image = self._all_vals[im]
         return image
 
@@ -815,6 +817,11 @@ def _get_color_lookup_func(self, data, max_label_val):
         if max_label_val * nbytes_low_discrepancy > max_nbytes:
             return self._lookup_with_low_discrepancy_image
         else:
+            if self._all_vals.size < max_label_val + 1:
+                self._all_vals = low_discrepancy_image(
+                    np.arange(max_label_val + 1), self._seed
+                )
+                self._all_vals[0] = 0
             return self._lookup_with_index
 
     def _raw_to_displayed(self, raw):


Test patch:
diff --git a/napari/layers/labels/_tests/test_labels.py b/napari/layers/labels/_tests/test_labels.py
--- a/napari/layers/labels/_tests/test_labels.py
+++ b/napari/layers/labels/_tests/test_labels.py
@@ -884,3 +884,24 @@ def test_cursor_size_with_negative_scale():
     layer = Labels(np.zeros((5, 5), dtype=int), scale=[-1, -1])
     layer.mode = 'paint'
     assert layer.cursor_size > 0
+
+
+def test_switching_display_func_during_slicing():
+    label_array = (5e6 * np.ones((2, 2, 2))).astype(np.uint64)
+    label_array[0, :, :] = [[0, 1], [2, 3]]
+    layer = Labels(label_array)
+    layer._dims_point = (1, 0, 0)
+    layer._set_view_slice()
+    assert layer._color_lookup_func == layer._lookup_with_low_discrepancy_image
+    assert layer._all_vals.size < 1026
+
+
+def test_add_large_colors():
+    label_array = (5e6 * np.ones((2, 2, 2))).astype(np.uint64)
+    label_array[0, :, :] = [[0, 1], [2, 3]]
+    layer = Labels(label_array)
+    assert len(layer._all_vals) == layer.num_colors
+
+    layer.show_selected_label = True
+    layer.selected_label = int(5e6)
+    assert layer._all_vals.size < 1026