Comseg fixes (#58)

Author: LouisK92

src/methods_transcript_assignment/comseg/config.vsh.yaml

@@ -1,3 +1,5 @@
+__merge__: /src/api/comp_method_transcript_assignment.yaml
+
 name: comseg
 label: "ComSeg Segmentation"
 summary: "Spatial segmentation using ComSeg method"
@@ -11,21 +13,16 @@ links:
 references:
   doi: "10.1038/s41592-020-01018-x"
 
-__merge__: /src/api/comp_method_segmentation.yaml
 
 arguments:
   - name: --transcripts_key
     type: string
     default: "transcripts"
     description: "Key for transcripts in the points layer"
-  - name: --shapes_key
-    type: string
-    default: "cell_boundaries"
-    description: "Key for cell boundaries in the shapes layer"
-  - name: --images_key
+  - name: --coordinate_system
     type: string
-    default: "morphology_mip"
-    description: "Key for morphology image in the images layer"
+    default: "global"
+    description: "Coordinate system for the transcripts"
   - name: --patch_width
     type: integer
     default: 1200

src/methods_transcript_assignment/comseg/script.py

@@ -1,18 +1,19 @@
+import dask
+import xarray as xr
 import spatialdata as sd
 import sopa
 import anndata as ad
 import pandas as pd
 import numpy as np
-from scipy import sparse
 
 ## VIASH START
 par = {
-    "input": "resources_test/task_ist_preprocessing/mouse_brain_combined/raw_ist.zarr",
-    "output": "transcripts.zarr",
-
+    "input_ist": "resources_test/task_ist_preprocessing/mouse_brain_combined/raw_ist.zarr",
+    "input_segmentation": "resources_test/task_ist_preprocessing/mouse_brain_combined/segmentation.zarr",
     "transcripts_key": "transcripts",
-    "shapes_key": "cell_boundaries",
-    "images_key": "morphology_mip",
+    "coordinate_system": "global",
+    "output": "temp/comseg/transcripts.zarr",
+
     "patch_width": 1200,
     "patch_overlap": 50,
     "transcript_patch_width": 200,
@@ -72,19 +73,36 @@ def _add_csr(X_partitions, geo_df, partition, gene_column, gene_names ):
     return adata
 
 
-# Read input SpatialData
-sdata = sd.read_zarr(par["input"])
+
+# Read input files
+print('Reading input files', flush=True)
+sdata = sd.read_zarr(par['input_ist'])
+sdata_segm = sd.read_zarr(par['input_segmentation'])
+
+
+# Convert the prior segmentation to polygons
+if isinstance(sdata_segm["segmentation"], xr.DataTree):
+    shapes_gdf = sopa.shapes.vectorize(sdata_segm["segmentation"]["scale0"].image)
+else:
+    shapes_gdf = sopa.shapes.vectorize(sdata_segm["segmentation"])
+
+sdata["segmentation_boundaries"] = sd.models.ShapesModel.parse(
+    shapes_gdf, transformations=sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True).copy()
+)
+
+# Make patches
 sopa.make_image_patches(sdata, patch_width=par["patch_width"], patch_overlap=par["patch_overlap"])
 
 transcript_patch_args = {
     "sdata": sdata,
     "write_cells_centroids": True,
     "patch_width": par["transcript_patch_width"],
+    "prior_shapes_key": "segmentation_boundaries",
 }
-transcript_patch_args["prior_shapes_key"] = par["shapes_key"]
 
 sopa.make_transcript_patches(**transcript_patch_args)
 
+# Run ComSeg
 config = {
     "dict_scale": {"x": 1, "y": 1, "z": 1},
     "mean_cell_diameter": par["mean_cell_diameter"],
@@ -96,21 +114,45 @@ def _add_csr(X_partitions, geo_df, partition, gene_column, gene_names ):
     "gene_column": par["gene_column"],
 }
 
-
 sopa.aggregation.transcripts._count_transcripts_aligned = fixed_count_transcripts_aligned
+# sopa.settings.parallelization_backend = 'dask'
 sopa.segmentation.comseg(sdata, config)
 
+# Assign transcripts to cell ids
+sopa.spatial.assign_transcript_to_cell(
+    sdata,
+    points_key="transcripts",
+    shapes_key="comseg_boundaries",
+    key_added="cell_id",
+    unassigned_value=0
+)
+
 # Create output SpatialData 
-sd_output = sd.SpatialData()
 
-cell_id_col = sdata["transcripts"][f"cell_id"]
-sdata.tables["table"]=ad.AnnData(obs=pd.DataFrame({"cell_id":cell_id_col}), var=sdata.tables["table"].var[[]])
-sdata_new = sd.SpatialData(
-    points=sdata.points,  
-    tables=sdata.tables   
-) 
+# Create objects for cells table
+print('Creating objects for cells table', flush=True)
+unique_cells = np.unique(sdata["transcripts"]["cell_id"])
+zero_idx = np.where(unique_cells == 0)
+if len(zero_idx[0]): 
+    unique_cells=np.delete(unique_cells, zero_idx[0][0])
+cell_id_col = pd.Series(unique_cells, name='cell_id', index=unique_cells)
+
+# Create transcripts only sdata
+print('Subsetting to transcripts cell id data', flush=True)
+sdata_transcripts_only = sd.SpatialData(
+    points={
+        "transcripts": sdata['transcripts']
+    },
+    tables={
+        "table": ad.AnnData(
+          obs=pd.DataFrame(cell_id_col),
+          var=sdata.tables["table"].var[[]]
+        )
+    }
+)
+
 
 output_path = par['output']
-sdata_new.write(output_path, overwrite=True)
+sdata_transcripts_only.write(output_path, overwrite=True)