Repository: scikit-image/scikit-image

Issue:
Segfault in peak_local_max with large numbed of segments
## Description

scikit-image dives to (absolutely uncatchable and untrackable) segfault in peak_local_max.

## Way to reproduce
```python
import numpy as np
from scipy.ndimage import distance_transform_edt
from skimage.feature import peak_local_max


def segment(binary_image):
    distance = distance_transform_edt(binary_image)
    peak_local_max(
        distance, min_distance=100, footprint=np.ones((3, 3)), labels=binary_image,
    )

for p in [0.05, 0.95, 0.001, 0.999]:
    print(p)
    segment(np.random.random([2048, 2048]) < p)


```


## Version information
```python
# Paste the output of the following python commands
from __future__ import print_function
import sys; print(sys.version)
import platform; print(platform.platform())
import skimage; print(f'scikit-image version: {skimage.__version__}')
import numpy; print(f'numpy version: {numpy.__version__}')
```

```python
3.8.10 (default, Sep 28 2021, 16:10:42) 
[GCC 9.3.0]
Linux-5.10.47-linuxkit-x86_64-with-glibc2.29
scikit-image version: 0.18.3
numpy version: 1.21.4
```



Discussion:
additional comment: looks like 0.17.2 does not have this issue
(which I can't use on mac because wheels are not provided for mac + python 3.8)
>  (which I can't use on mac because wheels are not provided for mac + python 3.8)

Um, they are...?

https://pypi.org/project/scikit-image/0.17.2/#files

Are you on an older macOS perhaps?

At any rate, though, thanks for the bug report, we need to fix it. We (by which I mean @rfezzani 🙏 🙏 🙏 ) fixed a large number of bugs with 0.18, so it'd be great to get you using that version. Currently your reproducer works fine for me on main + macOS, though, so I'm wondering whether you could try to run it on main and see if you still have the issue? I'll also try it on Linux later.
@jni it is reproduced in linux
I think you're right about wheels are available, but for some reasons they are not recognized. 
Mac os is fresh: 11.5.2
```python
~ pip install scikit-image== --only-binary=scikit-image
ERROR: Could not find a version that satisfies the requirement scikit-image== (from versions: 0.18.0rc0, 0.18.0rc1, 0.18.0rc2, 0.18.0, 0.18.1, 0.18.2rc1, 0.18.2rc2, 0.18.2, 0.18.3)
ERROR: No matching distribution found for scikit-image==
```
@jni tested on main branch and it did not fail (on mac). Do you have guesses where is was fixed?
Hmm, on main on *Linux* I get a "Killed" message. So I guess it *wasn't* fixed and it is platform specific? Which is, well, the worst. 😂 Can you check main on your Linux machine?

At any rate, to check potential changes to this function I recommend `git log skimage/feature/peak.py`. It shows a few changes since 0.18 (Dec 2020), but none that I think could account for a fixed segfault here.

However, I stepped through in a debugger and traced the Killed error to this line:

https://github.com/scikit-image/scikit-image/blob/9fea28b7d5708152e0b96dea2f3635fa2fe6d4a0/skimage/_shared/coord.py#L33

called with ~85,000 points as input. I can reproduce the error with just this snippet:

```python
import numpy as np
from scipy.spatial import cKDTree

coord = np.random.randint(low=0, high=1848, size=(85000, 2))
tree = cKDTree(coord)
indices = tree.query_ball_point(coord, r=100, p=np.inf)
```

(By the way, @rfezzani, why is it np.inf, shouldn't we use p=2??? But that doesn't change the error anyway.)

So, I don't know at this point whether there's anything to be done here or we should raise upstream to SciPy...
Thank you for your investigation @jni. https://github.com/scipy/scipy/issues/12956 reports memory issues with cKDTree but the memory saving strategy suggested there is already implemented in `scikit-image` :confused: 

> (By the way, @rfezzani, why is it np.inf, shouldn't we use p=2??? But that doesn't change the error anyway.)

`np.inf` is used for retro-compatibility with previous implementation :wink:.
@rfezzani 

> the memory saving strategy suggested there is already implemented in scikit-image 😕

Yeah, I remembered, but, do we actually? Is there a bug somewhere? Walking through the code now, how did I get to 85,000 points? I get to that in the first call to _ensure_spacing`...! See the code I see that it shouldn't happen with a batch size of 50...!

> np.inf is used for retro-compatibility with previous implementation 😉

I feel like that's a bug and we should use 2?
> Yeah, I remembered, but, do we actually? Is there a bug somewhere? Walking through the code now, how did I get to 85,000 points? I get to that in the first call to _ensure_spacing`...! See the code I see that it shouldn't happen with a batch size of 50...!

OK, I got it: we are not using a batch size, but a minimum batch size of 50! This comes for the strategy we adopted to manage a worst case https://github.com/scikit-image/scikit-image/pull/5062#pullrequestreview-535202840! Maybe should we change our mind.

> I feel like that's a bug and we should use 2?

I don't really see it as a bug, but we can change this and break the behaviour of v0.18

@rfezzani but I think we just found a bug in the implementation! Here's what happens when I look at the split indices in debugging:

```python
>>> split_idx                                                                                        
array([ 85098, 127647, 148921, 159558, 164876, 167535, 168864, 169528,                               
       169860, 170026, 170109])        
```                                                              
ie we are using *decreasing* batch sizes instead of increasing them!!!! :facepalm: From the discussion in #5062:

> So you build with the top 10, 20, 40, ... n peaks, removing points each time.

So, currently, we have an 85000 point cKDTree and check 85000 points, when in fact we should have a much *smaller* cKDTree by the time we get to the final 85000 points!
:facepalm: oups!...

Resolving patch:
diff --git a/skimage/_shared/coord.py b/skimage/_shared/coord.py
--- a/skimage/_shared/coord.py
+++ b/skimage/_shared/coord.py
@@ -59,7 +59,7 @@ def _ensure_spacing(coord, spacing, p_norm, max_out):
 
 
 def ensure_spacing(coords, spacing=1, p_norm=np.inf, min_split_size=50,
-                   max_out=None):
+                   max_out=None, *, max_split_size=2000):
     """Returns a subset of coord where a minimum spacing is guaranteed.
 
     Parameters
@@ -74,10 +74,19 @@ def ensure_spacing(coords, spacing=1, p_norm=np.inf, min_split_size=50,
         ``inf`` corresponds to the Chebyshev distance and 2 to the
         Euclidean distance.
     min_split_size : int
-        Minimum split size used to process ``coord`` by batch to save
+        Minimum split size used to process ``coords`` by batch to save
         memory. If None, the memory saving strategy is not applied.
     max_out : int
         If not None, only the first ``max_out`` candidates are returned.
+    max_split_size : int
+        Maximum split size used to process ``coords`` by batch to save
+        memory. This number was decided by profiling with a large number
+        of points. Too small a number results in too much looping in
+        Python instead of C, slowing down the process, while too large
+        a number results in large memory allocations, slowdowns, and,
+        potentially, in the process being killed -- see gh-6010. See
+        benchmark results `here
+        <https://github.com/scikit-image/scikit-image/pull/6035#discussion_r751518691>`_.
 
     Returns
     -------
@@ -94,9 +103,12 @@ def ensure_spacing(coords, spacing=1, p_norm=np.inf, min_split_size=50,
             batch_list = [coords]
         else:
             coord_count = len(coords)
-            split_count = int(np.log2(coord_count / min_split_size)) + 1
-            split_idx = np.cumsum(
-                [coord_count // (2 ** i) for i in range(1, split_count)])
+            split_idx = [min_split_size]
+            split_size = min_split_size
+            while coord_count - split_idx[-1] > max_split_size:
+                split_size *= 2
+                split_idx.append(split_idx[-1] + min(split_size,
+                                                     max_split_size))
             batch_list = np.array_split(coords, split_idx)
 
         output = np.zeros((0, coords.shape[1]), dtype=coords.dtype)


Test patch:
diff --git a/skimage/_shared/tests/test_coord.py b/skimage/_shared/tests/test_coord.py
--- a/skimage/_shared/tests/test_coord.py
+++ b/skimage/_shared/tests/test_coord.py
@@ -1,6 +1,9 @@
-import pytest
+import time
+
 import numpy as np
+import pytest
 from scipy.spatial.distance import pdist, minkowski
+
 from skimage._shared.coord import ensure_spacing
 
 
@@ -56,6 +59,24 @@ def test_ensure_spacing_batch_processing(p, size):
                           expected)
 
 
+def test_max_batch_size():
+    """Small batches are slow, large batches -> large allocations -> also slow.
+
+    https://github.com/scikit-image/scikit-image/pull/6035#discussion_r751518691
+    """
+    coords = np.random.randint(low=0, high=1848, size=(40000, 2))
+    tstart = time.time()
+    ensure_spacing(coords, spacing=100, min_split_size=50,
+                   max_split_size=2000)
+    dur1 = time.time() - tstart
+
+    tstart = time.time()
+    ensure_spacing(coords, spacing=100, min_split_size=50,
+                   max_split_size=20000)
+    dur2 = time.time() - tstart
+
+    assert dur1 < dur2
+
 
 @pytest.mark.parametrize("p", [1, 2, np.inf])
 @pytest.mark.parametrize("size", [30, 50, None])