SpatialData loader and writer

src/segger/cli/main.py

@@ -1484,7 +1484,7 @@ def _resolve_transcripts():
             print(
                 "Warning: spatialdata not installed. "
                 "Install with: pip install segger[spatialdata]"
-            )
+            )  # triggered whenever the SpatialDataWrite import fails
         return
 
     raise ValueError(f"Unsupported export format: {format}")

src/segger/data/data_module.py

@@ -39,6 +39,8 @@
 )
 from .tiling import QuadTreeTiling, SquareTiling
 from .partition import PartitionSampler
+from .utils.anndata import setup_anndata
+from .utils.heterodata import setup_heterodata
 
 
 
@@ -266,8 +268,6 @@ def load(self):
             return
 
         from ..io.preprocessor import get_preprocessor
-        from .utils.anndata import setup_anndata
-        from .utils.heterodata import setup_heterodata
 
         # Load standardized IST data with quality filtering
         pp = get_preprocessor(
@@ -376,32 +376,13 @@ def _load_from_spatialdata(self, path: Path):
         ImportError
             If spatialdata is not installed.
         """
-        from ..io.spatialdata_loader import SpatialDataLoader, load_from_spatialdata
-        from ..io.quality_filter import get_quality_filter
+        from ..io.spatialdata_loader import SpatialDataLoader
+        from ..io.quality_filter import SpatialDataQualityFilter
 
         tx_fields = StandardTranscriptFields()
         bd_fields = StandardBoundaryFields()
 
-        # Load from SpatialData
-        loader = SpatialDataLoader(path)
-        transcripts_lf = loader.transcripts(normalize=True)
-        boundaries = loader.boundaries(boundary_type="all")
-
-        # Apply quality filtering if needed
-        if self.min_qv is not None and loader.platform:
-            qf = get_quality_filter(loader.platform)
-            transcripts_lf = qf.filter(
-                transcripts_lf,
-                min_threshold=self.min_qv,
-                feature_column=tx_fields.feature,
-            )
-
-        # Collect to DataFrame
-        tx = self.tx = transcripts_lf.collect()
-        bd = self.bd = boundaries
-
-        # Continue with standard processing
-        # Mask transcripts to reference segmentation
+        # Set compartments and boundary type according to segmentation mode
         if self.segmentation_graph_mode == "nucleus":
             compartments = [tx_fields.nucleus_value]
             boundary_type = bd_fields.nucleus_value
@@ -417,13 +398,31 @@ def _load_from_spatialdata(self, path: Path):
                 f"'{self.segmentation_graph_mode}'."
             )
 
-        # Check if compartment column exists
+        # Load from SpatialData
+        loader = SpatialDataLoader(path)
+        transcripts_lf = loader.transcripts(normalize=True, platform=loader.platform)
+        boundaries = loader.boundaries(boundary_type=boundary_type)
+
+       # Apply quality filtering
+        qf = SpatialDataQualityFilter(loader.platform)
+        transcripts_lf = qf.filter(
+            transcripts_lf,
+            min_threshold=self.min_qv,
+            feature_column=tx_fields.feature,
+        )
+
+        # Collect to DataFrame
+        tx = self.tx = transcripts_lf.collect()
+        bd = self.bd = boundaries
+
+        # Mask trasncripts on the compartment
         if tx_fields.compartment in tx.columns:
             tx_mask = pl.col(tx_fields.compartment).is_in(compartments)
         else:
             # If no compartment info, use cell_id presence
             tx_mask = pl.col(tx_fields.cell_id).is_not_null()
 
+        # Mask boundaries on type
         if bd is not None and bd_fields.boundary_type in bd.columns:
             bd_mask = bd[bd_fields.boundary_type] == boundary_type
         else:

src/segger/data/tile_dataset.py

@@ -303,6 +303,10 @@ def _get_sample_size(self, idx: int) -> int:
                 f"'{self.__class__.__name__}' expected {self.mode} <= "
                 f"{self.max_num}, but found {num}. Increase 'max_num' or set "
                 f"'skip_too_big=True'"
+                # for final user the message should be something like:
+                # "Graph is too large for current batch size: increase '--max-nodes-per-batch"  or 
+                #  decrease '--max-nodes-per-tile' and adjust
+                #  '--tiling-margin-training', '--tiling-margin-prediction' accordingly."
             )
         return num
 

src/segger/data/utils/anndata.py

@@ -33,7 +33,7 @@ def anndata_from_transcripts(
     """TODO: Add description.
     """
     _lazy_imports()
-    # Remove non-nuclear transcript
+    # Remove transcripts without cell assignment
     tx = tx.filter(pl.col(cell_id_column).is_not_null())
     # Get sparse counts from transcripts
     if feature_vocab is None:

src/segger/data/utils/heterodata.py

@@ -161,6 +161,11 @@ def _coerce_categorical_numeric(
         tx_fields.cell_cluster,
         tx_fields.gene_cluster,
     ]
+
+    # Convert feature_name to string for joining
+    if transcripts[tx_fields.feature].dtype != pl.Utf8:
+        transcripts = transcripts.with_columns(pl.col(tx_fields.feature).cast(pl.Utf8))
+
     # Update transcripts with fields for training
 
     transcripts = (

src/segger/export/anndata_writer.py

@@ -192,7 +192,7 @@ def write(
         x_column: Optional[str] = "x",
         y_column: Optional[str] = "y",
         z_column: Optional[str] = "z",
-        overwrite: bool = False,
+        overwrite: bool = None,
         **kwargs,
     ) -> Path:
         """Write segmentation results to AnnData (.h5ad).
@@ -215,6 +215,9 @@ def write(
         output_dir.mkdir(parents=True, exist_ok=True)
         output_path = output_dir / output_name
 
+        # If overwrite is not passed explicitly, set based on path existence
+        overwrite = output_path.exists() if overwrite is None else overwrite
+
         if output_path.exists() and not overwrite:
             raise FileExistsError(
                 f"Output path exists: {output_path}. "

src/segger/export/spatialdata_writer.py

@@ -112,7 +112,7 @@ def write(
         x_column: str = "x",
         y_column: str = "y",
         z_column: Optional[str] = "z",
-        overwrite: bool = False,
+        overwrite: bool = None,
         **kwargs,
     ) -> Path:
         """Write segmentation results to SpatialData Zarr store.
@@ -171,6 +171,9 @@ def write(
         output_dir.mkdir(parents=True, exist_ok=True)
         output_path = output_dir / output_name
 
+        # If overwrite is not passed explicitly, set based on path existence
+        overwrite = output_path.exists() if overwrite is None else overwrite
+
         # Check if exists
         if output_path.exists() and not overwrite:
             raise FileExistsError(
@@ -253,7 +256,7 @@ def _create_spatialdata(
     ) -> "SpatialData":
         """Create SpatialData object from transcripts and boundaries."""
         import spatialdata
-        from spatialdata.models import PointsModel, ShapesModel
+        from spatialdata.models import PointsModel, ShapesModel, TableModel
         import geopandas as gpd
         import pandas as pd
         import dask.dataframe as dd
@@ -294,7 +297,7 @@ def _create_spatialdata(
                 **({"z": z_column} if has_z else {}),
             },
         )
-        elements[f"points/{self.points_key}"] = points
+        elements[self.points_key] = points
 
         # Shapes element (if boundaries provided or generated)
         if self.include_boundaries and self.boundary_method != "skip":
@@ -307,16 +310,17 @@ def _create_spatialdata(
             )
             if shapes is not None and len(shapes) > 0:
                 shapes_parsed = ShapesModel.parse(shapes)
-                elements[f"shapes/{self.shapes_key}"] = shapes_parsed
+                elements[self.shapes_key] = shapes_parsed
 
         # Create SpatialData
-        sdata = spatialdata.SpatialData.from_elements_dict(elements)
+        sdata = spatialdata.SpatialData.init_from_elements(elements)
 
         # Optional AnnData table
         if self.include_table:
-            region = self.shapes_key if f"shapes/{self.shapes_key}" in elements else None
+            region = self.shapes_key if self.shapes_key in elements else None
             region_key = self.table_region_key if region is not None else None
-            table = build_anndata_table(
+            instance_key = cell_id_column
+            table_data = build_anndata_table(
                 transcripts=transcripts,
                 cell_id_column=cell_id_column,
                 feature_column=feature_column,
@@ -327,6 +331,7 @@ def _create_spatialdata(
                 region=region,
                 region_key=region_key,
             )
+            table = TableModel.parse(table_data, region, region_key, instance_key, overwrite_metadata=True)
             sdata.tables[self.table_key] = table
 
         return sdata

src/segger/io/fields.py

@@ -23,7 +23,7 @@ class XeniumTranscriptFields:
     z: str = 'z_location'  # Z-coordinate (microns)
     feature: str = 'feature_name'
     cell_id: str = 'cell_id'
-    null_cell_id: str = 'UNASSIGNED'
+    null_cell_id: tuple = ('-1', 'UNASSIGNED')
     compartment: str = 'overlaps_nucleus'
     nucleus_value: int = 1
     quality: str = 'qv'  # Phred-scaled quality value (0-40)
@@ -35,6 +35,7 @@ class XeniumTranscriptFields:
         'DeprecatedCodeword_*',
         'UnassignedCodeword_*',
     ])
+    is_gene: str = 'is_gene'  # filter field in newer Xenium outputs
 
 @dataclass
 class XeniumBoundaryFields:

src/segger/io/preprocessor.py

@@ -469,6 +469,12 @@ def transcripts(self) -> pl.DataFrame:
             .with_row_index(name=std.row_index)
         )
 
+        # Force UTF8 if parquet stored it as Binary (Xenium 1.0.0)
+        if lf.collect_schema()[raw.feature] == pl.Binary:
+            lf = lf.with_columns(
+                pl.col(raw.feature).cast(pl.String)
+            )
+
         # Apply quality filtering using quality_filter module
         qf = XeniumQualityFilter()
         lf = qf.filter(
@@ -480,12 +486,11 @@ def transcripts(self) -> pl.DataFrame:
 
         # Standardize compartment labels
         lf = lf.with_columns(
-            pl.when(pl.col(raw.compartment) == raw.nucleus_value)
+            pl.when(pl.col(raw.compartment) ==  pl.lit(raw.nucleus_value, dtype=pl.UInt8))
             .then(std.nucleus_value)
             .when(
-                (pl.col(raw.compartment) != raw.nucleus_value) &
-                (pl.col(raw.cell_id) != raw.null_cell_id)
-            )
+                (pl.col(raw.compartment) !=  pl.lit(raw.nucleus_value, dtype=pl.UInt8)) &
+                ~pl.col(raw.cell_id).cast(pl.Utf8).is_in(raw.null_cell_id)            )
             .then(std.cytoplasmic_value)
             .otherwise(std.extracellular_value)
             .alias(std.compartment)
@@ -494,6 +499,7 @@ def transcripts(self) -> pl.DataFrame:
         # Standardize cell IDs
         lf = lf.with_columns(
             pl.col(raw.cell_id)
+            .cast(pl.Utf8)
             .replace(raw.null_cell_id, None)
             .alias(std.cell_id)
         )
@@ -574,10 +580,13 @@ def boundaries(self) -> gpd.GeoDataFrame:
             nuclei.reset_index(drop=False, names=std.id),
         ])
         # Convert index to string type (to join on AnnData)
-        bd.index = bd[std.id] + '_' + bd[std.boundary_type].map({
+        bd.index = bd[std.id].astype(str) + '_' + bd[std.boundary_type].map({
             std.nucleus_value: '0',
             std.cell_value: '1',
         })
+        # Convert std.id to string type (to join on AnnData)
+        if bd[std.id].dtype != "object":
+            bd[std.id] = bd[std.id].astype(str)
 
         return bd
 

src/segger/io/quality_filter.py

@@ -206,7 +206,10 @@ def filter(
         qv_col = quality_column or self.quality_column
 
         # Filter control probes
-        df = self._filter_by_patterns(df, feature_column)
+        if 'is_gene' in df.collect_schema():
+            df = df.filter(pl.col('is_gene'))  # use ad-hoc column (more robust)
+        else:
+            df = self._filter_by_patterns(df, feature_column)  # fallback to filter by patterns
 
         # Apply QV threshold
         if min_threshold is not None and min_threshold > 0: