fix_cli

gustaveroussy · Jun 10, 2024 · 68ba5ee · 68ba5ee
1 parent b1fb4e7
commit 68ba5ee
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 69 deletions.
diff --git a/sopa/cli/patchify.py b/sopa/cli/patchify.py
@@ -56,7 +56,7 @@ def baysor(
     config: str = typer.Option(
         default={},
         callback=ast.literal_eval,
-        help="Dictionnary of baysor parameters",
+        help="Dictionnary of baysor parameters, overwrite the config_path argument if provided",
     ),
     cell_key: str = typer.Option(
         None,
@@ -73,12 +73,19 @@ def baysor(
     ),
 ):
     """Prepare patches for transcript-based segmentation with baysor"""
-    return transcript_segmentation(
+    from sopa._constants import SopaKeys, SopaFiles
+    from .utils import _default_boundary_dir
+
+
+    if baysor_temp_dir is None:
+        baysor_temp_dir = _default_boundary_dir(sdata_path, SopaKeys.BAYSOR_BOUNDARIES)
+    return _transcript_segmentation(
         sdata_path=sdata_path,
-        method="baysor",
         patch_width_microns=patch_width_microns,
         patch_overlap_microns=patch_overlap_microns,
         temp_dir=baysor_temp_dir,
+        filename=SopaFiles.PATCHES_DIRS_BAYSOR,
+        config_name=SopaFiles.TOML_CONFIG_FILE,
         config_path=config_path,
         config=config,
         cell_key=cell_key,
@@ -94,7 +101,7 @@ def comseg(
     patch_overlap_microns: float = typer.Option(
         help="Number of overlapping microns between the patches. We advise to choose approximately twice the diameter of a cell"
     ),
-    baysor_temp_dir: str = typer.Option(
+    comseg_temp_dir: str = typer.Option(
         None,
         help="Temporary directory where baysor inputs and outputs will be saved. By default, uses `.sopa_cache/comseg_boundaries`",
     ),
@@ -105,7 +112,7 @@ def comseg(
     config: str = typer.Option(
         default={},
         callback=ast.literal_eval,
-        help="Dictionnary of ComSeg parameters",
+        help="Dictionnary of ComSeg parameters, overwrite the config_path argument if provided",
     ),
     cell_key: str = typer.Option(
         None,
@@ -118,13 +125,19 @@ def comseg(
     ),
 ):
     """Prepare patches for transcript-based segmentation with ComSeg"""
+    from sopa._constants import SopaKeys, SopaFiles
+    from .utils import _default_boundary_dir
 
-    return transcript_segmentation(
+    if comseg_temp_dir is None:
+        comseg_temp_dir = _default_boundary_dir(sdata_path, SopaKeys.COMSEG_BOUNDARIES)
+
+    return _transcript_segmentation(
         sdata_path=sdata_path,
-        method="comseg",
         patch_width_microns=patch_width_microns,
         patch_overlap_microns=patch_overlap_microns,
-        temp_dir=baysor_temp_dir,
+        temp_dir=comseg_temp_dir,
+        filename=SopaFiles.PATCHES_DIRS_COMSEG,
+        config_name=SopaFiles.JSON_CONFIG_FILE,
         config_path=config_path,
         config=config,
         cell_key=cell_key,
@@ -133,48 +146,35 @@ def comseg(
     )
 
 
-@app_patchify.command()
-def transcript_segmentation(
-    sdata_path: str = typer.Argument(help=SDATA_HELPER),
-    method: str = typer.Option(
-        "baysor",
-        help="Name of the method to use, choose in ['baysor', 'comseg']. for ComSeg, make sure to first run Cellpose or "
-        f"manually add the segmentation boundaries to the sdata.shapes as {SopaKeys.CELLPOSE_BOUNDARIES} key",
-    ),
-    patch_width_microns: float = typer.Option(help="Width (and height) of each patch in microns"),
-    patch_overlap_microns: float = typer.Option(
-        help="Number of overlapping microns between the patches. We advise to choose approximately twice the diameter of a cell"
-    ),
-    temp_dir: str = typer.Option(
-        None,
-        help="Temporary directory where baysor inputs and outputs will be saved. By default, uses `.sopa_cache/baysor_boundaries`",
-    ),
-    config_path: str = typer.Option(
-        None,
-        help="Path to the baysor config (you can also directly provide the argument via the `config` option)",
-    ),
-    config: str = typer.Option(
-        default={},
-        callback=ast.literal_eval,
-        help="Dictionnary of baysor parameters",
-    ),
-    cell_key: str = typer.Option(
-        None,
-        help="Optional column of the transcripts dataframe that indicates in which cell-id each transcript is, in order to use prior segmentation. "
-        f" Default is {SopaKeys.DEFAULT_CELL_KEY} if cell_key=None",
-    ),
-    unassigned_value: int = typer.Option(
-        None,
-        help="If --cell-key is provided, this is the value given to transcripts that are not inside any cell (if it's already 0, don't provide this argument)",
-    ),
-    use_prior: bool = typer.Option(
-        False,
-        help="Whether to use cellpose segmentation as a prior for baysor and comseg (if True, make sure to first run Cellpose or "
-        f"manually add the segmentation boundaries to the sdata.shapes as {SopaKeys.CELLPOSE_BOUNDARIES} key)",
-    ),
+def _transcript_segmentation(
+    sdata_path: str,
+    patch_width_microns: float,
+    patch_overlap_microns: float,
+    temp_dir: str,
+    filename: str,
+    config_name: str,
+    config_path: str,
+    config: str,
+    cell_key: str,
+    use_prior: bool,
+    unassigned_value: int,
 ):
-    """Prepare patches for transcript-based segmentation for the different available methods (baysor, comseg)"""
-    from sopa._constants import SopaFiles, SopaKeys
+    """Prepare patches for transcript-based segmentation for the different available methods (baysor, comseg)
+    args:
+        sdata_path (str) : Path to the SpatialData object
+        patch_width_microns (float) : Width (and height) of each patch in microns
+        patch_overlap_microns (str) : Number of overlapping microns between the patches. We advise to choose approximately twice the diameter of a cell
+        temp_dir (str) : Temporary directory where baysor inputs and outputs will be saved. By default, uses `.sopa_cache/baysor_boundaries`"
+        filename (str) : Name of the file to indicating the patch's index
+        config_name (str) : Name of the config file created for each patch
+        config_path (str): "Path to the config file (you can also directly provide the argument via the `config` option)"
+        config (str): "Dictionnary of parameters"
+        cell_key (str): "Optional column of the transcripts dataframe that indicates in which cell-id each transcript is, in order to use prior segmentation.
+        " Default is cell if cell_key=None"
+        use_prior (bool): "Whether to use cellpose segmentation as a prior for baysor and comseg (if True, make sure to first run Cellpose)"
+        unassigned_value (int): "If cell-key is provided, this is the value given to transcripts that are not inside any cell (if it's already 0, don't provide this argument)"
+
+    """
     from sopa._sdata import get_key
     from sopa.io.standardize import read_zarr_standardized, sanity_check
     from sopa.patches import Patches2D
@@ -187,23 +187,10 @@ def transcript_segmentation(
     assert (
         config or config_path is not None
     ), "Provide '--config-path', the path to a Baysor config file (toml) or comseg file (jsons)"
-    assert method in ["baysor", "comseg"], "method must be either 'baysor' or 'comseg'"
-
-    if temp_dir is None:
-        if method == "baysor":
-            temp_dir = _default_boundary_dir(sdata_path, SopaKeys.BAYSOR_BOUNDARIES)
-            filename = SopaFiles.PATCHES_DIRS_BAYSOR
-            config_name = SopaFiles.TOML_CONFIG_FILE
-        elif method == "comseg":
-            temp_dir = _default_boundary_dir(sdata_path, SopaKeys.COMSEG_BOUNDARIES)
-            filename = SopaFiles.PATCHES_DIRS_COMSEG
-            config_name = SopaFiles.JSON_CONFIG_FILE
-        else:
-            raise ValueError("method must be either 'baysor' or 'comseg'")
 
     df_key = get_key(sdata, "points")
     patches = Patches2D(sdata, df_key, patch_width_microns, patch_overlap_microns)
-    if method == "comseg":
+    if filename==SopaKeys.COMSEG_BOUNDARIES:
         patches.patchify_centroids(temp_dir)
         assert (
             use_prior

diff --git a/sopa/cli/segmentation.py b/sopa/cli/segmentation.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 import ast
-import logging
-from pathlib import Path
+
 
 import typer
 from tqdm import tqdm
@@ -11,7 +10,6 @@
 
 app_segmentation = typer.Typer()
 
-log = logging.getLogger(__name__)
 
 
 @app_segmentation.command()
@@ -188,19 +186,21 @@ def comseg(
     sdata_path: str = typer.Argument(help=SDATA_HELPER),
     patch_index: int = typer.Option(
         default=None,
-        help="Index of the patch on which the segmentation method should be run. NB: the number of patches is `len(sdata['sopa_patches'])`",
+        help="Index of the patch on which the segmentation method should be run.`",
     ),
     patch_dir: str = typer.Option(
         default=None,
-        help="Path to the temporary the segmentation method directory inside which we will store each individual patch segmentation. By default, saves into the `.sopa_cache/<method_name>` directory",
+        help="Path to the temporary the segmentation method directory inside which we will store each individual patch segmentation. By default, saves into the `.sopa_cache/comseg` directory",
     ),
 ):
     """Perform ComSeg segmentation. This can be done on all patches directly, or on one individual patch."""
     import json
 
     from sopa._constants import SopaFiles, SopaKeys
     from sopa.segmentation.methods import comseg_patch
-
+    import logging
+    log = logging.getLogger(__name__)
+    from pathlib import Path
     from .utils import _default_boundary_dir
 
     config_name = SopaFiles.JSON_CONFIG_FILE