fix: detect nuclei inside roi if provided (#28)

Author: bathienle

descriptor.yml

@@ -31,6 +31,12 @@ inputs:
       max_height: 4000
     description: The input image in which nuclei should be segmented
 
+  roi:
+    display_name: Input ROI
+    description: The ROI in the image in which the nuclei should be segmented
+    type: geometry
+    optional: true
+
   stardist_prob_t:
     display_name: Probability threshold
     description: Probability Threshold in range [0.0, 1.0] - higher values lead to fewer segmented objects, but will likely avoid false positives

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
   "imageio==2.34.1",
   "numpy==1.26.4",
   "pyyaml==6.0.1",
+  "shapely==2.1.2",
   "stardist==0.9.1",
   "tensorflow-cpu==2.16.1",
 ]

scripts/run.py

@@ -6,6 +6,7 @@
 import yaml
 from csbdeep.utils import normalize
 from imageio import imread
+from shapely.geometry import Point, Polygon
 from stardist.models import StarDist2D
 
 
@@ -14,6 +15,30 @@
 MODEL_DATA_DIR = "/models/"
 
 
+def is_nucleus_inside_roi(nucleus_coords: np.ndarray, roi_polygon: Polygon, image_height: int) -> bool:
+  """
+  Check if a nucleus is inside the ROI polygon.
+
+  Args:
+    nucleus_coords (np.ndarray): The nucleus coordinates in stardist format (2, N) where first row is x, second is y.
+    roi_polygon (Polygon): The ROI polygon from shapely.
+    image_height (int): The image height for coordinate transformation.
+
+  Returns:
+    bool: True if the nucleus centroid is inside the ROI, False otherwise.
+  """
+
+  # Reverse (y,x) → (x,y) and flip Y coordinates to match ROI coordinate system
+  x_coords = nucleus_coords[0]
+  y_coords = image_height - nucleus_coords[1]
+
+  # Calculate centroid
+  centroid_x = np.mean(x_coords)
+  centroid_y = np.mean(y_coords)
+
+  # Check if centroid is inside ROI
+  return roi_polygon.contains(Point(centroid_x, centroid_y))
+
 
 def from_stardist_to_geojson_string(stardist_polygroup: np.ndarray, image_height: int):
   """Converts a polygon coordinates ndarray generated by the stardist algorithm into a geojson string.
@@ -89,13 +114,22 @@ def write_array(array_path: str, array_data: Iterable[Any], format_fn: Callable[
 
 
 def main():
-  # Red inputs
+  # Read inputs
   stardist_norm_perc_low = read_parameter(os.path.join(INPUT_DIR, "stardist_norm_perc_low"), cast_fn=float, default=1.0)
   stardist_norm_perc_high = read_parameter(os.path.join(INPUT_DIR, "stardist_norm_perc_high"), cast_fn=float, default=99.0)
   stardist_prob_t = read_parameter(os.path.join(INPUT_DIR, "stardist_prob_t"), cast_fn=float, default=0.5)
   stardist_nms_t = read_parameter(os.path.join(INPUT_DIR, "stardist_nms_t"), cast_fn=float, default=0.5)
   image_path = os.path.join(INPUT_DIR, "image")
 
+  # Read ROI if it exists
+  roi_path = os.path.join(INPUT_DIR, "roi")
+  roi_polygon = None
+  if os.path.isfile(roi_path):
+    with open(roi_path, "r") as fp:
+      roi_content = fp.read().strip()
+      roi = geojson.loads(roi_content)
+      roi_polygon = Polygon(roi['coordinates'][0])
+
   # use local model file in ~/models/2D_versatile_HE/
   model = StarDist2D(None, name='2D_versatile_HE', basedir=MODEL_DATA_DIR)
 
@@ -110,7 +144,6 @@ def main():
     axis=(0, 1)  # normalize channels independently
   )
 
-
   # Stardist model prediction with thresholds
   _, details = model.predict_instances(
     img,
@@ -119,13 +152,30 @@ def main():
     n_tiles=model._guess_n_tiles(img)
   )
 
-  # writing ouputs
+  # Filter nuclei if ROI is provided
+  if roi_polygon is not None:
+    filtered_coords = []
+    filtered_probs = []
+
+    for i, nucleus_coords in enumerate(details['coord']):
+      if is_nucleus_inside_roi(nucleus_coords, roi_polygon, image_height):
+        filtered_coords.append(nucleus_coords)
+        filtered_probs.append(details['prob'][i])
+  else:
+    filtered_coords = list(details['coord'])
+    filtered_probs = details['prob'].tolist()
+
+  # writing outputs
   write_array(
     array_path=os.path.join(OUTPUT_DIR, "nuclei"),
-    array_data=details['coord'],
-    format_fn=lambda poly: from_stardist_to_geojson_string(poly, image_height)
+    array_data=filtered_coords,
+    format_fn=lambda poly: from_stardist_to_geojson_string(poly, image_height),
+  )
+  write_array(
+    array_path=os.path.join(OUTPUT_DIR, "probs"),
+    array_data=filtered_probs,
+    format_fn=str,
   )
-  write_array(array_path=os.path.join(OUTPUT_DIR, "probs"), array_data=details['prob'].tolist(), format_fn=str)
 
 
 if __name__ == "__main__":