Merge pull request #211 from ArneDefauw/visium_hd

Visium hd

Author: LucaMarconato

src/spatialdata_io/readers/visium_hd.py

@@ -16,8 +16,9 @@
 from geopandas import GeoDataFrame
 from imageio import imread as imread2
 from multiscale_spatial_image import MultiscaleSpatialImage
+from numpy.random import default_rng
 from spatial_image import SpatialImage
-from spatialdata import SpatialData
+from spatialdata import SpatialData, rasterize_bins, rasterize_bins_link_table_to_labels
 from spatialdata.models import Image2DModel, ShapesModel, TableModel
 from spatialdata.transformations import (
     Affine,
@@ -31,6 +32,8 @@
 from spatialdata_io._constants._constants import VisiumHDKeys
 from spatialdata_io._docs import inject_docs
 
+RNG = default_rng(0)
+
 
 @inject_docs(vx=VisiumHDKeys)
 def visium_hd(
@@ -39,6 +42,7 @@ def visium_hd(
     filtered_counts_file: bool = True,
     bin_size: int | list[int] | None = None,
     bins_as_squares: bool = True,
+    annotate_table_by_labels: bool = False,
     fullres_image_file: str | Path | None = None,
     load_all_images: bool = False,
     imread_kwargs: Mapping[str, Any] = MappingProxyType({}),
@@ -67,6 +71,9 @@ def visium_hd(
     bins_as_squares
         If `True`, the bins are represented as squares. If `False`, the bins are represented as circles. For a correct
         visualization one should use squares.
+    annotate_table_by_labels
+        If `True`, the tables will annotate labels layers representing the bins, if `False`, the tables will annotate
+        shapes layer.
     fullres_image_file
         Path to the full-resolution image. By default the image is searched in the ``{vx.MICROSCOPE_IMAGE!r}``
         directory.
@@ -89,6 +96,7 @@ def visium_hd(
     tables = {}
     shapes = {}
     images: dict[str, Any] = {}
+    labels: dict[str, Any] = {}
 
     if dataset_id is None:
         dataset_id = _infer_dataset_id(path)
@@ -347,7 +355,32 @@ def _get_bins(path_bins: Path) -> list[str]:
         affine1 = transform_matrices["spot_colrow_to_microscope_colrow"]
         set_transformation(image, Sequence([affine0, affine1]), "global")
 
-    return SpatialData(tables=tables, images=images, shapes=shapes)
+    sdata = SpatialData(tables=tables, images=images, shapes=shapes, labels=labels)
+
+    if annotate_table_by_labels:
+        for bin_size_str in bin_sizes:
+
+            shapes_name = dataset_id + "_" + bin_size_str
+
+            # add labels layer (rasterized bins).
+            labels_name = f"{dataset_id}_{bin_size_str}_labels"
+
+            labels_element = rasterize_bins(
+                sdata,
+                bins=shapes_name,
+                table_name=bin_size_str,
+                row_key=VisiumHDKeys.ARRAY_ROW,
+                col_key=VisiumHDKeys.ARRAY_COL,
+                value_key=None,
+                return_region_as_labels=True,
+            )
+
+            sdata[labels_name] = labels_element
+            rasterize_bins_link_table_to_labels(
+                sdata=sdata, table_name=bin_size_str, rasterized_labels_name=labels_name
+            )
+
+    return sdata
 
 
 def _infer_dataset_id(path: Path) -> str: