Merge pull request #92 from sagar87/fix/export_anndata

added colors to anndata object when exporting
sagar87 · Sep 3, 2024 · 5f5be5d · 5f5be5d
2 parents c4c9cc0 + 800b555
commit 5f5be5d
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/spatialproteomics/pp/preprocessing.py b/spatialproteomics/pp/preprocessing.py
@@ -1313,13 +1313,16 @@ def merge_segmentation(
 
         return xr.merge([obj, da])
 
-    def get_layer_as_df(self, layer: str = Layers.OBS, celltypes_to_str: bool = True):
+    def get_layer_as_df(
+        self, layer: str = Layers.OBS, celltypes_to_str: bool = True, idx_to_str: bool = False
+    ) -> pd.DataFrame:
         """
         Returns the specified layer as a pandas DataFrame.
 
         Parameters:
             layer (str): The name of the layer to retrieve. Defaults to Layers.OBS.
             celltypes_to_str (bool): Whether to convert celltype labels to strings. Defaults to True.
+            idx_to_str (bool): Whether to convert the index to strings. Defaults to False.
 
         Returns:
             pandas.DataFrame: The layer data as a DataFrame.
@@ -1336,6 +1339,9 @@ def get_layer_as_df(self, layer: str = Layers.OBS, celltypes_to_str: bool = True
             label_dict = self._obj.la._label_to_dict(Props.NAME)
             df[Features.LABELS] = df[Features.LABELS].apply(lambda x: label_dict[x])
 
+        if idx_to_str:
+            df.index = df.index.astype(str)
+
         return df
 
     def get_disconnected_cell(self) -> int:

diff --git a/spatialproteomics/tl/tool.py b/spatialproteomics/tl/tool.py
@@ -5,7 +5,7 @@
 import xarray as xr
 
 from ..base_logger import logger
-from ..constants import Dims, Layers
+from ..constants import Dims, Features, Layers, Props
 from ..pp.utils import _normalize
 from .utils import _convert_masks_to_data_array, _get_channels
 
@@ -466,7 +466,20 @@ def convert_to_anndata(
         adata.var_names = self._obj.coords[Dims.CHANNELS].values
 
         if obs_key in self._obj:
-            adata.obs = pd.DataFrame(self._obj[obs_key], columns=self._obj.coords[Dims.FEATURES])
+            adata.obs = self._obj.pp.get_layer_as_df(obs_key, idx_to_str=True)
+
+            # if we have labels and colors for them, we add them to the anndata object
+            if Dims.LABELS in self._obj.dims and Layers.PROPERTIES in self._obj:
+                properties = self._obj.pp.get_layer_as_df(Layers.PROPERTIES)
+                if Props.COLOR in properties.columns:
+                    # putting it into the anndata object
+                    adata.uns[f"{Features.LABELS}_colors"] = list(properties[Props.COLOR].values)
+
+            # to be compatible with squidpy out of the box, a spatial key is added to obsm if possible
+            if Features.X in adata.obs and Features.Y in adata.obs:
+                adata.obs[Features.X] = adata.obs[Features.X].astype(float)
+                adata.obs[Features.Y] = adata.obs[Features.Y].astype(float)
+                adata.obsm["spatial"] = np.array(adata.obs[[Features.X, Features.Y]])
 
         if additional_uns:
             for key, layer in additional_uns.items():