added min_size to mininstance sampler

constantinpape · Jul 11, 2024 · 549b47d · 549b47d
1 parent 352535b
commit 549b47d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/torch_em/data/sampler.py b/torch_em/data/sampler.py
@@ -1,5 +1,6 @@
 import numpy as np
 from typing import List
+import skimage.segmentation
 
 
 class MinForegroundSampler:
@@ -85,13 +86,18 @@ class MinInstanceSampler:
  def __init__(
  self,
  min_num_instances: int = 2,
- p_reject: float = 1.0
+ p_reject: float = 1.0,
+ min_size: int = None
  ):
  self.min_num_instances = min_num_instances
  self.p_reject = p_reject
+ self.min_size = min_size
 
  def __call__(self, x, y):
- uniques = np.unique(y)
+ uniques, sizes = np.unique(y, return_counts=True)
+ if self.min_size is not None:
+ filter_ids = uniques[sizes >= self.min_size]
+ uniques = filter_ids
  if len(uniques) >= self.min_num_instances:
  return True
  else: