From a7e294643df22f4c68a7a0ebb83f16ac3656886b Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Wed, 18 Dec 2024 22:06:22 -0500
Subject: [PATCH] feat(dust): add between thresholds denoted (lower,upper)

---
 automated_test.py | 27 +++++++++++++++++++++++++++
 cc3d/__init__.py  | 24 ++++++++++++++++--------
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/automated_test.py b/automated_test.py
index b7a142f..1b0ffb2 100644
--- a/automated_test.py
+++ b/automated_test.py
@@ -1184,6 +1184,33 @@ def test_dust_retain_all():
   )
   assert np.all(recovered == labels)
 
+@pytest.mark.parametrize("dtype", TEST_TYPES)
+@pytest.mark.parametrize("connectivity", (6,18,26))
+@pytest.mark.parametrize("order", ("C", "F"))
+@pytest.mark.parametrize("in_place", (False, True))
+@pytest.mark.parametrize("invert", (False, True))
+def test_dust_between_threshold(
+  dtype, connectivity, order, in_place, invert
+):
+  labels = np.zeros((100,100,10), dtype=np.uint8, order=order)
+  labels[:5,:5,:1] = 1
+  labels[20:40,20:40,:] = 2
+  labels[10:17,10:17,:] = 3
+  recovered = cc3d.dust(
+    labels, 
+    threshold=(25, 491), 
+    connectivity=connectivity,
+    in_place=in_place,
+    invert=invert,
+  )
+
+  res = [ int(x) for x in np.unique(recovered) ]
+
+  if invert:
+    assert res == [0,2]
+  else:
+    assert res == [0,1,3]
+
 @pytest.mark.parametrize("dtype", INT_TYPES)
 @pytest.mark.parametrize("connectivity", (6,18,26))
 def test_dust_random(dtype, connectivity):
diff --git a/cc3d/__init__.py b/cc3d/__init__.py
index da9c53d..9d6ee7a 100644
--- a/cc3d/__init__.py
+++ b/cc3d/__init__.py
@@ -1,5 +1,5 @@
 from typing import (
-  Dict, Union, Tuple, Iterator, 
+  Dict, Union, Tuple, List, Iterator, 
   Sequence, Optional, Any, BinaryIO
 )
 
@@ -19,7 +19,7 @@
 
 def dust(
   img:np.ndarray, 
-  threshold:Union[int,float], 
+  threshold:Union[int,float,Tuple[int,int],Tuple[float,float],List[int],List[float]], 
   connectivity:int = 26,
   in_place:bool = False,
   binary_image:bool = False,
@@ -32,15 +32,18 @@ def dust(
   can be read as a verb "to dust" the image.
 
   img: 2D or 3D image
-  threshold: discard components smaller than this in voxels
+  threshold: 
+    (int) discard components smaller than this in voxels
+    (tuple/list) keep components in range [lower, upper)
   connectivity: cc3d connectivity to use
   in_place: whether to modify the input image or perform
     dust 
   precomputed_ccl: for performance, avoid computing a CCL
     pass since the input is already a CCL output from this
     library.
-  invert: switch the operation from less than threshold to
-    greater than or equal to threshold.
+  invert: switch the threshold direction. For scalar input,
+    this means less than converts to greater than or equal to,
+    for ranged input, switch from between to outside of range.
 
   Returns: dusted image
   """
@@ -63,9 +66,14 @@ def dust(
   mask_sizes = stats["voxel_counts"]
   del stats
 
-  to_mask = [ 
-    i for i in range(1, N+1) if mask_sizes[i] < threshold 
-  ]
+  if isinstance(threshold, (tuple, list)):
+    to_mask = [ 
+      i for i in range(1, N+1) if not (threshold[0] <= mask_sizes[i] < threshold[1])
+    ]
+  else:
+    to_mask = [ 
+      i for i in range(1, N+1) if mask_sizes[i] < threshold 
+    ]
 
   if len(to_mask) == 0:
     return img