Workflow transform labels

scallops/cli/extract_crops.py

@@ -1,4 +1,5 @@
 import json
+import os
 from typing import Literal
 
 import dask.array as da
@@ -8,9 +9,13 @@
 import pyarrow.parquet as pq
 from array_api_compat import get_namespace
 from skimage.util import img_as_ubyte
-from zarr import Group
 
-from scallops.cli.features import _read_merged_or_objects, get_labels
+from scallops.cli.features import (
+    _find_labels,
+    _image_key_without_time_and_selected_time,
+    _read_merged_or_objects,
+)
+from scallops.cli.find_objects import get_path
 from scallops.cli.util import (
     _get_cli_logger,
     cli_metadata,
@@ -39,11 +44,10 @@ def single_crop(
     group: str,  # NOT USED
     file_list: list[str],
     metadata: dict,
-    labels_group: Group | None,
+    label_paths: list[str] | None,
     output_dir: str,
     output_sep: str,
-    merge_dir: str,
-    merge_dir_sep: str,
+    merge_paths: list[str],
     crop_size: tuple[int, int],
     output_format: Literal["tiff", "npy"],
     label_name: str,
@@ -57,38 +61,61 @@ def single_crop(
     force: bool,
 ):
     image_key = metadata["id"]
-
     output_dir = f"{output_dir}{output_sep}{label_name}{output_sep}{image_key}"
+    image_key_without_t, selected_timepoint = _image_key_without_time_and_selected_time(
+        metadata
+    )
 
-    output_parquet_path = f"{output_dir}.parquet"
+    output_parquet_path = get_path(
+        output_dir,
+        output_sep,
+        label_name,
+        image_key_without_t if image_key_without_t is not None else image_key,
+        selected_timepoint,
+        ".parquet",
+    )
     if not force and is_parquet_file(output_parquet_path):
         logger.info(f"Skipping features for {image_key} {label_name}")
         return
-
-    output_fs, _ = fsspec.core.url_to_fs(output_dir)
-    output_fs.makedirs(output_dir, exist_ok=True)
-    image = _images2fov(file_list, metadata, dask=True).squeeze().data
-    logger.info(f"{image_key} image shape {image.shape}")
-    label_image = None
-    if labels_group is not None:
-        zarr_labels = get_labels(
-            labels_group=labels_group,
-            name=image_key,
-            suffix=label_name,  # e.g. nuclei
-        )
-
-        if zarr_labels is None:
-            raise ValueError(f"Unable to read {label_name} labels for {image_key}.")
-        label_image = da.from_zarr(zarr_labels)
     merged_df = _read_merged_or_objects(
-        merge_dir=merge_dir,
-        merge_dir_sep=merge_dir_sep,
+        paths=merge_paths,
+        timepoint=selected_timepoint,
         label_name=label_name,
         image_key=image_key,
+        image_key_without_t=image_key_without_t,
         label_filter=label_filter,
+        add_timepoint_suffix=False,
     )
     if merged_df is None:
         raise ValueError(f"Unable to read merged data for {image_key}.")
+    output_fs, _ = fsspec.core.url_to_fs(output_dir)
+    output_fs.makedirs(output_dir, exist_ok=True)
+    image = _images2fov(file_list, metadata, dask=True).squeeze().data
+    logger.info(f"{image_key} image shape {image.shape}")
+    label_image = None
+    if label_paths is not None:
+        g, timepoints = _find_labels(
+            label_paths=label_paths,
+            image_key=image_key,
+            label_name=label_name,
+            image_key_without_t=image_key_without_t,
+            selected_timepoint=selected_timepoint,
+        )
+        if g is None:
+            raise ValueError(f"No labels found for {image_key}")
+
+        if len(timepoints) != 1:
+            raise ValueError(f"More than one timepoint found for {image_key}")
+        label_image = da.from_array(g[list(g.keys())[0]])
+        timepoint = timepoints[0]
+        if timepoint is not None and label_image.ndim == 3:
+            index = -1
+            for i in range(len(timepoints)):
+                if str(timepoints[i]) == str(timepoint):
+                    index = i
+                    break
+            label_image = label_image[index]
+
     n_labels_before_filtering = len(merged_df)
     if label_filter is not None:
         merged_df = merged_df.query(label_filter)
@@ -152,7 +179,6 @@ def single_crop(
     )
 
     output_metadata = cli_metadata() if not no_version else dict()
-
     table = pa.Table.from_pandas(merged_df, preserve_index=True)
     table = table.replace_schema_metadata(
         {
@@ -162,6 +188,7 @@ def single_crop(
     )
 
     fs, output_parquet_path = fsspec.url_to_fs(output_parquet_path)
+    fs.mkdirs(os.path.dirname(output_parquet_path.rstrip(fs.sep)), exist_ok=True)
     pq.write_table(
         table,
         output_parquet_path,

scallops/cli/extract_crops_main.py

@@ -1,7 +1,6 @@
 import argparse
 
 import fsspec
-import zarr
 from dask.bag import from_sequence
 
 from scallops.cli.arg_parser import _sort_groups
@@ -39,7 +38,7 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
 
     image_patterns = arguments.image_pattern
     output_dir = arguments.output
-    merge_dir = arguments.merge
+    merge_paths = arguments.merge
     subset = arguments.subset
     force = arguments.force
     groupby = arguments.groupby
@@ -71,24 +70,17 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
     if dask_server_url is None and arguments.dask_cluster is None:
         dask_cluster_parameters = _dask_workers_threads()
 
-    merge_dir_sep = None
-    if merge_dir is not None:
-        merge_dir_sep = fsspec.core.url_to_fs(merge_dir)[0].sep
-        merge_dir = merge_dir.rstrip(merge_dir_sep)
-
     output_fs, _ = fsspec.core.url_to_fs(output_dir)
     output_dir = output_dir.rstrip(output_fs.sep)
+    output_fs.mkdirs(output_dir, exist_ok=True)
 
-    labels_path = arguments.labels
+    label_paths = arguments.labels
     no_version = arguments.no_version
-    labels_group = None
+
     if not mask:
-        labels_path = None
-    if labels_path is None and mask:
+        label_paths = None
+    if label_paths is None and mask:
         raise ValueError("Labels must be provided when `mask` is true.")
-    if labels_path is not None:
-        label_root = zarr.open(labels_path, mode="r")
-        labels_group = label_root["labels"]
 
     image_seq = from_sequence(
         _set_up_experiment(
@@ -109,9 +101,8 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
             single_crop,
             output_dir=output_dir,
             output_sep=output_fs.sep,
-            merge_dir=merge_dir,
-            merge_dir_sep=merge_dir_sep,
-            labels_group=labels_group,
+            merge_paths=merge_paths,
+            label_paths=label_paths,
             label_filter=label_filter,
             label_name=label_name,
             percentile_normalize=percentile_normalize,
@@ -145,22 +136,24 @@ def _create_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> N
 
     required.add_argument(
         "--merge",
-        help="Path to directory containing output from `merge`",
+        nargs="+",
+        help="Path(s) to directory containing output from `find-objects`, `merge`, or `features`",
     )
     parser.add_argument(
         "--labels",
         dest="labels",
+        nargs="*",
         help="Path to zarr directory containing labels. Required when `mask` is true",
     )
     parser.add_argument(
         "--mask",
         action="store_true",
-        help="Set pixels not belonging to target label to zero",
+        help="Set crop pixels not belonging to target label to zero",
     )
     parser.add_argument(
         "--label-name",
         help="Name of labels to use. For example `nuclei` or `cell`",
-        default="cell",
+        default="nuclei",
     )
     parser.add_argument(
         "--crop-size",