Added mask flag to extract crops

scallops/cli/extract_crops.py

@@ -39,7 +39,7 @@ def single_crop(
     group: str,  # NOT USED
     file_list: list[str],
     metadata: dict,
-    labels_group: Group,
+    labels_group: Group | None,
     output_dir: str,
     output_sep: str,
     merge_dir: str,
@@ -69,14 +69,17 @@ def single_crop(
     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}")
-    zarr_labels = get_labels(
-        labels_group=labels_group,
-        name=image_key,
-        suffix=label_name,  # e.g. nuclei
-    )
+    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}.")
+        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,
@@ -98,7 +101,6 @@ def single_crop(
     )
     if len(merged_df) == 0:
         raise ValueError(f"No labels found for {image_key}.")
-    # e.g. CHAMMI-75
 
     if percentile_normalize is not None:
         if local_percentile_normalize:
@@ -135,7 +137,7 @@ def single_crop(
     image = da.map_blocks(img_as_ubyte, image)
     label_col = "label" if "label" in merged_df.columns else None
     merged_df = to_label_crops(
-        label_image=da.from_zarr(zarr_labels),
+        label_image=label_image,
         intensity_image=image,
         df=merged_df,
         label_col=label_col,

scallops/cli/extract_crops_main.py

@@ -46,8 +46,8 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
     crop_size = arguments.crop_size
     crop_size = (crop_size, crop_size)
     label_filter = arguments.label_filter
+    mask = arguments.mask
     percentile_min = arguments.percentile_min
-
     percentile_max = arguments.percentile_max
     output_format = arguments.output_format
     local_percentile_normalize = arguments.local_percentile_normalize
@@ -64,6 +64,8 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
             percentile_max = 100
         percentile_normalize = (percentile_min, percentile_max)
     gaussian_sigma = arguments.gaussian_sigma
+    if gaussian_sigma is not None and not mask:
+        raise ValueError("Please specify `mask` flag when providing `gaussian sigma`")
     label_name = arguments.label_name  # cell, cytosol, nuclei
     chunks = arguments.chunks
     if dask_server_url is None and arguments.dask_cluster is None:
@@ -79,9 +81,14 @@ def run_pipeline_extract_crops(arguments: argparse.Namespace):
 
     labels_path = arguments.labels
     no_version = arguments.no_version
-    assert labels_path is not None, "No labels provided"
-    label_root = zarr.open(labels_path, mode="r")
-    labels_group = label_root["labels"]
+    labels_group = None
+    if not mask:
+        labels_path = None
+    if labels_path 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(
@@ -133,20 +140,23 @@ def _create_parser(subparsers: argparse.ArgumentParser, default_help: bool) -> N
     required = parser.add_argument_group("required arguments")
     images_arg(required)
     output_dir_arg(required)
-    required.add_argument(
-        "--labels",
-        dest="labels",
-        required=True,
-        help="Path to zarr directory containing labels",
-    )
 
     image_pattern_arg(parser)
 
     required.add_argument(
         "--merge",
-        required=False,
         help="Path to directory containing output from `merge`",
     )
+    parser.add_argument(
+        "--labels",
+        dest="labels",
+        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",
+    )
     parser.add_argument(
         "--label-name",
         help="Name of labels to use. For example `nuclei` or `cell`",

scallops/io.py

@@ -954,7 +954,7 @@ def _extract_crops(
 
 def to_label_crops(
     intensity_image: zarr.Array | da.Array,
-    label_image: zarr.Array | da.Array,
+    label_image: zarr.Array | da.Array | None,
     df: pd.DataFrame,
     output_dir: str,
     crop_size: tuple[int, int] = (224, 224),
@@ -980,13 +980,14 @@ def to_label_crops(
     output_dir = output_dir.rstrip("/")
     is_dask_array = isinstance(intensity_image, da.Array)
     image_shape = intensity_image.shape[-2:]
-    label_shape = label_image.shape
-    assert label_shape == image_shape, "Label shape does not match image shape."
-    assert isinstance(label_image, da.Array) == is_dask_array, (
-        "Label image type does not match intensity image type."
-    )
-    if is_dask_array and intensity_image.chunksize[-2:] != label_image.chunks:
-        label_image = label_image.rechunk(intensity_image.chunksize[-2:])
+    if label_image is not None:
+        label_shape = label_image.shape
+        assert label_shape == image_shape, "Label shape does not match image shape."
+        assert isinstance(label_image, da.Array) == is_dask_array, (
+            "Label image type does not match intensity image type."
+        )
+        if is_dask_array and intensity_image.chunksize[-2:] != label_image.chunks:
+            label_image = label_image.rechunk(intensity_image.chunksize[-2:])
 
     if centroid_cols is None:
         centroid_cols = df.columns[
@@ -1038,7 +1039,8 @@ def to_label_crops(
 
     if not is_dask_array:
         intensity_image = delayed(intensity_image)
-        label_image = delayed(label_image)
+        if label_image is not None:
+            label_image = delayed(label_image)
 
     padding = 0 if gaussian_sigma is None else int(math.ceil(gaussian_sigma * 4))
 
@@ -1058,27 +1060,29 @@ def to_label_crops(
                 slice(
                     max(0, df_slice["crop-bbox-0"].min() - padding),
                     min(
-                        label_shape[0],
+                        image_shape[0],
                         df_slice["crop-bbox-2"].max() + padding,
                     ),
                 ),
                 slice(
                     max(0, df_slice["crop-bbox-1"].min() - padding),
                     min(
-                        label_shape[1],
+                        image_shape[1],
                         df_slice["crop-bbox-3"].max() + padding,
                     ),
                 ),
             )
             label_block = None
             if is_dask_array:
                 image_block = intensity_image[..., sl[0], sl[1]]
-                label_block = label_image[sl[0], sl[1]]
-
+                if label_image is not None:
+                    label_block = label_image[sl[0], sl[1]]
+            elif label_image is not None:
+                label_block = label_image
             results.append(
                 _extract_crops_delayed(
                     intensity_image=image_block if is_dask_array else intensity_image,
-                    label_image=label_block if is_dask_array else label_image,
+                    label_image=label_block,
                     output_dir=output_dir,
                     sl=sl,
                     label_col=label_col,
@@ -1277,8 +1281,10 @@ def _images2fov(
         if image is None:
             raise ValueError(f"{file_list[i]} could not be read.")
         if file_metadata is not None:
-            if "t" in file_metadata[i] and isinstance(
-                file_metadata[i]["t"], str
+            if (
+                "t" in file_metadata[i]
+                and isinstance(file_metadata[i]["t"], str)
+                and file_metadata[i]["t"].isdigit()
             ):  # convert to int
                 try:
                     # note we replace "_" with "-" so we don't convert "1_2" to 12 e.g.
@@ -1323,8 +1329,7 @@ def _images2fov(
             for dim in concat_dims
             if len(set([image.coords[dim].values[0] for image in images])) > 1
         ]
-        if len(concat_dims) > 0:
-            _match_size(images)
+
         if len(concat_dims) == 1:
             image = xr.concat(
                 images,

scallops/tests/test_features.py

@@ -33,8 +33,9 @@
 __data__ = __tests__.joinpath("data")
 
 
+@pytest.mark.parametrize("mask", [True, False])
 @pytest.mark.features
-def test_extract_crops_cmd(tmp_path, array_A1_102_cells, array_A1_102_alnpheno):
+def test_extract_crops_cmd(tmp_path, array_A1_102_cells, array_A1_102_alnpheno, mask):
     image = (
         array_A1_102_alnpheno.transpose(*("z", "c", "t", "y", "x")).rename(
             {"z": "t", "t": "z"}
@@ -76,6 +77,8 @@ def test_extract_crops_cmd(tmp_path, array_A1_102_cells, array_A1_102_alnpheno):
         "--output",
         crops_output_path,
     ]
+    if mask:
+        cmd.append("--mask")
 
     check_call(cmd)
     img = read_image(crops_output_path + "/cell/test/1523.tiff")
@@ -90,8 +93,9 @@ def test_to_label_crops(tmp_path, array_A1_102_cells, array_A1_102_alnpheno):
         .rename({"z": "t", "t": "z"})
         .isel(t=0, z=0)
     ).data  # ops swaps z and t in saved tif
-    output_dir_dask = str(tmp_path / "crops-dask")
-    output_dir_zarr = str(tmp_path / "crops-zarr")
+    output_dir_dask = str(tmp_path / "crops-mask-dask")
+    output_dir_zarr = str(tmp_path / "crops-mask-zarr")
+    output_dir_no_mask = str(tmp_path / "crops-dask")
 
     intensity_image = da.from_array(intensity_image).rechunk((-1, -1, 50, 50))
     objects_df = find_objects(label_image).compute()
@@ -127,16 +131,30 @@ def test_to_label_crops(tmp_path, array_A1_102_cells, array_A1_102_alnpheno):
     img_zarr = read_image(os.path.join(output_dir_zarr, "2603.tiff")).values.squeeze()
     np.testing.assert_array_equal(img_dask, img_zarr)
     # centroid: 1007.136364  579.090909
-    slice_2603 = intensity_image[
+    slice_2603_img = intensity_image[
         ..., slice(1007 - 15, 1007 + 15), slice(579 - 15, 579 + 15)
     ].compute()
     slice_2603_labels = label_image[
         ..., slice(1007 - 15, 1007 + 15), slice(579 - 15, 579 + 15)
     ].compute()
-    slice_2603 = slice_2603 * (slice_2603_labels == 2603)
+    slice_2603 = slice_2603_img * (slice_2603_labels == 2603)
     assert slice_2603.shape == (2, 30, 30)
     np.testing.assert_array_equal(img_dask, slice_2603, strict=True)
 
+    result_df_no_mask = to_label_crops(
+        intensity_image=intensity_image,
+        label_image=None,
+        df=objects_df.query("index==2603|index==17"),
+        crop_size=crop_size,
+        output_dir=output_dir_no_mask,
+    )
+    # 17 should be filtered b/c on tile edge
+    assert len(result_df_no_mask) == 1 and result_df_no_mask.index.values[0] == 2603
+    img_no_mask = read_image(
+        os.path.join(output_dir_no_mask, "2603.tiff")
+    ).values.squeeze()
+    np.testing.assert_array_equal(img_no_mask, slice_2603_img)
+
 
 def _label_features(
     intensity_image, label_image, features, label_filter=None, normalize=True