Storing patches embeddings as an AnnData object instead of images

Author: quentinblampey

CHANGELOG.md

@@ -4,6 +4,9 @@
 - Added H-optimus-0 model for H&E patches embeddings @stergioc (#208)
 - Can provide `qv_threshold` argument to the Xenium reader to filter low quality transcripts @callum-jpg (#210)
 
+### Changed
+- Storing patches embeddings as an `AnnData` object instead of images
+
 ### Fixed
 - Right sorting scales for WSI reader with openslide backend @stergioc (#209)
 - When a polygon cannot be simplified (for the Xenium Explorer), convert it to a circle (#206)

sopa/patches/_inference.py

@@ -3,7 +3,6 @@
 
 import numpy as np
 import torch
-from spatialdata.transformations import Scale, Sequence, get_transformation
 from xarray import DataArray, DataTree
 
 from . import models
@@ -76,13 +75,6 @@ def infer_bboxes(self, bboxes: np.ndarray) -> torch.Tensor:
 
         return embedding.cpu()  # shape (B, output_dim)
 
-    def get_patches_transformations(self, patch_overlap: float) -> dict[str, Sequence]:
-        image_transformations = get_transformation(self.image, get_all=True)
-
-        patch_step = self.patch_width - patch_overlap
-        to_image = Sequence([Scale([patch_step, patch_step], axes=("x", "y"))])
-        return {cs: to_image.compose_with(t) for cs, t in image_transformations.items()}
-
 
 def _get_extraction_parameters(
     image: DataArray | DataTree, level: int | None, magnification: int | None

sopa/patches/cluster.py

@@ -1,16 +1,12 @@
 from typing import Callable
 
-import anndata
-import numpy as np
 import scanpy as sc
+from anndata import AnnData
 from spatialdata import SpatialData
-from xarray import DataArray
 
-from sopa._constants import SopaKeys
 
-
-def leiden_clustering(X: np.ndarray, flavor: str = "igraph", **kwargs):
-    adata = anndata.AnnData(X=X)
+def leiden_clustering(adata: AnnData, flavor: str = "igraph", **kwargs):
+    adata = adata.copy()
     sc.pp.pca(adata)
     sc.pp.neighbors(adata)
     sc.tl.leiden(adata, flavor=flavor, **kwargs)
@@ -23,8 +19,8 @@ def leiden_clustering(X: np.ndarray, flavor: str = "igraph", **kwargs):
 
 
 def cluster_embeddings(
-    sdata: SpatialData,
-    element: DataArray | str,
+    sdata: SpatialData | None,
+    element: AnnData | str,
     method: Callable | str = "leiden",
     key_added: str = "cluster",
     **method_kwargs: str,
@@ -35,23 +31,18 @@ def cluster_embeddings(
         The clusters are added to the `key_added` column of the "inference_patches" shapes (`key_added='cluster'` by default).
 
     Args:
-        sdata: A `SpatialData` object
-        element: The `DataArray` containing the embeddings, or the name of the element
-        method: Callable that takes as an input an array of size `(n_patches x embedding_size)` and returns an array of clusters of size `n_patches`, or an available method name (`leiden`)
-        key_added: The key containing the clusters to be added to the patches `GeoDataFrame`
+        sdata: A `SpatialData` object. Can be `None` if element is an `AnnData` object.
+        element: The `AnnData` containing the embeddings, or the name of the element
+        method: Callable that takes as an AnnData object and returns an array of clusters of size `n_obs`, or an available method name (`leiden`)
+        key_added: The key containing the clusters to be added to the `element.obs`
         method_kwargs: kwargs provided to the method callable
     """
     if isinstance(element, str):
-        element: DataArray = sdata.images[element]
+        element: AnnData = sdata.tables[element]
 
     if isinstance(method, str):
         assert method in METHODS_DICT, f"Method {method} is not available. Use one of: {', '.join(METHODS_DICT.keys())}"
         method = METHODS_DICT[method]
 
-    gdf_patches = sdata[SopaKeys.EMBEDDINGS_PATCHES]
-
-    ilocs = np.array(list(gdf_patches.ilocs))
-    embeddings = element.compute().data[:, ilocs[:, 1], ilocs[:, 0]].T
-
-    gdf_patches[key_added] = method(embeddings, **method_kwargs)
-    gdf_patches[key_added] = gdf_patches[key_added].astype("category")
+    element.obs[key_added] = method(element, **method_kwargs)
+    element.obs[key_added] = element.obs[key_added].astype("category")

sopa/patches/infer.py

@@ -1,10 +1,10 @@
 import logging
 from typing import Callable
 
-import numpy as np
 import tqdm
+from anndata import AnnData
 from spatialdata import SpatialData
-from spatialdata.models import Image2DModel
+from spatialdata.models import TableModel
 from xarray import DataArray, DataTree
 
 from .._constants import SopaAttrs, SopaKeys
@@ -76,17 +76,23 @@ def compute_embeddings(
     if len(predictions.shape) == 1:
         predictions = torch.unsqueeze(predictions, 1)
 
-    output_image = np.zeros((predictions.shape[1], *patches.shape), dtype=np.float32)
-    for (loc_x, loc_y), pred in zip(patches.ilocs, predictions):
-        output_image[:, loc_y, loc_x] = pred
+    patches.add_shapes(key_added=SopaKeys.EMBEDDINGS_PATCHES)
 
-    output_image = DataArray(output_image, dims=("c", "y", "x"))
-    output_image = Image2DModel.parse(output_image, transformations=infer.get_patches_transformations(patch_overlap))
+    gdf = sdata[SopaKeys.EMBEDDINGS_PATCHES]
 
-    key_added = key_added or f"{infer.model_str}_embeddings"
-    add_spatial_element(sdata, key_added, output_image)
+    adata = AnnData(predictions.numpy())
+    adata.obs["region"] = SopaKeys.EMBEDDINGS_PATCHES
+    adata.obs["instance"] = gdf.index.values
+    adata = TableModel.parse(
+        adata,
+        region=SopaKeys.EMBEDDINGS_PATCHES,
+        region_key="region",
+        instance_key="instance",
+    )
+    adata.obsm["spatial"] = gdf.centroid.get_coordinates().values
 
-    patches.add_shapes(key_added=SopaKeys.EMBEDDINGS_PATCHES)
+    key_added = key_added or f"{infer.model_str}_embeddings"
+    add_spatial_element(sdata, key_added, adata)
 
 
 def _get_image_for_inference(sdata: SpatialData, image_key: str | None = None) -> DataArray | DataTree: