add argument return_background for get_centroids and get_element_instances in the case of labels

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning][].
 
 ## [0.x.x] - 2024-xx-xx
 
+-   Add `return_background` as argument to `get_centroids` and `get_element_instances` #621
+
 ## [0.2.1] - 2024-07-04
 
 ### Minor

src/spatialdata/_core/centroids.py

@@ -35,6 +35,7 @@ def _validate_coordinate_system(e: SpatialElement, coordinate_system: str) -> No
 def get_centroids(
     e: SpatialElement,
     coordinate_system: str = "global",
+    return_background: bool = False,
 ) -> DaskDataFrame:
     """
     Get the centroids of the geometries contained in a SpatialElement, as a new Points element.
@@ -45,6 +46,8 @@ def get_centroids(
         The SpatialElement. Only points, shapes (circles, polygons and multipolygons) and labels are supported.
     coordinate_system
         The coordinate system in which the centroids are computed.
+    return_background
+        If True, the centroid of the background label (0) is included in the output.
 
     Notes
     -----
@@ -69,7 +72,7 @@ def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
     -------
     pd.DataFrame
         A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
-        The index of the DataFrame is the collection of label values, sorted ascendingly.
+        The index of the DataFrame is the collection of label values, sorted in ascending order.
     """
     centroids: dict[int, float] = defaultdict(float)
     for i in xdata[axis]:
@@ -95,6 +98,7 @@ def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
 def _(
     e: DataArray | DataTree,
     coordinate_system: str = "global",
+    return_background: bool = False,
 ) -> DaskDataFrame:
     """Get the centroids of a Labels element (2D or 3D)."""
     model = get_model(e)
@@ -110,6 +114,8 @@ def _(
     for axis in get_axes_names(e):
         dfs.append(_get_centroids_for_axis(e, axis))
     df = pd.concat(dfs, axis=1)
+    if not return_background and 0 in df.index:
+        df = df.drop(index=0)  # drop the background label
     t = get_transformation(e, coordinate_system)
     centroids = PointsModel.parse(df, transformations={coordinate_system: t})
     return transform(centroids, to_coordinate_system=coordinate_system)

src/spatialdata/_core/query/relational_query.py

@@ -86,6 +86,7 @@ def _filter_table_by_element_names(table: AnnData | None, element_names: str | l
 @singledispatch
 def get_element_instances(
     element: SpatialElement,
+    return_background: bool = False,
 ) -> pd.Index:
     """
     Get the instances (index values) of the SpatialElement.
@@ -94,6 +95,8 @@ def get_element_instances(
     ----------
     element
         The SpatialElement.
+    return_background
+        If True, the background label (0) is included in the output.
 
     Returns
     -------
@@ -106,6 +109,7 @@ def get_element_instances(
 @get_element_instances.register(DataTree)
 def _(
     element: DataArray | DataTree,
+    return_background: bool = False,
 ) -> pd.Index:
     model = get_model(element)
     assert model in [Labels2DModel, Labels3DModel], "Expected a `Labels` element. Found an `Image` instead."
@@ -119,7 +123,10 @@ def _(
         xdata = next(iter(v))
         # can be slow
         instances = da.unique(xdata.data).compute()
-    return pd.Index(np.sort(instances))
+    index = pd.Index(np.sort(instances))
+    if not return_background and 0 in index:
+        return index.drop(0)  # drop the background label
+    return index
 
 
 @get_element_instances.register(GeoDataFrame)
@@ -568,7 +575,8 @@ def join_spatialelement_table(
     both the SpatialElement and table.
 
     For Points and Shapes elements every valid join for argument how is supported. For Labels elements only
-     the ``'left'`` and ``'right_exclusive'`` joins are supported.
+    the ``'left'`` and ``'right_exclusive'`` joins are supported.
+    For Labels, the background label (0) is not included in the output and it will not be returned.
 
     Parameters
     ----------

tests/core/query/test_relational_query.py

@@ -699,14 +699,15 @@ def test_labels_table_joins(full_sdata):
         table_name="table",
         how="left",
     )
-    assert all(table.obs["instance_id"] == range(100))
+
+    assert all(table.obs["instance_id"] == range(1, 100))
 
     full_sdata["table"].obs["instance_id"] = list(reversed(range(100)))
 
     element_dict, table = join_spatialelement_table(
         sdata=full_sdata, spatial_element_names="labels2d", table_name="table", how="left", match_rows="left"
     )
-    assert all(table.obs["instance_id"] == range(100))
+    assert all(table.obs["instance_id"] == range(1, 100))
 
     with pytest.warns(UserWarning, match="Element type"):
         join_spatialelement_table(
@@ -724,7 +725,8 @@ def test_labels_table_joins(full_sdata):
         sdata=full_sdata, spatial_element_names="labels2d", table_name="table", how="right_exclusive"
     )
     assert element_dict["labels2d"] is None
-    assert table is None
+    assert len(table) == 1
+    assert all(table.obs["instance_id"] == 0)  # the background value, which is filtered out effectively
 
 
 def test_points_table_joins(full_sdata):

tests/core/test_centroids.py

@@ -100,7 +100,10 @@ def test_get_centroids_shapes(shapes, coordinate_system: str, shapes_name: str):
 @pytest.mark.parametrize("coordinate_system", ["global", "aligned"])
 @pytest.mark.parametrize("is_multiscale", [False, True])
 @pytest.mark.parametrize("is_3d", [False, True])
-def test_get_centroids_labels(labels, coordinate_system: str, is_multiscale: bool, is_3d: bool):
+@pytest.mark.parametrize("return_background", [False, True])
+def test_get_centroids_labels(
+    labels, coordinate_system: str, is_multiscale: bool, is_3d: bool, return_background: bool
+):
     scale_factors = [2] if is_multiscale else None
     if is_3d:
         model = Labels3DModel
@@ -124,6 +127,8 @@ def test_get_centroids_labels(labels, coordinate_system: str, is_multiscale: boo
             },
             index=[0, 1, 2],
         )
+        if not return_background:
+            expected_centroids = expected_centroids.drop(index=0)
     else:
         array = np.array(
             [
@@ -145,11 +150,14 @@ def test_get_centroids_labels(labels, coordinate_system: str, is_multiscale: boo
 
     if coordinate_system == "aligned":
         set_transformation(element, transformation=affine, to_coordinate_system=coordinate_system)
-    centroids = get_centroids(element, coordinate_system=coordinate_system)
+    centroids = get_centroids(element, coordinate_system=coordinate_system, return_background=return_background)
 
-    labels_indices = get_element_instances(element)
+    labels_indices = get_element_instances(element, return_background=return_background)
     assert np.array_equal(centroids.index.values, labels_indices)
 
+    if not return_background:
+        assert 0 not in centroids.index
+
     if coordinate_system == "global":
         assert np.array_equal(centroids.compute().values, expected_centroids.values)
     else: