From 160f8775e7454bdd3c087b98f62c5aa1fce31f87 Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Mon, 25 Nov 2024 21:14:13 -0500
Subject: [PATCH] feat+perf(dust): add invert option and simplify operations

---
 cc3d/__init__.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/cc3d/__init__.py b/cc3d/__init__.py
index 04ebd10..f7b7e8c 100644
--- a/cc3d/__init__.py
+++ b/cc3d/__init__.py
@@ -24,6 +24,7 @@ def dust(
   in_place:bool = False,
   binary_image:bool = False,
   precomputed_ccl:bool = False,
+  invert:bool = False,
 ) -> np.ndarray:
   """
   Remove from the input image connected components
@@ -38,6 +39,8 @@ def 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.
 
   Returns: dusted image
   """
@@ -60,17 +63,21 @@ 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 invert:
+    to_retain = [ 
+      i for i in range(1, N+1) if mask_sizes[i] < threshold 
+    ]
+  else:
+    to_retain = [ 
+      i for i in range(1, N+1) if mask_sizes[i] >= threshold 
+    ]
 
-  if len(to_mask) == 0:
+  if len(to_retain) == N:
     return img
 
-  mask = np.isin(cc_labels, to_mask)
+  mask = np.isin(cc_labels, to_retain, assume_unique=True)
   del cc_labels
-  np.logical_not(mask, out=mask)
-  np.multiply(img, mask, out=img)
+  img *= mask
   return img.view(orig_dtype)
 
 def largest_k(