diff --git a/tests/small_test_analysis/config/aggregate_combo.tsv b/tests/small_test_analysis/config/aggregate_combo.tsv index f5e30427..0c121bc5 100644 --- a/tests/small_test_analysis/config/aggregate_combo.tsv +++ b/tests/small_test_analysis/config/aggregate_combo.tsv @@ -1,3 +1,5 @@ -cell_class channel_combo plate well -Interphase DAPI_COXIV_CENPA_WGA 1 A1 -Interphase DAPI_COXIV_CENPA_WGA 1 A2 +cell_class channel_combo compartment_combo plate well +Interphase DAPI_COXIV_CENPA_WGA cell-nucleus-cytoplasm 1 A1 +Interphase DAPI_COXIV_CENPA_WGA cell-nucleus-cytoplasm 1 A2 +Interphase DAPI nucleus 1 A1 +Interphase DAPI nucleus 1 A2 diff --git a/tests/small_test_analysis/config/cluster_combo.tsv b/tests/small_test_analysis/config/cluster_combo.tsv index 574c4b8a..56b3fc18 100644 --- a/tests/small_test_analysis/config/cluster_combo.tsv +++ b/tests/small_test_analysis/config/cluster_combo.tsv @@ -1,2 +1,3 @@ -cell_class channel_combo leiden_resolution -Interphase DAPI_COXIV_CENPA_WGA 2 +cell_class channel_combo compartment_combo leiden_resolution +Interphase DAPI_COXIV_CENPA_WGA cell-nucleus-cytoplasm 2 +Interphase DAPI nucleus 2 diff --git a/tests/small_test_analysis/config/config.yml b/tests/small_test_analysis/config/config.yml index e89f362a..a0ee086b 100644 --- a/tests/small_test_analysis/config/config.yml +++ b/tests/small_test_analysis/config/config.yml @@ -96,6 +96,7 @@ phenotype: mode: null nuclei_diameter: 41.97813156768016 reconcile: contained_in_cells + second_obj_detection: false merge: approach: fast merge_combo_fp: config/merge_combo.tsv @@ -172,6 +173,10 @@ aggregate: bootstrap_combinations: - cell_class: Interphase channel_combo: DAPI_COXIV_CENPA_WGA + compartment_combo: cell-nucleus-cytoplasm + - cell_class: Interphase + channel_combo: DAPI + compartment_combo: nucleus cluster: cluster_combo_fp: config/cluster_combo.tsv uniprot_data_fp: config/benchmark_clusters/uniprot_data.tsv diff --git a/tests/test_compartment_filter.py b/tests/test_compartment_filter.py new file mode 100644 index 00000000..78af892a --- /dev/null +++ b/tests/test_compartment_filter.py @@ -0,0 +1,161 @@ +"""Unit tests for compartment filtering in the aggregate module.""" + +import sys +from pathlib import Path + +import pandas as pd +import pytest + +# Make the workflow's `lib` package importable without installing the package +WORKFLOW_LIB_DIR = Path(__file__).resolve().parents[1] / "workflow" +sys.path.insert(0, str(WORKFLOW_LIB_DIR)) + +from lib.aggregate.cell_data_utils import ( + COMPARTMENT_PREFIXES, + SECOND_OBJ_EXTRA_COLS, + compartment_combo_subset, +) + +ALL_COMPARTMENTS_4 = ["cell", "nucleus", "cytoplasm", "second_obj"] +ALL_COMPARTMENTS_3 = ["cell", "nucleus", "cytoplasm"] + + +def _make_features_4c(): + return pd.DataFrame( + { + "cell_DAPI_mean": [1.0], + "cell_area": [2.0], + "nucleus_DAPI_mean": [3.0], + "nucleus_area": [4.0], + "cytoplasm_DAPI_mean": [5.0], + "cytoplasm_area": [6.0], + "second_obj_DAPI_mean": [7.0], + "second_obj_area": [8.0], + "total_second_obj_area": [10.0], + "mean_second_obj_diameter": [11.0], + "mean_distance_to_nucleus": [12.0], + } + ) + + +def test_all_compartments_kept(): + df = _make_features_4c() + out = compartment_combo_subset(df, ALL_COMPARTMENTS_4, ALL_COMPARTMENTS_4) + assert list(out.columns) == list(df.columns) + + +def test_nucleus_only_drops_other_prefixes(): + df = _make_features_4c() + out = compartment_combo_subset(df, ["nucleus"], ALL_COMPARTMENTS_4) + assert set(out.columns) == {"nucleus_DAPI_mean", "nucleus_area"} + + +def test_cell_plus_nucleus(): + df = _make_features_4c() + out = compartment_combo_subset(df, ["cell", "nucleus"], ALL_COMPARTMENTS_4) + expected = {"cell_DAPI_mean", "cell_area", "nucleus_DAPI_mean", "nucleus_area"} + assert set(out.columns) == expected + + +def test_second_obj_exclusion_drops_extra_cols(): + df = _make_features_4c() + out = compartment_combo_subset(df, ALL_COMPARTMENTS_3, ALL_COMPARTMENTS_4) + for c in SECOND_OBJ_EXTRA_COLS: + assert c not in out.columns + assert "second_obj_DAPI_mean" not in out.columns + assert "cell_area" in out.columns + + +def test_second_obj_included_keeps_extra_cols(): + df = _make_features_4c() + out = compartment_combo_subset(df, ["second_obj"], ALL_COMPARTMENTS_4) + for c in SECOND_OBJ_EXTRA_COLS: + assert c in out.columns + + +def test_compartment_prefixes_constant_shape(): + assert set(COMPARTMENT_PREFIXES) == {"cell", "nucleus", "cytoplasm", "second_obj"} + for name, prefix in COMPARTMENT_PREFIXES.items(): + assert prefix == f"{name}_" + + +from lib.aggregate.cell_data_utils import resolve_aggregate_combos + + +def test_resolve_defaults_4_when_detection_true(): + out = resolve_aggregate_combos([{"channels": ["DAPI"]}], second_obj_detection=True) + assert out == [ + { + "channels": ["DAPI"], + "compartments": ["cell", "nucleus", "cytoplasm", "second_obj"], + } + ] + + +def test_resolve_defaults_3_when_detection_false(): + out = resolve_aggregate_combos([{"channels": ["DAPI"]}], second_obj_detection=False) + assert out == [ + {"channels": ["DAPI"], "compartments": ["cell", "nucleus", "cytoplasm"]} + ] + + +def test_resolve_rejects_unknown_compartment(): + with pytest.raises(ValueError, match="unknown compartment"): + resolve_aggregate_combos( + [{"channels": ["DAPI"], "compartments": ["nucleous"]}], + second_obj_detection=True, + ) + + +def test_resolve_rejects_second_obj_when_detection_false(): + with pytest.raises(ValueError, match="second_obj_detection"): + resolve_aggregate_combos( + [{"channels": ["DAPI"], "compartments": ["second_obj"]}], + second_obj_detection=False, + ) + + +def test_resolve_rejects_empty_compartments(): + with pytest.raises(ValueError, match="at least one compartment"): + resolve_aggregate_combos( + [{"channels": ["DAPI"], "compartments": []}], + second_obj_detection=True, + ) + + +def test_resolve_rejects_empty_channels(): + with pytest.raises(ValueError, match="at least one channel"): + resolve_aggregate_combos( + [{"channels": [], "compartments": ["cell"]}], + second_obj_detection=True, + ) + + +def test_resolve_dedupes_compartments_within_combo(): + out = resolve_aggregate_combos( + [{"channels": ["DAPI"], "compartments": ["cell", "cell", "nucleus"]}], + second_obj_detection=True, + ) + assert out[0]["compartments"] == ["cell", "nucleus"] + + +def test_resolve_dedupes_duplicate_combos(): + out = resolve_aggregate_combos( + [ + {"channels": ["DAPI"], "compartments": ["cell"]}, + {"channels": ["DAPI"], "compartments": ["cell"]}, + ], + second_obj_detection=True, + ) + assert len(out) == 1 + + +def test_resolve_preserves_distinct_combos(): + out = resolve_aggregate_combos( + [ + {"channels": ["DAPI"], "compartments": ["cell"]}, + {"channels": ["DAPI"], "compartments": ["nucleus"]}, + ], + second_obj_detection=True, + ) + assert len(out) == 2 diff --git a/workflow/lib/aggregate/aggregate.py b/workflow/lib/aggregate/aggregate.py index a334708c..07f29864 100644 --- a/workflow/lib/aggregate/aggregate.py +++ b/workflow/lib/aggregate/aggregate.py @@ -16,6 +16,7 @@ def aggregate( method="mean", ps_probability_threshold=None, ps_percentile_threshold=None, + carry_cols: list[str] | None = None, ) -> tuple[np.ndarray, pd.DataFrame]: """Apply mean or median aggregation to replicate embeddings and perturbation scores for each perturbation. @@ -31,6 +32,13 @@ def aggregate( Defaults to "mean". ps_probability_threshold (float, optional): Threshold for filtering based on perturbation score. ps_percentile_threshold (float, optional): Percentile threshold for filtering based on perturbation score. + carry_cols (list[str], optional): Metadata columns functionally determined by + `pert_col` to preserve (one value per group) in the output. Typical use: + when `pert_col` is a construct ID (e.g. `cell_barcode_0`), carry the + human-readable gene symbol (`gene_symbol_0`) through so downstream + clustering / benchmarking / lookups can match by gene name. Raises + ValueError if a carry_col is missing from metadata or has more than one + unique value within a group. Returns: tuple: @@ -48,6 +56,21 @@ def aggregate( if aggr_func is None: raise ValueError(f"Invalid aggregation method: {method}") + if carry_cols is None: + carry_cols = [] + else: + missing = [c for c in carry_cols if c not in metadata.columns] + if missing: + raise ValueError( + f"carry_cols not found in metadata: {missing}. " + f"Available columns: {list(metadata.columns)}" + ) + overlap = [c for c in carry_cols if c == pert_col] + if overlap: + raise ValueError( + f"carry_cols overlap with pert_col: {overlap}. Remove duplicates." + ) + # filter by ps_probability_threshold; keep NaNs if ps_probability_threshold is not None: mask = metadata["perturbation_score"].isna() | ( @@ -81,6 +104,16 @@ def aggregate( if "perturbation_auc" in metadata.columns: agg_meta["perturbation_auc"] = group["perturbation_auc"].iloc[0] + # Carry through columns that are functionally determined by pert_col. + for c in carry_cols: + nuniq = group[c].nunique(dropna=False) + if nuniq > 1: + raise ValueError( + f"carry_col {c!r} has {nuniq} unique values within group " + f"{pert_col}={pert!r}; not functionally determined by {pert_col}." + ) + agg_meta[c] = group[c].iloc[0] + if ps_probability_threshold is not None or ps_percentile_threshold is not None: # aggregate perturbation score with same function pert_score = ( diff --git a/workflow/lib/aggregate/cell_data_utils.py b/workflow/lib/aggregate/cell_data_utils.py index 02c63e96..fc277caa 100644 --- a/workflow/lib/aggregate/cell_data_utils.py +++ b/workflow/lib/aggregate/cell_data_utils.py @@ -180,3 +180,123 @@ def get_feature_table_cols(feature_cols): selected_columns.extend(overlap_cols) return selected_columns + + +COMPARTMENT_PREFIXES = { + "cell": "cell_", + "nucleus": "nucleus_", + "cytoplasm": "cytoplasm_", + "second_obj": "second_obj_", +} + +# Per-cell summary columns added by aggregate_second_obj_data that don't carry +# the second_obj_ prefix but are derived from second-object data. +SECOND_OBJ_EXTRA_COLS = frozenset( + { + "total_second_obj_area", + "mean_second_obj_diameter", + "mean_distance_to_nucleus", + } +) + + +def compartment_combo_subset( + features: pd.DataFrame, compartment_combo: str, all_compartments: list[str] +) -> pd.DataFrame: + """Filter features to keep only columns belonging to the requested compartments. + + Args: + features (pd.DataFrame): Feature columns (metadata already split out). + compartment_combo (list[str]): Compartments to keep, e.g. ["cell", "nucleus"]. + all_compartments (list[str]): Compartments present in the input. Used to + determine which prefixes to exclude. + + Returns: + pd.DataFrame: features without columns belonging to excluded compartments. + """ + excluded = [c for c in all_compartments if c not in compartment_combo] + excluded_prefixes = [COMPARTMENT_PREFIXES[c] for c in excluded] + + cols_to_drop = [ + col + for col in features.columns + if any(col.startswith(p) for p in excluded_prefixes) + ] + if "second_obj" in excluded: + cols_to_drop += [c for c in SECOND_OBJ_EXTRA_COLS if c in features.columns] + + return features.drop(columns=cols_to_drop) + + +def resolve_aggregate_combos( + aggregate_combos: list[dict], second_obj_detection: bool +) -> list[dict]: + """Normalize and validate AGGREGATE_COMBOS entries. + + For each entry: + - Fills in default compartments (all 4 if detection on, 3 otherwise) when missing. + - Dedupes compartments within the combo (preserving order). + - Validates compartment names, non-empty channels/compartments, and + second_obj-vs-detection consistency. + After normalization, duplicate (channels, compartments) pairs are deduped. + + Args: + aggregate_combos (list[dict]): Each dict has "channels" (list[str]) and + optionally "compartments" (list[str]). + second_obj_detection (bool): From config["phenotype"]["second_obj_detection"]. + + Returns: + list[dict]: Normalized, validated, and de-duplicated combos. + + Raises: + ValueError: On any validation failure. + """ + valid_compartments = set(COMPARTMENT_PREFIXES) + default_compartments = ( + ["cell", "nucleus", "cytoplasm", "second_obj"] + if second_obj_detection + else ["cell", "nucleus", "cytoplasm"] + ) + + resolved = [] + seen = set() + for idx, combo in enumerate(aggregate_combos): + channels = list(combo.get("channels") or []) + if not channels: + raise ValueError( + f"AGGREGATE_COMBOS[{idx}] must specify at least one channel" + ) + + comps = combo.get("compartments") + if comps is None: + comps = list(default_compartments) + else: + comps = list(comps) + if not comps: + raise ValueError( + f"AGGREGATE_COMBOS[{idx}] must specify at least one compartment" + ) + + # Dedupe compartments while preserving order + deduped = [] + for c in comps: + if c not in valid_compartments: + raise ValueError( + f"AGGREGATE_COMBOS[{idx}] has unknown compartment {c!r}; " + f"must be one of {sorted(valid_compartments)}" + ) + if c == "second_obj" and not second_obj_detection: + raise ValueError( + f"AGGREGATE_COMBOS[{idx}] lists 'second_obj' but " + f"config['phenotype']['second_obj_detection'] is False" + ) + if c not in deduped: + deduped.append(c) + + key = (tuple(channels), tuple(deduped)) + if key in seen: + continue + seen.add(key) + resolved.append({"channels": channels, "compartments": deduped}) + + return resolved diff --git a/workflow/lib/aggregate/eval_aggregate.py b/workflow/lib/aggregate/eval_aggregate.py index 4be14560..6e67ce22 100644 --- a/workflow/lib/aggregate/eval_aggregate.py +++ b/workflow/lib/aggregate/eval_aggregate.py @@ -108,6 +108,24 @@ def plot_feature_distributions( Returns: matplotlib.figure.Figure: Figure containing the violin plots. """ + # Defensive: empty feature selection (e.g. a compartment_combo whose + # prefixes don't match any *_mean column in the merge data) would pass + # nrows=0 to plt.subplots and raise. Return a placeholder figure so + # eval_aggregate can still write its other outputs and the DAG proceeds. + if not original_feature_cols or not aligned_feature_cols: + fig, ax = plt.subplots(figsize=(8, 2)) + ax.text( + 0.5, + 0.5, + "No intensity features available for this compartment_combo; " + "feature-distribution plot skipped.", + ha="center", + va="center", + wrap=True, + ) + ax.axis("off") + return fig + # Melt original features df_orig = original_cell_data[["plate", "well"] + original_feature_cols].melt( id_vars=["plate", "well"], var_name="Feature", value_name="Value" diff --git a/workflow/lib/aggregate/filter.py b/workflow/lib/aggregate/filter.py index 803204a1..0307be05 100644 --- a/workflow/lib/aggregate/filter.py +++ b/workflow/lib/aggregate/filter.py @@ -5,14 +5,121 @@ - perturbation_filter: Remove cells without perturbation assignments - missing_values_filter: Handle missing values through dropping or imputation - intensity_filter: Remove outliers based on channel intensities using LocalOutlierFactor +- harmonize_pool_schema: Compute a consistent feature-column set for a pool of per-well parquets """ import pandas as pd import numpy as np +import pyarrow.dataset as ds +import pyarrow.parquet as pq from sklearn.impute import KNNImputer from sklearn.neighbors import LocalOutlierFactor +def harmonize_pool_schema( + paths: list[str], + metadata_cols: list[str], + drop_cols_threshold: float | None = None, +) -> tuple[list[str], list[str], dict]: + """Compute a consistent (metadata, feature) column set for a multi-file pool. + + When per-well filter decisions diverge — e.g. missing_values_filter drops + different columns in different wells because class composition varies — naive + pool reads via pyarrow's dataset union produce NaN blocks that break + downstream operations (PCA, center-scale). This helper produces a single + harmonized column set to use across every scan of the pool, matching the + drop-column convention from missing_values_filter. + + Steps: + 1. Schema intersection across `paths`. Columns present in some but not all + files are "lost to per-well schema mismatch" (typically driven by + class-composition differences across wells) and logged per-file. + 2. If `drop_cols_threshold` is provided, scan the intersected pool once and + drop any feature column whose pool-level NaN proportion is >= threshold. + Uses the same comparison as missing_values_filter. + Row-level NaN cleanup (drop_rows_threshold, impute) is intentionally + deferred to the caller so it can be applied on each working subset. + + Args: + paths (list[str]): parquet file paths forming the pool. + metadata_cols (list[str]): canonical metadata column names (only those + present in every file are returned). + drop_cols_threshold (float | None): pool-level NaN proportion threshold + above which a feature column is dropped. None disables this step. + + Returns: + tuple: + kept_metadata_cols (list[str]), + kept_feature_cols (list[str]), + report (dict) — keys: "schema_mismatch" (dict file -> missing cols), + "threshold_dropped" (list of col names). + """ + if not paths: + return [], [], {"schema_mismatch": {}, "threshold_dropped": []} + + per_file_schemas = {p: set(pq.read_schema(p).names) for p in paths} + union_cols = set().union(*per_file_schemas.values()) + intersection_cols = set.intersection(*per_file_schemas.values()) + schema_mismatch = { + p: sorted(union_cols - cols) + for p, cols in per_file_schemas.items() + if union_cols - cols + } + + if schema_mismatch: + total_lost = len(union_cols - intersection_cols) + print( + f"[pool] dropping {total_lost} column(s) present in some but not all " + f"of {len(paths)} input files (per-file missingness — typically driven " + f"by class composition differences across wells):" + ) + for p, missing in schema_mismatch.items(): + preview = ", ".join(missing[:5]) + ( + f", ...(+{len(missing) - 5} more)" if len(missing) > 5 else "" + ) + print(f" {p.split('/')[-1]}: missing {len(missing)} col(s): {preview}") + + # Preserve caller-provided order within metadata_cols; feature order is sorted + # for determinism since there's no natural ordering after an intersection. + kept_metadata_cols = [c for c in metadata_cols if c in intersection_cols] + kept_feature_cols = sorted(intersection_cols - set(metadata_cols)) + + threshold_dropped = [] + if drop_cols_threshold is not None and kept_feature_cols: + # Pin to one reference schema so pyarrow casts each file to it instead of + # auto-unifying (which raises ArrowInvalid when files share a column with a + # mismatched dtype). paths are non-empty and share the kept columns. + pool = ds.dataset( + paths, format="parquet", schema=pq.read_schema(paths[0]) + ).to_table(columns=kept_feature_cols) + # Compute pool-level NaN proportion per column without materializing full pandas. + total_rows = pool.num_rows + if total_rows > 0: + for col_name in list(kept_feature_cols): + nulls = pool.column(col_name).null_count + if nulls / total_rows >= drop_cols_threshold: + threshold_dropped.append(col_name) + if threshold_dropped: + print( + f"[pool] dropping {len(threshold_dropped)} column(s) with pool-level " + f"NaN proportion >= {drop_cols_threshold * 100:.1f}%: " + f"{', '.join(threshold_dropped[:10])}" + f"{', ...' if len(threshold_dropped) > 10 else ''}" + ) + kept_feature_cols = [ + c for c in kept_feature_cols if c not in threshold_dropped + ] + + return ( + kept_metadata_cols, + kept_feature_cols, + { + "schema_mismatch": schema_mismatch, + "threshold_dropped": threshold_dropped, + }, + ) + + def query_filter(metadata, features, queries): """Sequentially apply a list of query strings to filter metadata and features DataFrames. @@ -149,42 +256,38 @@ def missing_values_filter( if pd.api.types.is_integer_dtype(features[col]): features[col] = features[col].astype("float64") - # Identify rows with any NAs in the remaining columns - has_na_mask = features[remaining_cols_with_na].isna().any(axis=1) - na_rows_idx = features.index[has_na_mask] - non_na_rows_idx = features.index[~has_na_mask] + # Use positional (iloc) indexing throughout — features.index may have + # duplicate labels (e.g. when concat across wells doesn't reset), + # which breaks .loc-based reads/writes. + has_na_mask = features[remaining_cols_with_na].isna().any(axis=1).to_numpy() + na_positions = np.flatnonzero(has_na_mask) + non_na_positions = np.flatnonzero(~has_na_mask) + col_positions = [ + features.columns.get_loc(c) for c in remaining_cols_with_na + ] np.random.seed(42) - # Process NA rows in batches - for i in range(0, len(na_rows_idx), batch_size): - batch_na_idx = na_rows_idx[i : i + batch_size] + for i in range(0, len(na_positions), batch_size): + batch_pos = na_positions[i : i + batch_size] print( - f"Imputing for batch {i // batch_size + 1} with {len(batch_na_idx)} NA rows" + f"Imputing for batch {i // batch_size + 1} with {len(batch_pos)} NA rows" ) - # Sample non-NA rows randomly instead of stratified sampling - sampled_non_na_idx = np.random.choice( - non_na_rows_idx, - size=min(sample_size, len(non_na_rows_idx)), + sampled_non_na_pos = np.random.choice( + non_na_positions, + size=min(sample_size, len(non_na_positions)), replace=False, ) + combined_pos = np.concatenate([batch_pos, sampled_non_na_pos]) - # Combine sampled non-NA rows with current batch of NA rows - batch_idx = np.concatenate([batch_na_idx, sampled_non_na_idx]) - - # Perform KNN imputation on this batch imputer = KNNImputer(n_neighbors=5) imputed_values = imputer.fit_transform( - features.loc[batch_idx, remaining_cols_with_na] + features.iloc[combined_pos, col_positions].to_numpy() ) - # Update only the NA rows with imputed values - na_rows_in_batch = np.arange(len(batch_na_idx)) - features.loc[batch_na_idx, remaining_cols_with_na] = pd.DataFrame( - imputed_values[na_rows_in_batch], - index=batch_na_idx, - columns=remaining_cols_with_na, - ) + features.iloc[batch_pos, col_positions] = imputed_values[ + : len(batch_pos) + ] return metadata, features diff --git a/workflow/lib/aggregate/perturbation_score.py b/workflow/lib/aggregate/perturbation_score.py index 0c5ae6af..b03de8e7 100644 --- a/workflow/lib/aggregate/perturbation_score.py +++ b/workflow/lib/aggregate/perturbation_score.py @@ -24,6 +24,9 @@ def perturbation_score( metadata_cols: list[str], perturbation_name_col: str, control_key: str, + perturbation_id_col: str | None = None, + control_name_col: str | None = None, + batch_cols: list[str] | None = None, minimum_cell_count: int = 100, n_jobs: int = -1, ) -> None: @@ -36,20 +39,32 @@ def perturbation_score( Args: cell_data (pd.DataFrame): DataFrame containing cell data that will be modified in-place. metadata_cols (list[str]): List of metadata column names that will be updated to include 'perturbation_score'. - perturbation_name_col (str): Column name containing perturbation identifiers. + perturbation_name_col (str): Column name containing perturbation identifiers (what `gene` is drawn from; e.g. "gene_symbol_0" or "cell_barcode_0"). control_key (str): Prefix identifying control perturbations (e.g., 'nontargeting'). + perturbation_id_col (str, optional): Column name for unique perturbation IDs + used by prepare_alignment_data. Defaults to perturbation_name_col. + control_name_col (str, optional): Column used to identify controls via + control_key. Defaults to perturbation_name_col. + batch_cols (list[str], optional): Columns defining the batch grouping. + Defaults to ["plate", "well"]. minimum_cell_count (int, optional): Minimum number of cells required to process a perturbation. Defaults to 100. n_jobs (int, optional): Number of parallel jobs. -1 uses all available CPUs. Defaults to -1. """ + if perturbation_id_col is None: + perturbation_id_col = perturbation_name_col + if control_name_col is None: + control_name_col = perturbation_name_col + if batch_cols is None: + batch_cols = ["plate", "well"] + perturbation_col = cell_data[perturbation_name_col] + control_col = cell_data[control_name_col] + # Non-control perturbations: NaN-safe comparison via astype(str), then drop any NaN-origin "nan" entries. + is_control = control_col.astype(str).str.startswith(control_key) perturbed_genes = [ - gene - for gene in perturbation_col.unique().tolist() - if not gene.startswith(control_key) + gene for gene in perturbation_col[~is_control].dropna().unique().tolist() ] - nt_idx = perturbation_col.index[ - perturbation_col.str.startswith(control_key) - ].to_numpy() + nt_idx = perturbation_col.index[is_control].to_numpy() print(f"Processing {len(perturbed_genes)} genes with {n_jobs} parallel jobs...") @@ -83,8 +98,12 @@ def perturbation_score( ) keep_idx = np.union1d(gene_idx, nt_keep) - # Extract subset - gene_subset_df = cell_data.iloc[keep_idx].copy() + # Extract subset. keep_idx holds pandas index labels (built from + # perturbation_col.index[...] and nt_idx above); use .loc so this + # works whether cell_data has a RangeIndex or a preserved label + # index (the latter happens when the caller passes a groupby slice + # without reset_index). + gene_subset_df = cell_data.loc[keep_idx].copy() original_idx = gene_subset_df.index.copy() gene_subset_df = gene_subset_df.reset_index(drop=True) @@ -98,6 +117,11 @@ def perturbation_score( gene_subset_df, original_idx, metadata_cols, + perturbation_name_col, + control_key, + perturbation_id_col, + control_name_col, + batch_cols, minimum_cell_count, ) for gene, gene_idx, gene_subset_df, original_idx in batch_data @@ -122,7 +146,7 @@ def calculate_perturbation_scores( cell_data: pd.DataFrame, gene: str, feature_cols: list[str], - perturbation_col: str = "gene_symbol_0", + perturbation_col: str, n_differential_features: int = 200, minimum_cell_count: int = 100, ) -> tuple[pd.Series, float]: @@ -138,7 +162,7 @@ def calculate_perturbation_scores( cell_data (pd.DataFrame): DataFrame containing cell data with features and metadata. gene (str): The target gene perturbation to score against. feature_cols (list[str]): List of feature column names to use for scoring. - perturbation_col (str, optional): Column name containing perturbation labels. Defaults to "gene_symbol_0". + perturbation_col (str): Column name containing perturbation labels used to build the binary target. n_differential_features (int, optional): Number of top differential features to select. Defaults to 200. minimum_cell_count (int, optional): Minimum number of cells required for scoring. Defaults to 200. @@ -149,7 +173,17 @@ def calculate_perturbation_scores( if cell_data.shape[0] < minimum_cell_count: return pd.Series(np.nan, index=cell_data.index), np.nan - y = (cell_data[perturbation_col] == gene).astype(int).to_numpy() + # NaN-safe binary target. Comparisons on nullable string dtypes can return + # pandas.NA for NaN rows; convert to plain bool via fillna(False) so + # astype(int) always yields strictly {0, 1}. + y = ( + cell_data[perturbation_col] + .eq(gene) + .fillna(False) + .astype(bool) + .astype(int) + .to_numpy() + ) X_all = cell_data[feature_cols].to_numpy() # select top-k differential features (ANOVA F-test) @@ -178,16 +212,26 @@ def _process_gene_subset( gene_subset_df: pd.DataFrame, original_idx: pd.Index, metadata_cols: list[str], + perturbation_name_col: str, + control_key: str, + perturbation_id_col: str, + control_name_col: str, + batch_cols: list[str], minimum_cell_count: int, ) -> tuple[str, np.ndarray, pd.Series, float] | None: """Process a pre-sliced gene subset and return perturbation scores. Args: - gene: Gene symbol being processed. - gene_idx: Original indices of gene cells in the full dataset. - gene_subset_df: Pre-sliced DataFrame with gene + control cells (reset index). + gene: Perturbation identifier being scored (drawn from perturbation_name_col). + gene_idx: Original indices of perturbation cells in the full dataset. + gene_subset_df: Pre-sliced DataFrame with perturbation + control cells (reset index). original_idx: Original indices before reset (for mapping scores back). metadata_cols: Metadata column names. + perturbation_name_col: Column naming each perturbation unit (what `gene` is drawn from). + control_key: Prefix identifying control rows in control_name_col. + perturbation_id_col: Column used as the unique perturbation ID in prepare_alignment_data. + control_name_col: Column used by centerscale to detect controls via control_key. + batch_cols: Columns defining the batch grouping. minimum_cell_count: Minimum cells required. Returns: @@ -203,19 +247,20 @@ def _process_gene_subset( metadata, features = prepare_alignment_data( metadata, features, - ["plate", "well"], - "gene_symbol_0", - "nontargeting", - "cell_barcode_0", + batch_cols, + perturbation_name_col, + control_key, + perturbation_id_col, ) features = features.astype(np.float32) features = centerscale_on_controls( features, metadata, - "gene_symbol_0", - "nontargeting", + perturbation_name_col, + control_key, "batch_values", + control_col=control_name_col, ) features = pd.DataFrame(features, columns=feature_cols) gene_subset_df = pd.concat([metadata, features], axis=1) @@ -225,7 +270,7 @@ def _process_gene_subset( gene_subset_df, gene, feature_cols, - perturbation_col="gene_symbol_0", + perturbation_col=perturbation_name_col, ) print( diff --git a/workflow/lib/aggregate/second_obj_utils.py b/workflow/lib/aggregate/second_obj_utils.py new file mode 100644 index 00000000..0653e21d --- /dev/null +++ b/workflow/lib/aggregate/second_obj_utils.py @@ -0,0 +1,215 @@ +"""Utilities for aggregating secondary object data into cell-level features. + +This module provides functions to merge cell-level data with per-object secondary +object data using configurable aggregation strategies. +""" + +import numpy as np +import pandas as pd +from typing import Literal, List + + +def aggregate_second_obj_data( + cells_df: pd.DataFrame, + second_objs_df: pd.DataFrame, + agg_strategy: Literal["none", "single", "all", "average"], +) -> pd.DataFrame: + """Aggregate secondary object data according to specified strategy. + + Merges cell-level data with secondary-object-level data using different + strategies to handle the one-to-many relationship between cells and + secondary objects. + + Parameters + ---------- + cells_df : pd.DataFrame + Cell-level data from final_merge. Must contain merge keys: + 'plate', 'well', 'tile', 'cell_0'. + second_objs_df : pd.DataFrame + Per-object secondary object data from merge_phenotype_second_objs. + Must contain merge keys: 'plate', 'well', 'tile', 'cell_id'. + agg_strategy : {"none", "single", "all", "average"} + Aggregation strategy: + - "none": Return cells_df unchanged + - "single": Add features only for cells with exactly 1 secondary object; + NaN for cells with 0 or 2+ objects (no rows are dropped) + - "all": Create numbered columns per object (feature_1, feature_2, etc.) + - "average": Mean of numeric features across all secondary objects per cell + + Returns: + ------- + pd.DataFrame + Cell-level data with secondary object features merged according to strategy. + All cells from cells_df are always preserved. + + Raises: + ------ + ValueError + If required merge keys are missing or strategy is invalid. + """ + valid_strategies = ["none", "single", "all", "average"] + if agg_strategy not in valid_strategies: + raise ValueError( + f"Unknown strategy: {agg_strategy}. Must be one of {valid_strategies}" + ) + + if agg_strategy == "none": + return cells_df.copy() + + cells_merge_keys = ["plate", "well", "tile", "cell_0"] + second_objs_merge_keys = ["plate", "well", "tile", "cell_id"] + + _validate_merge_keys(cells_df, cells_merge_keys, "cells") + _validate_merge_keys(second_objs_df, second_objs_merge_keys, "second_objs") + + strategy_functions = { + "single": _aggregate_single, + "all": _aggregate_all, + "average": _aggregate_average, + } + + return strategy_functions[agg_strategy]( + cells_df.copy(), second_objs_df, cells_merge_keys, second_objs_merge_keys + ) + + +def _validate_merge_keys( + df: pd.DataFrame, required_keys: List[str], df_name: str +) -> None: + """Validate that required merge keys exist in dataframe.""" + missing_keys = [key for key in required_keys if key not in df.columns] + if missing_keys: + raise ValueError(f"Missing merge keys {missing_keys} in {df_name} dataframe") + + +def _get_feature_cols(second_objs_df: pd.DataFrame, merge_keys: List[str]) -> List[str]: + """Get secondary object feature columns (everything except merge keys and second_obj_id).""" + exclude = set(merge_keys) | {"second_obj_id"} + return [col for col in second_objs_df.columns if col not in exclude] + + +def _prepare_second_objs( + second_objs_df: pd.DataFrame, + second_objs_merge_keys: List[str], + cells_merge_keys: List[str], +) -> pd.DataFrame: + """Rename cell_id to cell_0 for merging with cells dataframe.""" + df = second_objs_df.copy() + df = df.rename(columns={"cell_id": "cell_0"}) + return df + + +def _aggregate_single( + cells_df: pd.DataFrame, + second_objs_df: pd.DataFrame, + cells_merge_keys: List[str], + second_objs_merge_keys: List[str], +) -> pd.DataFrame: + """Single strategy: populate features only for cells with exactly 1 secondary object. + + All cells are preserved. Cells with 0 or 2+ secondary objects get NaN for + all secondary object feature columns. + """ + second_objs = _prepare_second_objs( + second_objs_df, second_objs_merge_keys, cells_merge_keys + ) + feature_cols = _get_feature_cols(second_objs_df, second_objs_merge_keys) + + # Count secondary objects per cell + obj_counts = ( + second_objs.groupby(cells_merge_keys).size().reset_index(name="_obj_count") + ) + + # Tag cells with their object count + cells_with_count = cells_df.merge(obj_counts, on=cells_merge_keys, how="left") + cells_with_count["_obj_count"] = cells_with_count["_obj_count"].fillna(0) + + # Split into single-object and other cells + single_mask = cells_with_count["_obj_count"] == 1 + single_cells = cells_with_count[single_mask].drop(columns=["_obj_count"]) + other_cells = cells_with_count[~single_mask].drop(columns=["_obj_count"]) + + # Merge features for single-object cells + merged_single = single_cells.merge( + second_objs[cells_merge_keys + feature_cols], + on=cells_merge_keys, + how="left", + ) + + # Add NaN columns for other cells + if feature_cols: + nan_df = pd.DataFrame(np.nan, index=other_cells.index, columns=feature_cols) + other_cells = pd.concat([other_cells, nan_df], axis=1) + + return pd.concat([merged_single, other_cells], ignore_index=True) + + +def _aggregate_all( + cells_df: pd.DataFrame, + second_objs_df: pd.DataFrame, + cells_merge_keys: List[str], + second_objs_merge_keys: List[str], +) -> pd.DataFrame: + """All strategy: create numbered columns for each secondary object. + + Creates feature_1, feature_2, etc. columns. Cells with fewer objects than + the maximum get NaN for the missing object columns. + """ + second_objs = _prepare_second_objs( + second_objs_df, second_objs_merge_keys, cells_merge_keys + ) + feature_cols = _get_feature_cols(second_objs_df, second_objs_merge_keys) + + # Number secondary objects within each cell + second_objs["_obj_num"] = second_objs.groupby(cells_merge_keys).cumcount() + 1 + + max_objs = second_objs["_obj_num"].max() if len(second_objs) > 0 else 0 + + result_df = cells_df.copy() + + for obj_num in range(1, max_objs + 1): + current = second_objs[second_objs["_obj_num"] == obj_num].copy() + col_mapping = {col: f"{col}_{obj_num}" for col in feature_cols} + current = current.rename(columns=col_mapping) + + merge_cols = cells_merge_keys + list(col_mapping.values()) + result_df = result_df.merge( + current[merge_cols], on=cells_merge_keys, how="left" + ) + + return result_df + + +def _aggregate_average( + cells_df: pd.DataFrame, + second_objs_df: pd.DataFrame, + cells_merge_keys: List[str], + second_objs_merge_keys: List[str], +) -> pd.DataFrame: + """Average strategy: mean of numeric features across all secondary objects per cell. + + Non-numeric columns use the first value. Cells with no secondary objects get NaN. + """ + second_objs = _prepare_second_objs( + second_objs_df, second_objs_merge_keys, cells_merge_keys + ) + feature_cols = _get_feature_cols(second_objs_df, second_objs_merge_keys) + + if not feature_cols: + return cells_df.copy() + + # Identify numeric vs non-numeric feature columns + numeric_cols = ( + second_objs[feature_cols].select_dtypes(include=[np.number]).columns.tolist() + ) + non_numeric_cols = [col for col in feature_cols if col not in numeric_cols] + + agg_dict = {} + for col in numeric_cols: + agg_dict[col] = "mean" + for col in non_numeric_cols: + agg_dict[col] = "first" + + averaged = second_objs.groupby(cells_merge_keys).agg(agg_dict).reset_index() + + return cells_df.merge(averaged, on=cells_merge_keys, how="left") diff --git a/workflow/lib/cluster/cluster_eval.py b/workflow/lib/cluster/cluster_eval.py index 525548f5..d8b5edc1 100644 --- a/workflow/lib/cluster/cluster_eval.py +++ b/workflow/lib/cluster/cluster_eval.py @@ -19,6 +19,7 @@ def find_optimal_resolution( root_fp, channel_combo, + compartment_combo, cell_class, use_filtered=False, metric="balanced", @@ -34,6 +35,7 @@ def find_optimal_resolution( Args: root_fp (Path): Root output directory (config["all"]["root_fp"]). channel_combo (str): Channel combination name. + compartment_combo (str): Compartment combination name. cell_class (str): Cell class name. use_filtered (bool): Whether to look in filtered/ subdirectory. metric (str): Metric to optimize. Options: @@ -65,9 +67,16 @@ def find_optimal_resolution( # Build base cluster path if use_filtered: - base_path = root_fp / "cluster" / channel_combo / cell_class / "filtered" + base_path = ( + root_fp + / "cluster" + / channel_combo + / compartment_combo + / cell_class + / "filtered" + ) else: - base_path = root_fp / "cluster" / channel_combo / cell_class + base_path = root_fp / "cluster" / channel_combo / compartment_combo / cell_class # Auto-discover resolutions if not provided if resolutions is None: @@ -430,6 +439,7 @@ def analyze_all_resolutions( root_fp, cell_classes, channel_combos, + compartment_combos, use_filtered=False, metric="balanced", ideal_size_range=(15, 25), @@ -438,7 +448,7 @@ def analyze_all_resolutions( top_n_table=10, verbose=True, ): - """Analyze optimal resolutions for all cell class/channel combinations. + """Analyze optimal resolutions for all cell class/channel/compartment combinations. Convenience wrapper that finds optimal resolutions, displays results, and returns a summary for all combinations. @@ -447,6 +457,7 @@ def analyze_all_resolutions( root_fp (Path): Root output directory (config["all"]["root_fp"]). cell_classes (list): List of cell class names (e.g., ["Interphase", "Mitotic"]). channel_combos (list): List of channel combinations. + compartment_combos (list): List of compartment combinations. use_filtered (bool): Whether to look in filtered/ subdirectory. metric (str): Metric to optimize ("balanced", "combined", etc.). ideal_size_range (tuple): (min, max) target cluster size. @@ -467,56 +478,58 @@ def analyze_all_resolutions( for cell_class in cell_classes: for channel_combo in channel_combos: - try: - result = find_optimal_resolution( - root_fp=root_fp, - channel_combo=channel_combo, - cell_class=cell_class, - use_filtered=use_filtered, - metric=metric, - ideal_size_range=ideal_size_range, - size_metric=size_metric, - ) - key = f"{cell_class}_{channel_combo}" - optimal_resolutions[key] = result - - if verbose: - print(f"\n{'=' * 80}") - print(f"{cell_class} / {channel_combo}") - print(f"{'=' * 80}") - print(f" Optimal resolution: {result['optimal_resolution']}") - print(f" Optimization metric: {result['metric_used']}") - print(f" Size metric used: {result['size_metric_used']}") - print(f" Target size range: {result['ideal_size_range']}") - - # Show formatted decision table - print( - f"\nTop {top_n_table} resolutions (sorted by {metric} score):" - ) - table = format_resolution_table( - result["all_results"], top_n=top_n_table - ) - display(table) - - if show_plots: - # Show metrics comparison plot - fig = plot_resolution_comparison( - result["all_results"], metric=metric - ) - plt.suptitle( - f"{cell_class} / {channel_combo}", - fontsize=14, - fontweight="bold", - ) - plt.tight_layout() - plt.show() - - except Exception as e: - if verbose: - print( - f"\n{cell_class} / {channel_combo}: No benchmark results found" + for compartment_combo in compartment_combos: + try: + result = find_optimal_resolution( + root_fp=root_fp, + channel_combo=channel_combo, + compartment_combo=compartment_combo, + cell_class=cell_class, + use_filtered=use_filtered, + metric=metric, + ideal_size_range=ideal_size_range, + size_metric=size_metric, ) - print(f" Error: {e}") + key = f"{cell_class}_{channel_combo}_{compartment_combo}" + optimal_resolutions[key] = result + + if verbose: + print(f"\n{'=' * 80}") + print(f"{cell_class} / {channel_combo} / {compartment_combo}") + print(f"{'=' * 80}") + print(f" Optimal resolution: {result['optimal_resolution']}") + print(f" Optimization metric: {result['metric_used']}") + print(f" Size metric used: {result['size_metric_used']}") + print(f" Target size range: {result['ideal_size_range']}") + + # Show formatted decision table + print( + f"\nTop {top_n_table} resolutions (sorted by {metric} score):" + ) + table = format_resolution_table( + result["all_results"], top_n=top_n_table + ) + display(table) + + if show_plots: + # Show metrics comparison plot + fig = plot_resolution_comparison( + result["all_results"], metric=metric + ) + plt.suptitle( + f"{cell_class} / {channel_combo} / {compartment_combo}", + fontsize=14, + fontweight="bold", + ) + plt.tight_layout() + plt.show() + + except Exception as e: + if verbose: + print( + f"\n{cell_class} / {channel_combo} / {compartment_combo}: No benchmark results found" + ) + print(f" Error: {e}") # Generate summary table if verbose: diff --git a/workflow/lib/phenotype/align_channels.py b/workflow/lib/phenotype/align_channels.py index eea1c008..f32ab08f 100644 --- a/workflow/lib/phenotype/align_channels.py +++ b/workflow/lib/phenotype/align_channels.py @@ -1,13 +1,15 @@ """Module for aligning channels in phenotype. -Uses NumPy and scikit-image to provide image -alignment between sequencing cycles, apply percentile-based filtering, fill masked -areas with noise, and perform various transformations to enhance image data quality. +Uses NumPy and scikit-image to provide image alignment between sequencing cycles. """ import numpy as np from lib.shared.image_utils import remove_channels -from lib.shared.align import apply_window, calculate_offsets, apply_offsets +from lib.shared.align import ( + apply_window, + calculate_offsets, + apply_offsets, +) def align_phenotype_channels( @@ -19,6 +21,7 @@ def align_phenotype_channels( window=2, remove_channel=False, verbose=False, + return_metrics=False, ): """Rigid alignment of phenotype channels based on target and source channels. @@ -38,9 +41,14 @@ def align_phenotype_channels( verbose (bool, optional): If True, print detailed alignment information including calculated offsets for source and rider channels. Useful for debugging alignment issues. Defaults to False. + return_metrics (bool, optional): If True, return alignment quality metrics in addition + to aligned data. Defaults to False. Returns: np.ndarray: Phenotype data aligned across specified channels. + If return_metrics=True, returns tuple of (aligned_data, metrics_dict) where + metrics_dict contains: + - 'offset': list, the [y, x] offset that was applied """ # Handle stacked vs unstacked data if image_data.ndim == 4: @@ -50,24 +58,28 @@ def align_phenotype_channels( data_ = image_data.copy() stack = False - # Calculate alignment offsets + # Calculate alignment offsets using phase cross-correlation windowed = apply_window(data_[[target, source]], window) - offsets = calculate_offsets(windowed, upsample_factor=upsample_factor) + offsets, _ = calculate_offsets(windowed, upsample_factor=upsample_factor) + + final_offset = offsets[1] # Handle riders and create full offsets array if not isinstance(riders, list): riders = [riders] full_offsets = np.zeros((data_.shape[0], 2)) - full_offsets[[source] + riders] = offsets[1] + full_offsets[[source] + riders] = final_offset if verbose: print("\n=== Phenotype Channel Alignment Offsets ===") print(f" Target channel (index {target}): no shift (reference)") - print(f" Source channel (index {source}): shift = {offsets[1]} pixels (y, x)") + print( + f" Source channel (index {source}): shift = {final_offset} pixels (y, x)" + ) if riders: for rider_idx in riders: print( - f" Rider channel (index {rider_idx}): shift = {offsets[1]} pixels (y, x)" + f" Rider channel (index {rider_idx}): shift = {final_offset} pixels (y, x)" ) # Apply alignment @@ -88,6 +100,15 @@ def align_phenotype_channels( elif remove_channel == "source": aligned = remove_channels(aligned, source) + # Return with metrics if requested + if return_metrics: + metrics_dict = { + "offset": final_offset.tolist() + if hasattr(final_offset, "tolist") + else list(final_offset), + } + return aligned, metrics_dict + return aligned diff --git a/workflow/lib/phenotype/extract_phenotype_second_objs.py b/workflow/lib/phenotype/extract_phenotype_second_objs.py new file mode 100644 index 00000000..64f213c6 --- /dev/null +++ b/workflow/lib/phenotype/extract_phenotype_second_objs.py @@ -0,0 +1,384 @@ +"""Helper function to extract phenotype features from CellProfiler-like data for secondary objects.""" + +from itertools import combinations, permutations, product + +import numpy as np +import pandas as pd +import skimage.measure +import skimage.morphology +import skimage.filters +import skimage.feature +import skimage.segmentation +from scipy import ndimage as ndi + +from lib.external.cp_emulator import ( + grayscale_features_multichannel, + correlation_features_multichannel, + shape_features, + grayscale_columns_multichannel, + correlation_columns_multichannel, + shape_columns, + neighbor_measurements, +) +from lib.shared.feature_extraction import extract_features, extract_features_bare +from lib.shared.log_filter import log_ndi +from lib.phenotype.constants import DEFAULT_METADATA_COLS + + +def extract_phenotype_second_objs( + data_phenotype, + second_objs, + wildcards, + second_obj_cell_mapping_df=None, + second_obj_channels="all", + foci_channel=None, + channel_names=["dapi", "tubulin", "gh2ax", "phalloidin"], +): + """Extract phenotype features for secondary objects with multi-channel functionality. + + Updated version with proper column ordering matching cp_multichannel. + + Args: + data_phenotype (numpy.ndarray): Phenotype data array of shape (..., CHANNELS, I, J). + second_objs (numpy.ndarray): Secondary object segmentation mask with unique integers for each object. + second_obj_cell_mapping_df (pandas.DataFrame): DataFrame containing the mapping between secondary objects and cells. + wildcards (dict): Dictionary containing wildcards. + second_obj_channels (str or list): List of channel indices to consider for secondary object analysis or 'all'. + foci_channel (int, optional): Index of the channel containing foci information. + channel_names (list): List of channel names. + + Returns: + pandas.DataFrame: DataFrame containing extracted phenotype features for each secondary object. + """ + # If secondary objects are empty, return an empty DataFrame + if np.sum(second_objs) == 0: + print("No secondary objects found for feature extraction.") + return pd.DataFrame(columns=["second_obj_id", "cell_id"]) + + # Check if all channels should be used + if second_obj_channels == "all": + try: + second_obj_channels = list(range(data_phenotype.shape[-3])) + except: + second_obj_channels = [0] + + dfs = [] + + # Define features + features = grayscale_features_multichannel.copy() + features.update(correlation_features_multichannel) + features.update(shape_features) + + # Define function to create column map + def make_column_map(channels): + columns = {} + # Create columns for grayscale features + for feat, out in grayscale_columns_multichannel.items(): + columns.update( + { + f"{feat}_{n}": f"{channel_names[ch]}_{renamed}" + for n, (renamed, ch) in enumerate(product(out, channels)) + } + ) + # Create columns for correlation features + for feat, out in correlation_columns_multichannel.items(): + if feat == "lstsq_slope": + iterator = permutations + else: + iterator = combinations + columns.update( + { + f"{feat}_{n}": renamed.format( + first=channel_names[first], second=channel_names[second] + ) + for n, (renamed, (first, second)) in enumerate( + product(out, iterator(channels, 2)) + ) + } + ) + # Add shape columns + columns.update(shape_columns) + return columns + + # Create column map for secondary objects + second_obj_columns = make_column_map(second_obj_channels) + + # Extract secondary object features for all channels + dfs.append( + extract_features( + data_phenotype[..., second_obj_channels, :, :], + second_objs, + dict(), # Pass empty dict instead of wildcards here + features, + multichannel=True, + ) + .rename(columns=second_obj_columns) + .set_index("label") + .add_prefix("second_obj_") + ) + + # Extract foci features within secondary objects if foci channel is provided + if foci_channel is not None: + foci = find_foci_in_second_objs( + data_phenotype[..., foci_channel, :, :], + second_objs, + remove_border_foci=True, + ) + + if foci is not None: + dfs.append( + extract_features_bare(foci, second_objs, features=foci_features) + .set_index("label") + .add_prefix(f"second_obj_{channel_names[foci_channel]}_") + ) + + # Extract secondary object neighbor measurements + dfs.append( + neighbor_measurements(second_objs, distances=[1]) + .set_index("label") + .add_prefix("second_obj_") + ) + + # Concatenate secondary object features + second_obj_features = pd.concat(dfs, axis=1, join="outer", sort=False).reset_index() + + # Combine with second_obj_cell_mapping_df + if second_obj_cell_mapping_df is not None: + second_obj_df = pd.merge( + second_obj_cell_mapping_df, + second_obj_features.rename(columns={"label": "second_obj_id"}), + on="second_obj_id", + how="left", + suffixes=("_map", "_feat"), # left, right + ) + + # If both exist, make a single second_obj_area column (prefer features) + # If other features are present in both dataframes, modify the next next four lines to reflect the column names and add _map and _feat suffixes + if {"second_obj_area_map", "second_obj_area_feat"} <= set( + second_obj_df.columns + ): + second_obj_df["second_obj_area"] = second_obj_df[ + "second_obj_area_feat" + ].combine_first(second_obj_df["second_obj_area_map"]) + second_obj_df = second_obj_df.drop( + columns=["second_obj_area_map", "second_obj_area_feat"] + ) + else: + second_obj_df = second_obj_features.rename(columns={"label": "second_obj_id"}) + + # Add wildcards metadata at the END (they'll be reordered later) + for k, v in sorted(wildcards.items()): + second_obj_df[k] = v + + # Apply column ordering + second_obj_df = order_dataframe_columns_second_objs(second_obj_df) + + return second_obj_df + + +def order_dataframe_columns_second_objs( + df, metadata_cols=None, label_cols=["second_obj_id", "cell_id"] +): + """Reorder DataFrame columns to put metadata first, then features for secondary objects. + + Args: + df (pandas.DataFrame): DataFrame to reorder + metadata_cols (list): List of metadata column names to put first + label_cols (list): Names of the label columns (second_obj_id, cell_id) + + Returns: + pandas.DataFrame: DataFrame with reordered columns + """ + if metadata_cols is None: + metadata_cols = DEFAULT_METADATA_COLS + + # Start with label columns + ordered_cols = [] + for col in label_cols: + if col in df.columns: + ordered_cols.append(col) + + # Add metadata columns that exist in the DataFrame + for col in metadata_cols: + if col in df.columns and col not in ordered_cols: + ordered_cols.append(col) + + # Categorize remaining feature columns + remaining_cols = [col for col in df.columns if col not in ordered_cols] + + # Group features by type + second_obj_features = [ + col for col in remaining_cols if col.startswith("second_obj_") + ] + + # Add any other columns that don't fit the above patterns + other_features = [ + col for col in remaining_cols if not col.startswith("second_obj_") + ] + + # Combine in desired order + ordered_cols.extend(other_features) # Any additional metadata/wildcards + ordered_cols.extend(second_obj_features) + + return df[ordered_cols] + + +def find_foci_in_second_objs( + data, second_objs, radius=3, threshold=10, remove_border_foci=False +): + """Detect foci within secondary objects using a white tophat filter and other processing steps. + + Args: + data (numpy.ndarray): Input image data. + second_objs (numpy.ndarray): Secondary object segmentation mask. + radius (int, optional): Radius of the disk used in the white tophat filter. Default is 3. + threshold (float, optional): Threshold value for identifying foci in the processed image. Default is 10. + remove_border_foci (bool, optional): Flag to remove foci touching the secondary object border. Default is False. + + Returns: + labeled (numpy.ndarray): Labeled segmentation mask of foci within secondary objects. + """ + # If no secondary objects, return None + if np.sum(second_objs) == 0: + return None + + # Create a binary mask for all secondary objects + second_obj_mask = second_objs > 0 + + # Mask the input data to only consider pixels within secondary objects + masked_data = np.zeros_like(data) + masked_data[second_obj_mask] = data[second_obj_mask] + + # Apply white tophat filter to highlight foci + tophat = skimage.morphology.white_tophat( + masked_data, footprint=skimage.morphology.disk(radius) + ) + + # Apply Laplacian of Gaussian to the filtered image + tophat_log = log_ndi(tophat, sigma=radius) + + # Threshold the image to create a binary mask + mask = tophat_log > threshold + + # Remove small objects from the mask + mask = skimage.morphology.remove_small_objects(mask, min_size=(radius**2)) + + # Ensure we only keep foci within secondary objects + mask = mask & second_obj_mask + + # Label connected components in the mask + labeled = skimage.measure.label(mask) + + # Apply watershed algorithm to refine segmentation + labeled = apply_watershed(labeled, smooth=1) + + if remove_border_foci: + # Create a border mask for secondary objects + second_obj_border = skimage.segmentation.find_boundaries(second_obj_mask) + # Remove foci touching the secondary object border + labeled = remove_border(labeled, second_obj_border) + + return labeled + + +def apply_watershed(img, smooth=4): + """Apply the watershed algorithm to the given image to refine segmentation. + + Args: + img (numpy.ndarray): Input binary image. + smooth (float, optional): Size of Gaussian kernel used to smooth the distance map. Default is 4. + + Returns: + result (numpy.ndarray): Labeled image after watershed segmentation. + """ + # If empty image, return as is + if np.sum(img) == 0: + return img + + # Compute the distance transform of the image + distance = ndi.distance_transform_edt(img) + + if smooth > 0: + # Apply Gaussian smoothing to the distance transform + distance = skimage.filters.gaussian(distance, sigma=smooth) + + # Identify local maxima in the distance transform + local_max_coords = skimage.feature.peak_local_max( + distance, footprint=np.ones((3, 3)), exclude_border=False + ) + + # Create a boolean mask for peaks + local_max = np.zeros_like(distance, dtype=bool) + if len(local_max_coords) > 0: # Check if any peaks were found + local_max[tuple(local_max_coords.T)] = ( + True # Convert coordinates to a boolean mask + ) + + # Label the local maxima + markers = ndi.label(local_max)[0] + + # Apply watershed algorithm to the distance transform + result = skimage.segmentation.watershed(-distance, markers, mask=img) + else: + # If no peaks found, return the original image + result = img + + return result.astype(np.uint16) + + +def remove_border(labels, mask, dilate=2): + """Remove labeled regions that touch the border of the given mask. + + Args: + labels (numpy.ndarray): Labeled image. + mask (numpy.ndarray): Mask indicating the border regions. + dilate (int, optional): Number of dilation iterations to apply to the mask. Default is 2. + + Returns: + labels (numpy.ndarray): Labeled image with border regions removed. + """ + # Dilate the mask to ensure regions touching the border are included + if dilate > 0: + mask = skimage.morphology.binary_dilation(mask, np.ones((dilate, dilate))) + + # Identify labels that need to be removed + remove = np.unique(labels[mask]) + + # Remove the identified labels from the labeled image + labels = labels.copy() + labels.flat[np.in1d(labels, remove)] = 0 + + return labels + + +# Define foci features specific to secondary objects +foci_features = { + "foci_count": lambda r: count_labels(r.intensity_image), + "foci_area": lambda r: (r.intensity_image > 0).sum(), + "foci_area_ratio": lambda r: (r.intensity_image > 0).sum() / r.area + if r.area > 0 + else 0, +} + + +def count_labels(labels, return_list=False): + """Count the unique non-zero labels in a labeled segmentation mask. + + Args: + labels (numpy array): Labeled segmentation mask. + return_list (bool): Flag indicating whether to return the list of unique labels along with the count. + + Returns: + int or tuple: Number of unique non-zero labels. If return_list is True, returns a tuple containing the count + and the list of unique labels. + """ + # Get unique labels in the segmentation mask + uniques = np.unique(labels) + # Remove the background label (0) + ls = np.delete(uniques, np.where(uniques == 0)) + # Count the unique non-zero labels + num_labels = len(ls) + # Return the count or both count and list of unique labels based on return_list flag + if return_list: + return num_labels, ls + return num_labels diff --git a/workflow/lib/phenotype/segment_secondary_object.py b/workflow/lib/phenotype/segment_secondary_object.py new file mode 100644 index 00000000..13462879 --- /dev/null +++ b/workflow/lib/phenotype/segment_secondary_object.py @@ -0,0 +1,1745 @@ +"""Segment secondary objects using thresholding or ML methods and visualize results. + +This module provides functions for segmenting and visualizing secondary objects in microscopy images. +Both traditional threshold-based and machine learning (ML) segmentation methods are supported, +with a shared post-processing pipeline that ensures consistent output formats. + +Architecture: + - segment_second_objs(): Basic segmentation (thresholding + declumping) + - segment_second_objs_ml(): ML-based segmentation template (Cellpose, StarDist, etc.) + - _postprocess_secondary_objects(): Shared post-processing for both methods + * Size filtering (Feret diameter or area) + * Cell association (spatial overlap) + * Cell summary statistics + * Cytoplasm mask updates + +Implementing Custom ML Segmentation: + Users implementing segment_second_objs_ml() only need to: + 1. Extract the target channel from the image + 2. Run their ML model to get a labeled mask (e.g., Cellpose, StarDist) + 3. Return the labeled mask to the shared post-processing pipeline + +Key Functions: + - segment_second_objs(): Basic segmentation + - segment_second_objs_ml(): ML-based segmentation template (user implements) + - _postprocess_secondary_objects(): Shared post-processing pipeline + - create_second_obj_boundary_visualization(): Visualize segmentation results + - create_second_obj_standard_visualization(): Standard visualization panel + +Helper Functions: + - apply_threshold_method(): Thresholding methods (Otsu, Li) + - apply_declumping(): CellProfiler-compatible declumping + - get_feret_diameters(): Compute Feret diameters + - create_empty_results(): Generate empty result structures + - get_spatial_overlap_candidates(): Spatial indexing for cell-object association + +""" + +import numpy as np +import pandas as pd +from scipy import ndimage +from skimage import filters, morphology, measure, segmentation, feature, util, exposure +from skimage.segmentation import mark_boundaries +import matplotlib.pyplot as plt +from microfilm.microplot import Microimage +from lib.shared.configuration_utils import create_micropanel +from lib.shared.segment_cellpose import ( + prepare_cellpose, + create_cellpose_model, + CELLPOSE_VERSION, + CELLPOSE_4X, +) +import cv2 + +# Check if Cellpose is available (for error messaging) +try: + import cellpose + + CELLPOSE_AVAILABLE = True +except ImportError: + CELLPOSE_AVAILABLE = False + + +def segment_second_objs_ml( + image, + second_obj_channel_index, + cell_masks=None, + cytoplasm_masks=None, + # Post-processing parameters (shared) + second_obj_min_size=10, + second_obj_max_size=200, + size_filter_method="feret", + max_objects_per_cell=120, + overlap_threshold=0.1, + nuclei_centroids=None, + max_total_objects=1000, + # Preprocessing parameters + logscale=True, + # ML-specific parameters - users add more as needed + **ml_params, +): + """Segment secondary objects using ML models (Cellpose, StarDist, etc.). + + This function implements ML-based segmentation for secondary objects with support + for both Cellpose and StarDist models. Users can choose the appropriate model + based on their object morphology: + - Cellpose: Better for irregular shapes (vacuoles, organelles with varying morphology) + - StarDist: Better for round/star-convex objects (nuclei-like structures) + + The shared post-processing pipeline (_postprocess_secondary_objects) will handle: + - Size filtering (Feret diameter or area) + - Cell association (spatial overlap) + - Cell summary statistics + - Cytoplasm mask updates + + Parameters + ---------- + image : ndarray + Multichannel image data with shape [channels, height, width] + second_obj_channel_index : int + Index of the channel used for secondary object detection + cell_masks : ndarray + Cell segmentation masks with unique integers for each cell + cytoplasm_masks : ndarray, optional + Cytoplasm segmentation masks. If provided, secondary object + regions will be removed from cytoplasm masks + + second_obj_min_size : float + Minimum size for valid secondary objects + second_obj_max_size : float + Maximum size for valid secondary objects + size_filter_method : str + Size filtering method ("feret" or "area") + max_objects_per_cell : int + Maximum secondary objects allowed per cell + overlap_threshold : float + Minimum overlap ratio to associate object with cell (0.0-1.0) + nuclei_centroids : dict, DataFrame, or None + Cell nuclei centroids for distance calculations + max_total_objects : int or None + Failsafe limit on detected objects + logscale : bool + Apply log scaling and normalization preprocessing to the target channel + before segmentation. This matches the preprocessing used in segment_cellpose + and improves segmentation performance. Default is True. + + **ml_params : dict + Additional ML model parameters. Required and optional parameters depend on ml_method: + + Common parameters: + - second_obj_method : str (required) + ML model to use: "cellpose" or "stardist" + - gpu : bool (default: False) + Whether to use GPU acceleration + + For second_obj_method="cellpose": + - second_obj_cellpose_model : str (default: 'cyto3') + Cellpose model type ('cyto3', 'cyto2', 'cyto', 'nuclei', etc.) + - second_obj_diameter : float or None (default: None) + Expected diameter of objects in pixels. If None, estimated automatically + - second_obj_flow_threshold : float (default: 0.4) + Flow error threshold for Cellpose segmentation + - second_obj_cellprob_threshold : float (default: 0.0) + Cell probability threshold for Cellpose + + For second_obj_method="stardist": + - second_obj_stardist_model : str (default: '2D_versatile_fluo') + StarDist pretrained model name + - second_obj_prob_threshold : float (default: 0.5) + Probability threshold for object detection + - second_obj_nms_threshold : float (default: 0.4) + Non-maximum suppression threshold + + Returns: + ------- + tuple + - second_obj_masks: Labeled mask of secondary objects [height, width] + - cell_second_obj_table: Dict with 'cell_summary' and 'second_obj_cell_mapping' DataFrames + - updated_cytoplasm_masks: Cytoplasm masks with secondary objects removed (if provided) + + Raises: + ------ + ValueError + If ml_method is not 'cellpose' or 'stardist', or if required packages are not installed + + Notes: + ----- + - All post-processing is handled by _postprocess_secondary_objects() + - Output format is guaranteed to match segment_second_objs() + - Requires cellpose or stardist packages to be installed + """ + # Extract and preprocess target channel + if logscale: + # Use prepare_cellpose for preprocessing (ensures consistency with training) + rgb = prepare_cellpose( + image, + dapi_index=second_obj_channel_index, # Dummy - will use cyto channel + cyto_index=second_obj_channel_index, # Target channel + helper_index=None, + logscale=True, + ) + target_channel = rgb[1] # Extract green (log scaled + normalized) + else: + target_channel = image[second_obj_channel_index].copy() + + # Get ML method + ml_method = ml_params.get("second_obj_method", None) + if ml_method is None: + raise ValueError( + "second_obj_method must be specified in ml_params. " + "Valid options: 'cellpose' or 'stardist'" + ) + + gpu = ml_params.get("gpu", False) + + # Route to appropriate ML model + if ml_method == "cellpose": + # Cellpose parameters + model_type = ml_params.get("second_obj_cellpose_model", "cyto3") + diameter = ml_params.get("second_obj_diameter", None) + flow_threshold = ml_params.get("second_obj_flow_threshold", 0.4) + cellprob_threshold = ml_params.get("second_obj_cellprob_threshold", 0.0) + + print( + f"Running Cellpose {model_type} model for secondary object segmentation..." + ) + if diameter is not None: + print(f" Using diameter: {diameter:.1f} pixels") + else: + print(f" Diameter will be estimated automatically") + print(f" Flow threshold: {flow_threshold}") + print(f" Cell probability threshold: {cellprob_threshold}") + print(f" GPU: {gpu}") + + # Check Cellpose availability + if not CELLPOSE_AVAILABLE: + raise ImportError( + "Cellpose is required for ML-based secondary object segmentation. " + "Install it with: pip install cellpose" + ) + + # Initialize Cellpose model (handles version detection and validation) + model = create_cellpose_model(model_type, gpu=gpu) + + # Run Cellpose segmentation + # Note: CellposeModel.eval() returns 3 values (masks, flows, styles) + # Diameter must be specified explicitly for Cellpose 4.x + labeled_mask, flows, styles = model.eval( + target_channel, + diameter=diameter, + flow_threshold=flow_threshold, + cellprob_threshold=cellprob_threshold, + ) + + print(f"Cellpose detected {len(np.unique(labeled_mask)) - 1} secondary objects") + if diameter is not None: + print(f"Using diameter: {diameter:.1f} pixels") + + elif ml_method == "stardist": + # StarDist parameters + model_type = ml_params.get("second_obj_stardist_model", "2D_versatile_fluo") + prob_thresh = ml_params.get("second_obj_prob_threshold", 0.5) + nms_thresh = ml_params.get("second_obj_nms_threshold", 0.4) + + print( + f"Running StarDist {model_type} model for secondary object segmentation..." + ) + print(f" Probability threshold: {prob_thresh}") + print(f" NMS threshold: {nms_thresh}") + print(f" GPU: {gpu}") + + # Import StarDist + try: + from stardist.models import StarDist2D + except ImportError: + raise ImportError( + "StarDist is required for ML-based secondary object segmentation. " + "Install it with: pip install stardist" + ) + + # Initialize StarDist model + model = StarDist2D.from_pretrained(model_type) + + # Run StarDist segmentation + labeled_mask, details = model.predict_instances( + target_channel, prob_thresh=prob_thresh, nms_thresh=nms_thresh + ) + + print(f"StarDist detected {len(np.unique(labeled_mask)) - 1} secondary objects") + + else: + raise ValueError( + f"Unknown ml_method: {ml_method}. Valid options: 'cellpose' or 'stardist'" + ) + + # Shared post-processing pipeline + return _postprocess_secondary_objects( + second_obj_masks=labeled_mask, # ML model output + cell_masks=cell_masks, + cytoplasm_masks=cytoplasm_masks, + second_obj_min_size=second_obj_min_size, + second_obj_max_size=second_obj_max_size, + size_filter_method=size_filter_method, + max_objects_per_cell=max_objects_per_cell, + overlap_threshold=overlap_threshold, + nuclei_centroids=nuclei_centroids, + max_total_objects=max_total_objects, + image=image, + second_obj_channel_index=second_obj_channel_index, + ) + + +def estimate_second_obj_diameter( + image, second_obj_channel_index, method="cellpose", **kwargs +): + """Estimate the diameter of secondary objects in an image channel. + + This is a convenience function to help users estimate appropriate diameter + parameters for ML-based secondary object segmentation. + + Parameters + ---------- + image : ndarray + Multichannel image data with shape [channels, height, width] + second_obj_channel_index : int + Index of the channel containing secondary objects + method : str + Method to use for diameter estimation: + - "cellpose": Use Cellpose's built-in diameter estimation (default) + - "manual": Manually measure from image statistics + **kwargs : dict + Additional parameters for the estimation method: + - For method="cellpose": + - model_type : str (default: 'cyto3') + - gpu : bool (default: False) + + Returns: + ------- + diameter : float + Estimated diameter in pixels + + Examples: + -------- + >>> diameter = estimate_second_obj_diameter( + ... aligned_image, + ... second_obj_channel_index=7, + ... method="cellpose", + ... model_type="cyto3" + ... ) + >>> print(f"Estimated diameter: {diameter:.1f} pixels") + """ + target_channel = image[second_obj_channel_index] + + if method == "cellpose": + # Check Cellpose availability + if not CELLPOSE_AVAILABLE: + raise ImportError( + "Cellpose is required for diameter estimation. " + "Install it with: pip install cellpose" + ) + + # Cellpose 4.x does not support automatic diameter estimation + if CELLPOSE_4X: + raise NotImplementedError( + "Automatic diameter estimation is not supported with Cellpose 4.x. " + "Please specify second_obj_diameter explicitly in your config, " + "or use method='manual' for threshold-based estimation, " + "or downgrade to Cellpose 3.x: pip install cellpose==3.1.0" + ) + + model_type = kwargs.get("model_type", "cyto3") + gpu = kwargs.get("gpu", False) + + print(f"Estimating secondary object diameter using Cellpose {model_type}...") + + # Cellpose 3.x: Use the old API which supports diameter estimation + from cellpose import models as cellpose_models + + model = cellpose_models.Cellpose(gpu=gpu, model_type=model_type) + + # Run segmentation with automatic diameter estimation + _, _, _, diameter = model.eval( + target_channel, + diameter=None, # Auto-estimate + channels=[0, 0], + ) + + print(f"Estimated diameter: {diameter:.1f} pixels") + return float(diameter) + + elif method == "manual": + # Simple estimation based on image statistics + # Threshold the image and measure typical object sizes + from skimage import filters, measure + from scipy import ndimage + + # Apply Otsu threshold + thresh = filters.threshold_otsu(target_channel) + binary = target_channel > thresh + + # Label objects + labeled, _ = ndimage.label(binary) + regions = measure.regionprops(labeled) + + if len(regions) == 0: + print("No objects detected for diameter estimation") + return None + + # Calculate median equivalent diameter + diameters = [r.equivalent_diameter for r in regions] + diameter = np.median(diameters) + + print( + f"Estimated diameter (median of {len(regions)} objects): {diameter:.1f} pixels" + ) + return float(diameter) + + else: + raise ValueError(f"Unknown method: {method}. Use 'cellpose' or 'manual'") + + +def apply_threshold_method(image, method="otsu_two_peak"): + """Apply specified thresholding method to an image. + + Parameters + ---------- + image : ndarray + Input image (should be preprocessed with log transform and smoothing) + method : str + Thresholding method to use. + Options: + - 'otsu_two_peak': Standard Otsu thresholding (2-class) + - 'otsu_three_peak_mid_bg': 3-class Otsu, middle class as background + - 'otsu_three_peak_mid_fg': 3-class Otsu, middle class as foreground + - 'min_cross_entropy': Minimum cross entropy (Li) thresholding + + Returns: + ------- + threshold : float + Computed threshold value + binary_mask : ndarray + Binary mask after thresholding + """ + if method == "otsu_two_peak": + # Standard two-class Otsu + threshold = filters.threshold_otsu(image) + binary_mask = image > threshold + + elif method == "otsu_three_peak_mid_bg": + # Three-class Otsu, treat middle intensity class as background + threshold = filters.threshold_multiotsu(image, classes=3) + # Keep only the highest intensity class (threshold[1] separates mid from high) + binary_mask = image > threshold[1] + + elif method == "otsu_three_peak_mid_fg": + # Three-class Otsu, treat middle intensity class as foreground + threshold = filters.threshold_multiotsu(image, classes=3) + # Keep both middle and high intensity classes (threshold[0] separates low from mid) + binary_mask = image > threshold[0] + + elif method == "min_cross_entropy": + # Minimum cross entropy (Li) method + threshold = filters.threshold_li(image) + binary_mask = image > threshold + + else: + raise ValueError( + f"Unknown threshold method: {method}. " + f"Valid options: 'otsu_two_peak', 'otsu_three_peak_mid_bg', " + f"'otsu_three_peak_mid_fg', 'min_cross_entropy'" + ) + + return threshold, binary_mask + + +def segment_second_objs( + image, + second_obj_channel_index, + cell_masks=None, + cytoplasm_masks=None, + # Size filtering + second_obj_min_size=10, + second_obj_max_size=200, + size_filter_method="feret", + # Pre-processing + threshold_smoothing_scale=1.3488, + threshold_method="otsu_two_peak", + use_morphological_opening=True, + opening_disk_radius=1, + fill_holes="both", + # Declumping method (CellProfiler standard) + declump_method="shape", + declump_mode="watershed", + # Seed detection (CellProfiler naming) + suppress_local_maxima=20, + maxima_reduction_factor=None, + # Shape-based refinement (independent from declump_method) + use_shape_refinement=False, + proportion_threshold=0.4, + # Cell association + max_objects_per_cell=120, + overlap_threshold=0.1, + nuclei_centroids=None, + # Failsafe + max_total_objects=1000, + # Debugging + return_threshold_output=False, +): + """Segment secondary objects within cells using CellProfiler-compatible thresholding and declumping. + + Args: + image (numpy.ndarray): Multichannel image data with shape [channels, height, width]. + second_obj_channel_index (int): Index of the channel used for secondary object detection. + cell_masks (numpy.ndarray): Cell segmentation masks with unique integers for each cell. + cytoplasm_masks (numpy.ndarray, optional): Cytoplasm segmentation masks with unique integers. + If provided, secondary object regions will be removed from cytoplasm masks. + + second_obj_min_size (float, optional): Minimum size for valid secondary objects (default: 10). + Interpreted as Feret diameter or area depending on size_filter_method. + second_obj_max_size (float, optional): Maximum size for valid secondary objects (default: 200). + size_filter_method (str, optional): Size filtering method (default: "feret"). + - "feret": Use Feret diameters (min and max widths of rotated bounding box) + - "area": Use pixel area (CellProfiler standard) + + threshold_smoothing_scale (float, optional): Sigma for Gaussian smoothing before thresholding. Default is 1.3488. + threshold_method (str, optional): Thresholding method to use (default: "otsu_two_peak"). + Options: + - "otsu_two_peak": Standard 2-class Otsu thresholding + - "otsu_three_peak_mid_bg": 3-class Otsu, keeps only highest intensity class + - "otsu_three_peak_mid_fg": 3-class Otsu, keeps middle and high intensity classes + - "min_cross_entropy": Minimum cross entropy (Li) thresholding + use_morphological_opening (bool, optional): Apply opening to separate weakly connected objects (default: True). + opening_disk_radius (int, optional): Radius of disk structuring element for opening (default: 1). + fill_holes (str, optional): When to fill holes in segmented objects (default: "both"). + Options: + - "threshold": Fill holes only after thresholding (before declumping) + - "declump": Fill holes only after declumping (per-label filling) + - "both": Fill holes after both thresholding and declumping + - "none": Do not fill holes at any stage + + declump_method (str, optional): Method for separating clumped objects (default: "shape"). + CellProfiler standard methods: + - "none": No declumping (connected components only) + - "shape": Distance transform peaks (radial distance) + - "intensity": Local intensity maxima + - "shape_intensity": Combined distance + intensity peaks + + declump_mode (str, optional): Watershed segmentation mode (default: "watershed"). + - "watershed": Standard watershed from markers + - "propagate": Distance propagation variant + - "none": Use markers only without watershed + + suppress_local_maxima (int, optional): Minimum spacing between seed points in pixels (default: 20). + CellProfiler parameter. Controls spatial separation of detected peaks. + + maxima_reduction_factor (float or None, optional): H-minima threshold for suppressing weak peaks (default: None). + Range: 0.0-1.0. Higher values = more aggressive suppression. + If None, no h-minima filtering applied. + Formula: h = maxima_reduction_factor * (peak_map_max - peak_map_min) + This is applied DURING seed detection (before watershed). + + use_shape_refinement (bool, optional): Apply boundary/perimeter quality control after declumping (default: False). + Custom feature not in CellProfiler. When enabled, evaluates watershed splits + and rejects splits where the dividing boundary is long relative to perimeter. + This is applied AFTER watershed declumping as a refinement step. + + proportion_threshold (float, optional): Boundary/perimeter ratio threshold for shape refinement (default: 0.4). + Only used when use_shape_refinement=True. + Splits accepted if boundary_length / perimeter < proportion_threshold. + + max_objects_per_cell (int, optional): Maximum secondary objects allowed per cell (default: 120). + overlap_threshold (float, optional): Minimum overlap ratio to associate object with cell (default: 0.1). + nuclei_centroids (dict or DataFrame, optional): Cell nuclei centroids for distance calculations. + Format: {nuclei_id: (i, j)} or DataFrame with columns 'i', 'j'. + + max_total_objects (int or None, optional): Failsafe limit on detected objects (default: 1000). + Returns empty results if exceeded to avoid processing over-segmented images. + + return_threshold_output (bool, optional): If True, returns intermediate thresholding results + for debugging and visualization (default: False). When enabled, the threshold_output + dictionary contains: + - 'binary_mask': Binary mask after thresholding, hole filling, and opening (before declumping) + - 'threshold_value': Computed threshold value from apply_threshold_method() + - 'preprocessed_channel': Log-transformed and Gaussian-smoothed channel used for thresholding + + Returns: + tuple: Returns depend on return_threshold_output flag: + + If return_threshold_output=False (default): + - second_obj_masks (numpy.ndarray): Labeled mask of secondary objects + - cell_second_obj_table (dict): Dictionary with DataFrames containing associations + - updated_cytoplasm_masks (numpy.ndarray): Updated cytoplasm masks (if cytoplasm_masks provided) + + If return_threshold_output=True: + - second_obj_masks (numpy.ndarray): Labeled mask of secondary objects + - cell_second_obj_table (dict): Dictionary with DataFrames containing associations + - updated_cytoplasm_masks (numpy.ndarray): Updated cytoplasm masks (if cytoplasm_masks provided) + - threshold_output (dict): Intermediate thresholding results with keys: + - 'binary_mask': Binary mask before declumping + - 'threshold_value': Threshold value used + - 'preprocessed_channel': Preprocessed channel image + """ + # Extract the secondary object channel + second_obj_img = image[second_obj_channel_index] + second_obj_img = np.clip(second_obj_img, a_min=0, a_max=None) + + # Apply log transform and smoothing + second_obj_log = exposure.adjust_log(second_obj_img + 1) + second_obj_smooth = filters.gaussian( + second_obj_log, sigma=threshold_smoothing_scale + ) + + # Apply selected threshold method + thresh, binary_mask = apply_threshold_method( + second_obj_smooth, method=threshold_method + ) + + # Fill holes after thresholding (if enabled) + if fill_holes in ["threshold", "both"]: + binary_mask = ndimage.binary_fill_holes(binary_mask) + + # Early exit if no objects found + if not np.any(binary_mask): + print("No objects detected after thresholding") + empty_results = create_empty_results( + cell_masks, cytoplasm_masks, nuclei_centroids + ) + + # Handle return_threshold_output for empty case + if return_threshold_output: + threshold_output = { + "binary_mask": binary_mask, + "threshold_value": thresh, + "preprocessed_channel": second_obj_smooth, + } + # Add threshold_output to the tuple + if cytoplasm_masks is not None: + return (*empty_results, threshold_output) + else: + return (*empty_results, threshold_output) + else: + return empty_results + + # Failsafe: Check for excessive objects early + if max_total_objects is not None: + temp_labeled, num_components = ndimage.label(binary_mask) + if num_components > max_total_objects: + print( + f"Failsafe triggered: Detected {num_components} objects (limit: {max_total_objects})" + ) + print("Returning empty results to avoid processing over-segmented image") + empty_results = create_empty_results( + cell_masks, cytoplasm_masks, nuclei_centroids + ) + + # Handle return_threshold_output for failsafe case + if return_threshold_output: + threshold_output = { + "binary_mask": binary_mask, + "threshold_value": thresh, + "preprocessed_channel": second_obj_smooth, + } + if cytoplasm_masks is not None: + return (*empty_results, threshold_output) + else: + return (*empty_results, threshold_output) + else: + return empty_results + + # Morphological opening + if use_morphological_opening: + binary_mask = apply_morphological_opening( + binary_mask, opening_disk_radius=opening_disk_radius + ) + + # Capture intermediate state for visualization + if return_threshold_output: + threshold_binary_mask = binary_mask.copy() + threshold_value_stored = thresh + threshold_preprocessed_channel = second_obj_smooth.copy() + + # Declumping + declumped = apply_declumping( + binary_mask, + second_obj_smooth, + declump_method=declump_method, + declump_mode=declump_mode, + suppress_local_maxima=suppress_local_maxima, + maxima_reduction_factor=maxima_reduction_factor, + ) + + print( + f"After declumping ({declump_method}): {len(np.unique(declumped)) - 1} objects" + ) + + # Optionally apply shape-based refinement (independent from declump_method) + if use_shape_refinement: + print("Applying shape-based boundary/perimeter refinement...") + declumped = shape_based_declumping( + declumped > 0, + second_obj_img=second_obj_img, + min_distance=suppress_local_maxima, + proportion_threshold=proportion_threshold, + ) + print(f"After shape refinement: {len(np.unique(declumped)) - 1} objects") + + # Fill holes after declumping (if enabled) + if fill_holes in ["declump", "both"]: + unique_labels = np.unique(declumped[declumped > 0]) + for label in unique_labels: + mask = declumped == label + filled = ndimage.binary_fill_holes(mask) + declumped[filled] = label + + # Apply shared post-processing pipeline: size filtering, cell association, statistics, cytoplasm updates + post_results = _postprocess_secondary_objects( + second_obj_masks=declumped, + cell_masks=cell_masks, + cytoplasm_masks=cytoplasm_masks, + second_obj_min_size=second_obj_min_size, + second_obj_max_size=second_obj_max_size, + size_filter_method=size_filter_method, + max_objects_per_cell=max_objects_per_cell, + overlap_threshold=overlap_threshold, + nuclei_centroids=nuclei_centroids, + max_total_objects=max_total_objects, + image=image, + second_obj_channel_index=second_obj_channel_index, + ) + + # Handle return_threshold_output flag for debugging + if return_threshold_output: + # Create threshold output dictionary + threshold_output = { + "binary_mask": threshold_binary_mask, + "threshold_value": threshold_value_stored, + "preprocessed_channel": threshold_preprocessed_channel, + } + # Append threshold_output to results tuple + return (*post_results, threshold_output) + else: + # Standard return (backward compatible) + return post_results + + +def create_second_obj_boundary_visualization( + image, + second_obj_channel_index, + cell_masks, + second_obj_masks, + channel_names=None, + channel_cmaps=None, +): + """Create enhanced visualization showing cells and secondary objects. + + Args: + image (numpy.ndarray): Multichannel image data with shape [channels, height, width]. + second_obj_channel_index (int): Index of the channel used for secondary object detection. + cell_masks (numpy.ndarray): Cell segmentation masks with unique integers for each cell. + second_obj_masks (numpy.ndarray): Secondary object segmentation masks with original secondary object IDs. + channel_names (list of str, optional): Names for each channel in the image. + channel_cmaps (list of str, optional): Color maps for each channel in the image. + + Returns: + matplotlib.figure.Figure: The created micropanel figure showing the cell boundaries (green) + and secondary object boundaries (magenta) overlaid on the image. + """ + if channel_names is None or len(channel_names) <= second_obj_channel_index: + channel_name = f"Channel {second_obj_channel_index}" + else: + channel_name = channel_names[second_obj_channel_index] + + # Get secondary object channel + second_obj_img = image[second_obj_channel_index].copy() + + # Create a copy of the original image for the merged view with boundaries + merged_img = image.copy() + + # Function to add boundaries to an image + def add_boundaries(base_image, base_is_multichannel=True): + # Determine the shape based on whether base_image is multichannel or single channel + if base_is_multichannel: + # For multichannel image, keep as is + enhanced_img = base_image.copy() + height, width = base_image.shape[1], base_image.shape[2] + num_channels = base_image.shape[0] + else: + # For single channel image, expand to 3 channels + height, width = base_image.shape + num_channels = 3 + # Create 3-channel image with the base image in all channels + enhanced_img = np.zeros((num_channels, height, width), dtype=np.float32) + base_norm = base_image / (base_image.max() if base_image.max() > 0 else 1.0) + for c in range(num_channels): + enhanced_img[c] = base_norm + + # Add cell boundaries (green) + if base_is_multichannel: + # For multichannel image, we need to create a temporary RGB image + # to use mark_boundaries, then extract the green channel + temp_img = np.zeros((height, width, 3), dtype=np.float32) + for c in range(min(3, num_channels)): + temp_img[:, :, c] = enhanced_img[c] / ( + enhanced_img[c].max() if enhanced_img[c].max() > 0 else 1.0 + ) + + cell_boundary_img = mark_boundaries( + temp_img, + cell_masks, + color=(0, 1, 0), # Green for cells + mode="thick", + ) + + # Update the green channel with cell boundaries - make them more prominent + cell_boundary_intensity = ( + 1.2 * enhanced_img[1].max() + ) # Increase intensity by 20% + enhanced_img[1] = np.maximum( + enhanced_img[1], cell_boundary_img[:, :, 1] * cell_boundary_intensity + ) + # Cap values at 1.0 if normalized + if enhanced_img.dtype == np.float32 or enhanced_img.dtype == np.float64: + enhanced_img[1] = np.minimum( + enhanced_img[1], + 1.0 if enhanced_img[1].max() <= 1.0 else enhanced_img[1].max(), + ) + else: + # For single channel image, directly add boundaries to green channel + cell_boundary = mark_boundaries( + base_image, + cell_masks, + color=(0, 1, 0), # Green for cells + mode="thick", + ) + enhanced_img[1] = np.maximum(enhanced_img[1], cell_boundary[:, :, 1]) + + # Add secondary object boundaries (magenta: red + blue) + if base_is_multichannel: + # For multichannel image, create temporary RGB again + second_obj_boundary_img = mark_boundaries( + temp_img, + second_obj_masks > 0, # Binary mask + color=(1, 0, 1), # Magenta for secondary objects + mode="thick", + ) + + # Update red and blue channels with secondary object boundaries + enhanced_img[0] = np.maximum( + enhanced_img[0], + second_obj_boundary_img[:, :, 0] * enhanced_img[0].max(), + ) + if num_channels > 2: # Make sure we have a blue channel + enhanced_img[2] = np.maximum( + enhanced_img[2], + second_obj_boundary_img[:, :, 2] * enhanced_img[2].max(), + ) + else: + # For single channel, add boundaries to red and blue channels + second_obj_boundary = mark_boundaries( + base_image, + second_obj_masks > 0, # Binary mask + color=(1, 0, 1), # Magenta for secondary objects + mode="thick", + ) + enhanced_img[0] = np.maximum(enhanced_img[0], second_obj_boundary[:, :, 0]) + enhanced_img[2] = np.maximum(enhanced_img[2], second_obj_boundary[:, :, 2]) + + return enhanced_img + + # Create merged microimage with boundaries + merged_with_boundaries = add_boundaries(merged_img) + merged_microimage = Microimage( + merged_with_boundaries, channel_names="Merged", cmaps=channel_cmaps + ) + + # Create secondary object channel microimage with boundaries + # Convert single channel to 3D for processing + second_obj_3d = add_boundaries(second_obj_img, base_is_multichannel=False) + boundaries_microimage = Microimage( + second_obj_3d, + channel_names=f"{channel_name}", + cmaps=["pure_red", "pure_green", "pure_blue"], + ) + + # Create the micropanel + microimages = [merged_microimage, boundaries_microimage] + panel = create_micropanel(microimages, add_channel_label=True) + + return panel + + +def create_second_obj_standard_visualization( + aligned_image, + second_obj_channel_index, + second_obj_channel_name, + second_obj_masks, + threshold_output=None, + label_color="magenta", +): + """Create standard visualization panel for secondary object segmentation. + + Parameters + ---------- + aligned_image : ndarray + Multichannel aligned image [channels, height, width] + second_obj_channel_index : int + Index of the channel used for secondary object detection + second_obj_channel_name : str + Name of the secondary object channel (e.g., "CDPK1") + second_obj_masks : ndarray + Labeled mask of segmented secondary objects + threshold_output : dict, optional + Dictionary containing threshold debugging output with keys: + - 'preprocessed_channel': Log-transformed and Gaussian-smoothed channel + - 'binary_mask': Binary mask after thresholding + If None, creates simple 1x2 panel. If provided, creates 2x2 panel. + label_color : str, optional + Color for channel labels (default: 'magenta') + + Returns: + ------- + panel : Micropanel + Micropanel object with visualizations + """ + from lib.shared.configuration_utils import random_cmap + + # Build secondary object colormap + second_obj_cmap = random_cmap(num_colors=len(np.unique(second_obj_masks))) + + if threshold_output: + # 2x2 grid when threshold_output is provided + micro_images = [ + Microimage( + threshold_output["preprocessed_channel"], + channel_names="Preprocessed", + cmaps="gray", + ), + Microimage( + threshold_output["binary_mask"], + channel_names="Threshold Binary", + cmaps="gray", + ), + Microimage( + aligned_image[second_obj_channel_index], + channel_names=f"{second_obj_channel_name} (Raw)", + cmaps="gray", + ), + Microimage( + second_obj_masks, + cmaps=second_obj_cmap, + channel_names="Secondary Objects", + ), + ] + num_cols = 2 + else: + # 1x2 grid when threshold_output is None + micro_images = [ + Microimage( + aligned_image[second_obj_channel_index], + channel_names=f"{second_obj_channel_name} (Raw)", + cmaps="gray", + ), + Microimage( + second_obj_masks, + cmaps=second_obj_cmap, + channel_names="Secondary Objects", + ), + ] + num_cols = 2 + + panel = create_micropanel( + micro_images, + add_channel_label=True, + num_cols=num_cols, + ) + + # Set all channel labels to specified color + for ax in panel.fig.axes: + for text in ax.texts: + text.set_color(label_color) + + return panel + + +def get_feret_diameters(coords): + """Compute the minimum and maximum Feret diameters of a 2D shape. + + The Feret diameters are calculated using OpenCV's minAreaRect, which finds + the smallest-area rotated bounding rectangle that encloses the input coordinates. + + Parameters + ---------- + coords : ndarray of shape (N, 2) + An array of (x, y) coordinates representing the pixels or contour of a region. + + Returns: + ------- + feret_min : float + The shortest distance between two parallel lines tangent to the object + (i.e., the minimum Feret diameter). + + feret_max : float + The longest distance between two parallel lines tangent to the object + (i.e., the maximum Feret diameter). + + Notes: + ----- + - This method assumes the input coordinates define a planar shape (e.g., from a binary mask or regionprops). + - The returned values are in the same units as the input coordinates (typically pixels). + - Internally uses OpenCV's cv2.minAreaRect for fast and robust measurement. + """ + cnt = coords.astype(np.int32) + rect = cv2.minAreaRect(cnt) + w, h = rect[1] + return min(w, h), max(w, h) + + +def get_feret_diameters(coords): + """Compute the minimum and maximum Feret diameters of a 2D shape. + + The Feret diameters are calculated using OpenCV's minAreaRect, which finds + the smallest-area rotated bounding rectangle that encloses the input coordinates. + + Parameters + ---------- + coords : ndarray of shape (N, 2) + An array of (x, y) coordinates representing the pixels or contour of a region. + + Returns: + ------- + feret_min : float + The shortest distance between two parallel lines tangent to the object + (i.e., the minimum Feret diameter). + + feret_max : float + The longest distance between two parallel lines tangent to the object + (i.e., the maximum Feret diameter). + + Notes: + ----- + - This method assumes the input coordinates define a planar shape (e.g., from a binary mask or regionprops). + - The returned values are in the same units as the input coordinates (typically pixels). + - Internally uses OpenCV's cv2.minAreaRect for fast and robust measurement. + """ + cnt = coords.astype(np.int32) + rect = cv2.minAreaRect(cnt) + w, h = rect[1] + return min(w, h), max(w, h) + + +def apply_morphological_opening(binary_mask, opening_disk_radius=1): + """Apply morphological opening to separate weakly connected secondary objects. + + Parameters + ---------- + binary_mask : ndarray + Binary mask of secondary objects + opening_disk_radius : int + Radius of disk structuring element (larger = more aggressive) + + Returns: + ------- + opened_mask : ndarray + Morphologically opened mask + """ + footprint = morphology.disk(max(1, opening_disk_radius)) + opened = morphology.binary_opening(binary_mask, footprint=footprint) + + # Recover small objects that were removed by opening + removed = binary_mask & ~opened + small_objects, num = ndimage.label(removed) + + # Only recover objects at least as large as the structuring element + min_recoverable_size = np.pi * opening_disk_radius**2 + for i in range(1, num + 1): + obj_mask = small_objects == i + if np.sum(obj_mask) >= min_recoverable_size: + opened |= obj_mask + + return opened + + +def apply_h_minima_suppression(peak_map, h_factor): + """Apply h-minima transform to suppress weak local maxima. + + This complements spatial suppression (min_distance) by filtering peaks + based on their prominence/height in the distance or intensity map. + + Parameters + ---------- + peak_map : ndarray + Distance transform or intensity image + h_factor : float + Height threshold factor (0.0-1.0) + h = h_factor * (peak_map.max() - peak_map.min()) + Higher values = more aggressive suppression + + Returns: + ------- + filtered_map : ndarray + Map with weak maxima suppressed + """ + if h_factor <= 0 or h_factor > 1: + raise ValueError(f"h_factor must be in (0, 1], got {h_factor}") + + # Calculate absolute height threshold + h = h_factor * (peak_map.max() - peak_map.min()) + + # Apply h-minima transform + filtered_map = morphology.h_minima(peak_map, h=h) + + return filtered_map + + +def apply_declumping( + binary_mask, + second_obj_smooth, + declump_method, + declump_mode, + suppress_local_maxima, + maxima_reduction_factor, +): + """Apply declumping based on CellProfiler-compatible method selection. + + Parameters + ---------- + binary_mask : ndarray + Binary mask of secondary objects + second_obj_smooth : ndarray + Smoothed intensity image (log + Gaussian filtered) + declump_method : str + "none", "shape", "intensity", "shape_intensity", "distance" + declump_mode : str + "watershed", "propagate", "none" + suppress_local_maxima : int + Minimum distance between peaks (spatial constraint) + maxima_reduction_factor : float or None + H-minima threshold (0.0-1.0), None=disabled + + Returns: + ------- + declumped : ndarray + Labeled mask after declumping + + Notes: + ----- + Shape refinement is NOT handled here - it's applied as optional refinement + after this function in the main pipeline. + """ + # Method 1: No declumping + if declump_method == "none": + declumped, _ = ndimage.label(binary_mask) + return declumped + + # Method 2: Shape-based (distance transform) + if declump_method in ["shape", "distance"]: + peak_map = ndimage.distance_transform_edt(binary_mask) + + # Method 3: Intensity-based + elif declump_method == "intensity": + # Use smoothed intensity within mask + peak_map = second_obj_smooth.copy() + peak_map[~binary_mask] = 0 + + # Method 4: Combined shape + intensity + elif declump_method == "shape_intensity": + # Normalize both maps to [0, 1] and average + distance_map = ndimage.distance_transform_edt(binary_mask) + distance_norm = distance_map / (distance_map.max() + 1e-10) + + intensity_map = second_obj_smooth.copy() + intensity_map[~binary_mask] = 0 + intensity_norm = intensity_map / (intensity_map.max() + 1e-10) + + peak_map = (distance_norm + intensity_norm) / 2 + + else: + raise ValueError(f"Unknown declump_method: {declump_method}") + + # Apply h-minima suppression if requested + if maxima_reduction_factor is not None: + peak_map = apply_h_minima_suppression(peak_map, maxima_reduction_factor) + + # Detect local maxima + local_max = feature.peak_local_max( + peak_map, + min_distance=suppress_local_maxima, + labels=binary_mask, + exclude_border=False, + ) + + # Create markers + markers = np.zeros_like(binary_mask, dtype=int) + if len(local_max) == 0: + # No peaks found, return connected components + declumped, _ = ndimage.label(binary_mask) + return declumped + + markers[tuple(local_max.T)] = np.arange(1, len(local_max) + 1) + + # Apply declump_mode + if declump_mode == "none": + # Use markers only (no watershed) + declumped = markers.copy() + + elif declump_mode == "watershed": + # Standard watershed with negative distance + if declump_method in ["shape", "shape_intensity"]: + # Use distance transform for watershed + distance = ndimage.distance_transform_edt(binary_mask) + declumped = segmentation.watershed(-distance, markers, mask=binary_mask) + else: + # For pure intensity, watershed on negative intensity + intensity = second_obj_smooth.copy() + intensity[~binary_mask] = intensity.max() + declumped = segmentation.watershed(intensity, markers, mask=binary_mask) + + elif declump_mode == "propagate": + # Propagate from seeds using positive distance + distance = ndimage.distance_transform_edt(binary_mask) + declumped = segmentation.watershed(distance, markers, mask=binary_mask) + + else: + raise ValueError(f"Unknown declump_mode: {declump_mode}") + + # Recover unassigned regions + missing = (declumped == 0) & binary_mask + if np.any(missing): + labeled_missing, _ = ndimage.label(missing) + if declumped.max() > 0: + labeled_missing[labeled_missing > 0] += declumped.max() + declumped += labeled_missing + + return declumped + + +def shape_based_declumping( + binary_mask, second_obj_img=None, min_distance=20, proportion_threshold=0.4 +): + """Split connected components only when the separating boundary is short relative to the region perimeter. + + Parameters + ---------- + binary_mask : ndarray + Input binary secondary object mask + second_obj_img : ndarray, optional + Intensity image (currently unused, kept for API compatibility) + min_distance : int + Minimum distance between peaks for watershed markers + proportion_threshold : float + If boundary_length / perimeter < proportion_threshold, accept the split + Example: 0.12 means cut must be < 12% of perimeter to split + + Returns: + ------- + labeled : ndarray + Labeled mask after shape-based declumping + """ + labeled_out = np.zeros_like(binary_mask, dtype=int) + next_label = 1 + + # Label connected regions + regions_lab, n = ndimage.label(binary_mask) + + for region_label in range(1, n + 1): + region_mask = regions_lab == region_label + if region_mask.sum() == 0: + continue + + # Distance transform and find peaks + dist = ndimage.distance_transform_edt(region_mask) + peaks = feature.peak_local_max( + dist, min_distance=min_distance, labels=region_mask, exclude_border=False + ) + + # If only one peak, keep as single object + if len(peaks) <= 1: + labeled_out[region_mask] = next_label + next_label += 1 + continue + + # Create markers and apply watershed + markers = np.zeros_like(region_mask, dtype=int) + markers[tuple(peaks.T)] = np.arange(1, len(peaks) + 1) + local_watershed = segmentation.watershed(-dist, markers, mask=region_mask) + + # Vectorized boundary detection + lab = local_watershed + + # Detect boundaries by comparing with neighbors + # Vertical boundaries (compare rows) + vertical_boundary = ( + (lab[:-1, :] != lab[1:, :]) & (lab[:-1, :] > 0) & (lab[1:, :] > 0) + ) + + # Horizontal boundaries (compare columns) + horizontal_boundary = ( + (lab[:, :-1] != lab[:, 1:]) & (lab[:, :-1] > 0) & (lab[:, 1:] > 0) + ) + + # Count total boundary pixels + # We need to count them separately since they have different shapes + boundary_length = np.sum(vertical_boundary) + np.sum(horizontal_boundary) + + prop = measure.regionprops(region_mask.astype(np.uint8))[0] + perimeter = prop.perimeter if prop.perimeter > 0 else 1.0 + + # Accept split if boundary is short relative to perimeter + if (boundary_length / perimeter) < proportion_threshold: + sublabels = np.unique(local_watershed[local_watershed > 0]) + for s in sublabels: + labeled_out[local_watershed == s] = next_label + next_label += 1 + else: + # Reject split, keep as single object + labeled_out[region_mask] = next_label + next_label += 1 + + return labeled_out + + +def create_empty_results(cell_masks, cytoplasm_masks, nuclei_centroids=None): + """Helper function to create empty results when no secondary objects are found. + + Parameters + ---------- + cell_masks : ndarray + Cell segmentation masks + cytoplasm_masks : ndarray, optional + Cytoplasm segmentation masks + nuclei_centroids : dict or DataFrame, optional + Nuclei centroids information + + Returns: + ------- + tuple + Empty secondary object masks, cell_second_obj_table dict, and optionally cytoplasm_masks + """ + cell_ids = np.unique(cell_masks[cell_masks > 0]) + empty_second_obj_masks = np.zeros_like(cell_masks) + + cell_summary = [] + for cell_id in cell_ids: + cell_area = np.sum(cell_masks == cell_id) + summary_entry = { + "cell_id": cell_id, + "has_second_obj": False, + "num_second_objs": 0, + "second_obj_ids": [], + "cell_area": cell_area, + "total_second_obj_area": 0, + "second_obj_area_ratio": 0, + "mean_second_obj_diameter": None, + } + + # Add cell nucleus distance fields if nuclei_centroids was provided + if nuclei_centroids is not None: + summary_entry["mean_distance_to_nucleus"] = None + + cell_summary.append(summary_entry) + + cell_second_obj_table = { + "cell_summary": pd.DataFrame(cell_summary), + "second_obj_cell_mapping": pd.DataFrame(), + } + + if cytoplasm_masks is not None: + return empty_second_obj_masks, cell_second_obj_table, cytoplasm_masks + else: + return empty_second_obj_masks, cell_second_obj_table + + +def get_spatial_overlap_candidates(second_obj_regions, cell_masks): + """Use bounding boxes to pre-filter which cells could overlap with each secondary object. + + Parameters + ---------- + second_obj_regions : dict + Dictionary mapping second_obj_id to regionprops + cell_masks : ndarray + Cell segmentation masks + + Returns: + ------- + candidates : dict + Dictionary mapping second_obj_id to list of candidate cell_ids + """ + # Get all cell regions with their bounding boxes + cell_regions = measure.regionprops(cell_masks) + cell_bboxes = { + r.label: r.bbox for r in cell_regions + } # (min_row, min_col, max_row, max_col) + + candidates = {} + + for second_obj_id, vac_region in second_obj_regions.items(): + vac_bbox = vac_region.bbox # (min_row, min_col, max_row, max_col) + + # Find cells whose bounding boxes intersect with this secondary object's bbox + overlapping_cells = [] + for cell_id, cell_bbox in cell_bboxes.items(): + # Check if bounding boxes overlap + if not ( + vac_bbox[2] < cell_bbox[0] # second_obj above cell + or vac_bbox[0] > cell_bbox[2] # second_obj below cell + or vac_bbox[3] < cell_bbox[1] # second_obj left of cell + or vac_bbox[1] > cell_bbox[3] + ): # second_obj right of cell + overlapping_cells.append(cell_id) + + candidates[second_obj_id] = overlapping_cells + + return candidates + + +def _postprocess_secondary_objects( + second_obj_masks, + cell_masks, + cytoplasm_masks, + second_obj_min_size, + second_obj_max_size, + size_filter_method, + max_objects_per_cell, + overlap_threshold, + nuclei_centroids, + max_total_objects, + image=None, + second_obj_channel_index=None, +): + """Apply post-processing pipeline to secondary object masks. + + This function performs the shared post-processing steps for both + basic thresholding and ML-based secondary object segmentation: + 1. Size filtering (Feret diameter or area) + 2. Cell association (spatial overlap) + 3. Cell summary statistics + 4. Cytoplasm mask updates + + Parameters + ---------- + second_obj_masks : ndarray + Labeled mask of secondary objects (integer labels, background=0) + cell_masks : ndarray + Cell segmentation masks with unique integers for each cell + cytoplasm_masks : ndarray or None + Cytoplasm segmentation masks. If provided, secondary object + regions will be removed from cytoplasm masks + second_obj_min_size : float + Minimum size for valid secondary objects + second_obj_max_size : float + Maximum size for valid secondary objects + size_filter_method : str + Size filtering method ("feret" or "area") + max_objects_per_cell : int + Maximum secondary objects allowed per cell + overlap_threshold : float + Minimum overlap ratio to associate object with cell (0.0-1.0) + nuclei_centroids : dict, DataFrame, or None + Cell nuclei centroids for distance calculations. + Format: {nuclei_id: (i, j)} or DataFrame with 'i', 'j' columns + max_total_objects : int or None + Failsafe limit on detected objects. Returns empty results if exceeded + image : ndarray, optional + Multichannel image [channels, height, width]. + Only needed if nuclei_centroids provided (for distance calculations) + second_obj_channel_index : int, optional + Index of secondary object channel. + Only needed if nuclei_centroids provided (for distance calculations) + + Returns: + ------- + tuple + - second_obj_masks: Filtered and renumbered secondary object masks + - cell_second_obj_table: Dict with 'cell_summary' and 'second_obj_cell_mapping' DataFrames + - updated_cytoplasm_masks: Cytoplasm masks with secondary objects removed (or None) + + Notes: + ----- + - This function is shared by both segment_second_objs() and segment_second_objs_ml() + - Input masks should already be labeled (not binary) + - Empty input masks are handled gracefully + """ + # Handle empty input + if not np.any(second_obj_masks): + print("No objects detected in input masks") + return create_empty_results(cell_masks, cytoplasm_masks, nuclei_centroids) + + # Failsafe: Check for excessive objects early + num_input_objects = len(np.unique(second_obj_masks)) - 1 # Exclude background + if max_total_objects is not None and num_input_objects > max_total_objects: + print( + f"Failsafe triggered: Detected {num_input_objects} objects (limit: {max_total_objects})" + ) + print("Returning empty results to avoid processing over-segmented image") + return create_empty_results(cell_masks, cytoplasm_masks, nuclei_centroids) + + # Filter by size + print(f"Filtering by {size_filter_method}...") + regions = measure.regionprops(second_obj_masks) + valid_labels = [] + + if size_filter_method == "feret": + # Feret diameter filtering + for region in regions: + coords = region.coords[:, [1, 0]] # (x, y) format + if len(coords) < 3: + continue + + feret_min, feret_max = get_feret_diameters(coords) + if second_obj_min_size <= feret_min and feret_max <= second_obj_max_size: + valid_labels.append(region.label) + + elif size_filter_method == "area": + # Area-based filtering (CellProfiler standard) + for region in regions: + if second_obj_min_size <= region.area <= second_obj_max_size: + valid_labels.append(region.label) + + else: + raise ValueError(f"Unknown size_filter_method: {size_filter_method}") + + if not valid_labels: + print(f"No valid secondary objects found after {size_filter_method} filtering") + return create_empty_results(cell_masks, cytoplasm_masks, nuclei_centroids) + + print( + f"After {size_filter_method} filtering: {len(valid_labels)} valid secondary objects" + ) + + # Create valid secondary objects mask with renumbered labels + labeled_second_objs = np.zeros_like(second_obj_masks) + for i, lbl in enumerate(valid_labels, start=1): + labeled_second_objs[second_obj_masks == lbl] = i + + num_second_objs = len(valid_labels) + + # Get cell IDs + cell_ids = np.unique(cell_masks[cell_masks > 0]) + + # Prepare nuclei centroids - this is for cell nuclei distance calculations + nuclei_centroids_dict = None + if nuclei_centroids is not None: + if isinstance(nuclei_centroids, pd.DataFrame): + nuclei_centroids_dict = { + row.get("nuclei_id", idx): (row["i"], row["j"]) + for idx, row in nuclei_centroids.iterrows() + } + else: + nuclei_centroids_dict = nuclei_centroids + + # Pre-compute region properties for all secondary objects + second_obj_regions = { + region.label: region for region in measure.regionprops(labeled_second_objs) + } + + # Pre-compute which cells could overlap with each secondary object + print("Computing spatial overlap candidates...") + overlap_candidates = get_spatial_overlap_candidates(second_obj_regions, cell_masks) + + # Initialize tracking variables + second_obj_cell_mapping = [] + second_objs_per_cell = {cell_id: 0 for cell_id in cell_ids} + + # Process each secondary object + print("Processing secondary object-cell associations...") + for second_obj_id in range(1, num_second_objs + 1): + if second_obj_id not in second_obj_regions: + continue + + region = second_obj_regions[second_obj_id] + second_obj_mask = labeled_second_objs == second_obj_id + second_obj_area = region.area + second_obj_centroid = region.centroid + + # Calculate equivalent diameter for this secondary object + second_obj_diameter = 2 * np.sqrt(second_obj_area / np.pi) + + # Initialize mapping entry with basic info + mapping_entry = { + "second_obj_id": second_obj_id, + "second_obj_area": second_obj_area, + "second_obj_diameter": second_obj_diameter, + } + + # Calculate distance to nearest cell nucleus + if nuclei_centroids_dict is not None: + min_dist = np.inf + nearest_nucleus_id = None + for nuc_id, nuc_centroid in nuclei_centroids_dict.items(): + dist = np.sqrt( + (second_obj_centroid[0] - nuc_centroid[0]) ** 2 + + (second_obj_centroid[1] - nuc_centroid[1]) ** 2 + ) + if dist < min_dist: + min_dist = dist + nearest_nucleus_id = nuc_id + + mapping_entry["distance_to_nucleus"] = ( + min_dist if min_dist != np.inf else None + ) + mapping_entry["nearest_nucleus_id"] = nearest_nucleus_id + + # Find best overlapping cell + best_cell_id = None + best_overlap = 0 + + # Check spatial overlap candidates + candidate_cells = overlap_candidates.get(second_obj_id, []) + + for cell_id in candidate_cells: + if second_objs_per_cell[cell_id] >= max_objects_per_cell: + continue + + # Calculate overlap efficiently + cell_mask = cell_masks == cell_id + overlap = np.sum(second_obj_mask & cell_mask) + + if overlap > 0: + overlap_ratio = overlap / second_obj_area + if overlap_ratio >= overlap_threshold and overlap_ratio > best_overlap: + best_overlap = overlap_ratio + best_cell_id = cell_id + + # Add successful associations + if best_cell_id is not None: + mapping_entry["cell_id"] = best_cell_id + mapping_entry["overlap_ratio"] = best_overlap + + second_obj_cell_mapping.append(mapping_entry) + second_objs_per_cell[best_cell_id] += 1 + + # Create secondary object-cell mapping DataFrame + second_obj_cell_df = pd.DataFrame(second_obj_cell_mapping) + + # Create cell summary + if second_obj_cell_mapping: + # Group by cell_id once for efficiency + grouped = second_obj_cell_df.groupby("cell_id") + cell_summary = [] + + for cell_id in cell_ids: + cell_area = np.sum(cell_masks == cell_id) + + # Initialize basic cell summary + summary_entry = { + "cell_id": cell_id, + "cell_area": cell_area, + } + + # Check if cell_id has associated secondary objects + if cell_id in grouped.groups: + cell_second_objs = grouped.get_group(cell_id) + + # Calculate cell-level statistics + total_second_obj_area = cell_second_objs["second_obj_area"].sum() + mean_diameter = cell_second_objs["second_obj_diameter"].mean() + + summary_entry.update( + { + "has_second_obj": True, + "num_second_objs": len(cell_second_objs), + "second_obj_ids": list(cell_second_objs["second_obj_id"]), + "total_second_obj_area": total_second_obj_area, + "second_obj_area_ratio": total_second_obj_area / cell_area + if cell_area > 0 + else 0, + "mean_second_obj_diameter": mean_diameter, + } + ) + + # Add cell nucleus distance fields if nuclei_centroids was provided + if nuclei_centroids_dict is not None: + mean_distance = ( + cell_second_objs["distance_to_nucleus"].dropna().mean() + if not cell_second_objs["distance_to_nucleus"].dropna().empty + else None + ) + summary_entry["mean_distance_to_nucleus"] = mean_distance + + else: # Cell without secondary objects + summary_entry.update( + { + "has_second_obj": False, + "num_second_objs": 0, + "second_obj_ids": [], + "total_second_obj_area": 0, + "second_obj_area_ratio": 0, + "mean_second_obj_diameter": None, + } + ) + + # Add cell nucleus distance fields if nuclei_centroids was provided + if nuclei_centroids_dict is not None: + summary_entry["mean_distance_to_nucleus"] = None + + cell_summary.append(summary_entry) + + else: + # Handle case with no secondary objects + cell_summary = [] + for cell_id in cell_ids: + cell_area = np.sum(cell_masks == cell_id) + summary_entry = { + "cell_id": cell_id, + "has_second_obj": False, + "num_second_objs": 0, + "second_obj_ids": [], + "cell_area": cell_area, + "total_second_obj_area": 0, + "second_obj_area_ratio": 0, + "mean_second_obj_diameter": None, + } + + # Add cell nucleus distance fields if nuclei_centroids was provided + if nuclei_centroids_dict is not None: + summary_entry["mean_distance_to_nucleus"] = None + + cell_summary.append(summary_entry) + + # Create final results + cell_summary_df = pd.DataFrame(cell_summary) + cell_second_obj_table = { + "cell_summary": cell_summary_df, + "second_obj_cell_mapping": second_obj_cell_df, + } + + # Create associated secondary object masks + associated_second_objs = np.zeros_like(labeled_second_objs) + for mapping in second_obj_cell_mapping: + second_obj_id = mapping["second_obj_id"] + second_obj_mask = labeled_second_objs == second_obj_id + associated_second_objs[second_obj_mask] = second_obj_id + + # Print statistics + total_kept = len(second_obj_cell_mapping) + print( + f"Kept {total_kept} out of {num_second_objs} detected secondary objects " + f"({total_kept / num_second_objs * 100:.1f}%)" + ) + print( + f"Discarded {num_second_objs - total_kept} secondary objects that didn't meet diameter criteria or cell overlap" + ) + + # Process cytoplasm masks if provided + updated_cytoplasm_masks = None + if cytoplasm_masks is not None: + updated_cytoplasm_masks = cytoplasm_masks.copy() + for mapping in second_obj_cell_mapping: + second_obj_id = mapping["second_obj_id"] + cell_id = mapping["cell_id"] + second_obj_mask = associated_second_objs == second_obj_id + cytoplasm_mask = updated_cytoplasm_masks == cell_id + updated_cytoplasm_masks[cytoplasm_mask & second_obj_mask] = 0 + print( + f"Updated cytoplasm masks by removing {len(second_obj_cell_mapping)} secondary object regions" + ) + + # Return results + if updated_cytoplasm_masks is not None: + return ( + associated_second_objs, + cell_second_obj_table, + updated_cytoplasm_masks, + ) + else: + return associated_second_objs, cell_second_obj_table diff --git a/workflow/lib/sbs/align_cycles.py b/workflow/lib/sbs/align_cycles.py index cbf5b895..a1ea5a61 100644 --- a/workflow/lib/sbs/align_cycles.py +++ b/workflow/lib/sbs/align_cycles.py @@ -286,7 +286,7 @@ def align_it(x): target = apply_window(aligned[:, sbs_channels], window=window).max(axis=1) normed = normalize_by_percentile(target, q_norm=q_norm) normed[normed > cutoff] = cutoff - offsets = calculate_offsets(normed, upsample_factor=upsample_factor) + offsets, _ = calculate_offsets(normed, upsample_factor=upsample_factor) if verbose: print("\n=== Cycle Alignment Offsets (sbs_mean method) ===") @@ -333,7 +333,7 @@ def align_within_cycle(data_, upsample_factor=4, window=1, q1=0, q2=90): # Filter the input data based on percentiles filtered = filter_percentiles(apply_window(data_, window), q1=q1, q2=q2) # Calculate offsets using the filtered data - offsets = calculate_offsets(filtered, upsample_factor=upsample_factor) + offsets, _ = calculate_offsets(filtered, upsample_factor=upsample_factor) # Apply the calculated offsets to the original data and return the result return apply_offsets(data_, offsets) @@ -356,7 +356,7 @@ def align_between_cycles( """ # Calculate offsets from the target channel target = apply_window(data[:, channel_index], window) - offsets = calculate_offsets(target, upsample_factor=upsample_factor) + offsets, _ = calculate_offsets(target, upsample_factor=upsample_factor) # Apply the calculated offsets to all channels warped = [] diff --git a/workflow/lib/shared/align.py b/workflow/lib/shared/align.py index d1aaca8d..3b0ce333 100644 --- a/workflow/lib/shared/align.py +++ b/workflow/lib/shared/align.py @@ -1,8 +1,6 @@ """Shared functions for aligning images. -Uses NumPy and scikit-image to provide image -alignment between sequencing cycles, apply percentile-based filtering, fill masked -areas with noise, and perform various transformations to enhance image data quality. +Uses NumPy and scikit-image to provide image alignment between sequencing cycles. """ import numpy as np @@ -65,26 +63,31 @@ def calculate_offsets(data_, upsample_factor): upsample_factor (int): Upsampling factor for cross-correlation. Returns: - np.ndarray: Offset values between images. + tuple: (offsets, errors) where offsets is np.ndarray of offset values + and errors is np.ndarray of correlation errors (0 = perfect correlation). """ # Set the target frame as the first frame in the data target = data_[0] - # Initialize an empty list to store offsets + + # Initialize empty lists to store offsets and errors offsets = [] + errors = [] # Iterate through each frame in the data for i, src in enumerate(data_): - # If it's the first frame, add a zero offset + # If it's the first frame, add a zero offset and zero error if i == 0: offsets += [(0, 0)] + errors += [0.0] else: # Calculate the offset between the current frame and the target frame - offset, _, _ = skimage.registration.phase_cross_correlation( + offset, error, _ = skimage.registration.phase_cross_correlation( src, target, upsample_factor=upsample_factor, normalization=None ) - # Add the offset to the list + # Add the offset and error to the lists offsets += [offset] - # Convert the list of offsets to a numpy array and return - return np.array(offsets) + errors += [error] + # Convert the lists to numpy arrays and return + return np.array(offsets), np.array(errors) @applyIJ diff --git a/workflow/lib/shared/cellpose_training.py b/workflow/lib/shared/cellpose_training.py new file mode 100644 index 00000000..cbbd773e --- /dev/null +++ b/workflow/lib/shared/cellpose_training.py @@ -0,0 +1,776 @@ +"""Cellpose Fine-Tuning Utilities. + +Functions for loading training data, augmentation, training, evaluation, +and visualization of Cellpose models for secondary object segmentation. +""" + +from pathlib import Path +from typing import List, Tuple, Optional, Union, Dict +import numpy as np +from tifffile import imread +from cellpose import models, train, io +from skimage.transform import rotate +from skimage.measure import regionprops, label +import matplotlib.pyplot as plt +from matplotlib.colors import ListedColormap +import random + +# Reuse utilities from segment_cellpose (single source of truth) +from lib.shared.segment_cellpose import prepare_cellpose, create_cellpose_model + + +def load_training_data( + image_paths: List[Union[str, Path]], + mask_paths: List[Union[str, Path]], + mode: str = "secondary_obj", + channel_index: Optional[int] = None, + dapi_index: Optional[int] = None, + cyto_index: Optional[int] = None, + helper_index: Optional[int] = None, + logscale: bool = True, +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """Load paired images and masks for Cellpose training. + + Supports different preprocessing modes to match deployment functions: + - "secondary_obj": For segment_second_objs_ml (single channel, log scaling) + - "cells": For segment_cellpose with cells=True (3-channel RGB) + - "nuclei": For segment_cellpose with cells=False (DAPI only) + + Parameters + ---------- + image_paths : List[str | Path] + Paths to image files (TIFF format, can be multi-channel). + mask_paths : List[str | Path] + Paths to mask files (NPY format, labeled masks where each object + has a unique integer ID and background is 0). + mode : str + Preprocessing mode. Options: + - "secondary_obj": For segment_second_objs_ml (requires channel_index) + - "cells": For segment_cellpose cells (requires dapi_index, cyto_index) + - "nuclei": For segment_cellpose nuclei only (requires dapi_index) + channel_index : int, optional + Channel index for secondary_obj mode. + dapi_index : int, optional + DAPI channel index for cells/nuclei modes. + cyto_index : int, optional + Cytoplasm channel index for cells mode. + helper_index : int, optional + Helper channel index for cells mode (optional). + logscale : bool + Apply log scaling preprocessing. Default True. + + Returns: + ------- + images : List[np.ndarray] + List of preprocessed image arrays (uint8). + - secondary_obj/nuclei: 2D arrays [height, width] + - cells: 3D arrays [3, height, width] + masks : List[np.ndarray] + List of 2D labeled mask arrays (int32). + + Raises: + ------ + ValueError + If number of images and masks don't match, required indices not provided, + or if dimensions mismatch. + """ + if len(image_paths) != len(mask_paths): + raise ValueError( + f"Number of images ({len(image_paths)}) must match " + f"number of masks ({len(mask_paths)})" + ) + + # Validate mode and required parameters + if mode == "secondary_obj": + if channel_index is None: + raise ValueError("channel_index is required for mode='secondary_obj'") + elif mode == "cells": + if dapi_index is None or cyto_index is None: + raise ValueError("dapi_index and cyto_index are required for mode='cells'") + elif mode == "nuclei": + if dapi_index is None: + raise ValueError("dapi_index is required for mode='nuclei'") + else: + raise ValueError( + f"Unknown mode: {mode}. Valid options: 'secondary_obj', 'cells', 'nuclei'" + ) + + images = [] + masks = [] + + for img_path, mask_path in zip(image_paths, mask_paths): + # Load image + img = imread(str(img_path)) + + # Apply mode-specific preprocessing using prepare_cellpose (single source of truth) + if mode == "secondary_obj": + # Use prepare_cellpose with target channel as cyto, extract green channel + # Green channel (index 1) has log scaling + max normalization + rgb = prepare_cellpose( + img, + dapi_index=channel_index, # Dummy - will use cyto + cyto_index=channel_index, # Target channel + helper_index=None, + logscale=logscale, + ) + processed_img = rgb[1] # Extract green (log scaled + normalized) + + elif mode == "cells": + # Full RGB output from prepare_cellpose + processed_img = prepare_cellpose( + img, + dapi_index=dapi_index, + cyto_index=cyto_index, + helper_index=helper_index, + logscale=logscale, + ) + + elif mode == "nuclei": + # Use prepare_cellpose and extract DAPI (blue) channel + rgb = prepare_cellpose( + img, + dapi_index=dapi_index, + cyto_index=dapi_index, # Dummy + helper_index=None, + logscale=False, # DAPI uses percentile norm, not log scale + ) + processed_img = rgb[2] # Extract blue (DAPI with percentile norm) + + # Load mask (support both .npy and .tif/.tiff formats) + mask_path_str = str(mask_path) + if mask_path_str.endswith(".npy"): + mask = np.load(mask_path_str) + else: + mask = imread(mask_path_str) + if mask.ndim > 2: + raise ValueError(f"Mask at {mask_path} should be 2D, got {mask.ndim}D") + + # Validate dimensions match (compare 2D shapes) + img_shape_2d = ( + processed_img.shape[-2:] if processed_img.ndim > 2 else processed_img.shape + ) + if img_shape_2d != mask.shape: + raise ValueError( + f"Image shape {img_shape_2d} doesn't match mask shape {mask.shape} " + f"for {img_path}" + ) + + images.append(processed_img) + masks.append(mask.astype(np.int32)) + + print(f"Loaded {len(images)} image-mask pairs (mode={mode})") + return images, masks + + +def augment_training_data( + images: List[np.ndarray], + masks: List[np.ndarray], + rotations: bool = True, + flips: bool = True, + intensity_scaling: bool = True, + intensity_range: Tuple[float, float] = (0.8, 1.2), + noise: bool = False, + noise_std: float = 0.02, +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """Apply data augmentation to expand small training datasets. + + For a dataset of N images, this can produce up to 8N augmented samples + (4 rotations x 2 flip states). + + Parameters + ---------- + images : List[np.ndarray] + List of 2D image arrays. + masks : List[np.ndarray] + List of 2D labeled mask arrays. + rotations : bool + Apply 90, 180, 270 degree rotations. + flips : bool + Apply horizontal and vertical flips. + intensity_scaling : bool + Apply random intensity scaling. + intensity_range : Tuple[float, float] + Range for intensity scaling factor. + noise : bool + Add Gaussian noise. + noise_std : float + Standard deviation of Gaussian noise. + + Returns: + ------- + aug_images : List[np.ndarray] + Augmented images (includes originals). + aug_masks : List[np.ndarray] + Augmented masks (includes originals). + """ + aug_images = [] + aug_masks = [] + + for img, mask in zip(images, masks): + # Start with original + variants = [(img.copy(), mask.copy())] + + # Rotations (90, 180, 270 degrees) + if rotations: + for k in [1, 2, 3]: # k*90 degrees + rot_img = np.rot90(img, k) + rot_mask = np.rot90(mask, k) + variants.append((rot_img.copy(), rot_mask.copy())) + + # Flips + if flips: + current_variants = variants.copy() + for v_img, v_mask in current_variants: + # Horizontal flip + flip_img = np.fliplr(v_img) + flip_mask = np.fliplr(v_mask) + variants.append((flip_img.copy(), flip_mask.copy())) + + # Intensity modifications (applied to each variant) + final_variants = [] + for v_img, v_mask in variants: + # Original intensity + final_variants.append((v_img.copy(), v_mask.copy())) + + # Intensity scaling + if intensity_scaling: + scale = random.uniform(intensity_range[0], intensity_range[1]) + scaled_img = np.clip(v_img * scale, 0, 1) + final_variants.append((scaled_img.astype(np.float32), v_mask.copy())) + + # Noise + if noise: + noisy_img = v_img + np.random.normal(0, noise_std, v_img.shape) + noisy_img = np.clip(noisy_img, 0, 1) + final_variants.append((noisy_img.astype(np.float32), v_mask.copy())) + + for v_img, v_mask in final_variants: + aug_images.append(v_img) + aug_masks.append(v_mask) + + print( + f"Augmentation: {len(images)} original → {len(aug_images)} samples " + f"({len(aug_images) / len(images):.1f}x expansion)" + ) + return aug_images, aug_masks + + +def prepare_cellpose_training( + images: List[np.ndarray], + masks: List[np.ndarray], + test_fraction: float = 0.1, + seed: int = 42, +) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], List[np.ndarray]]: + """Split data into training and test sets for Cellpose. + + Parameters + ---------- + images : List[np.ndarray] + List of 2D image arrays. + masks : List[np.ndarray] + List of 2D labeled mask arrays. + test_fraction : float + Fraction of data to use for testing (0.0 to 1.0). + seed : int + Random seed for reproducibility. + + Returns: + ------- + train_images : List[np.ndarray] + train_masks : List[np.ndarray] + test_images : List[np.ndarray] + test_masks : List[np.ndarray] + """ + n_samples = len(images) + n_test = max(1, int(n_samples * test_fraction)) + + # Shuffle indices + np.random.seed(seed) + indices = np.random.permutation(n_samples) + + test_indices = indices[:n_test] + train_indices = indices[n_test:] + + train_images = [images[i] for i in train_indices] + train_masks = [masks[i] for i in train_indices] + test_images = [images[i] for i in test_indices] + test_masks = [masks[i] for i in test_indices] + + print(f"Split: {len(train_images)} training, {len(test_images)} test samples") + return train_images, train_masks, test_images, test_masks + + +def train_cellpose( + train_images: List[np.ndarray], + train_masks: List[np.ndarray], + test_images: Optional[List[np.ndarray]] = None, + test_masks: Optional[List[np.ndarray]] = None, + base_model: str = "cpsam", + n_epochs: int = 500, + learning_rate: float = 0.1, + weight_decay: float = 1e-5, + batch_size: int = 8, + save_path: Union[str, Path] = "models", + model_name: str = "cpsam_secondary_obj", + gpu: bool = True, + channels: List[int] = None, +) -> models.CellposeModel: + """Fine-tune a Cellpose model on custom training data. + + Parameters + ---------- + train_images : List[np.ndarray] + Training images (2D arrays). + train_masks : List[np.ndarray] + Training masks (labeled 2D arrays). + test_images : List[np.ndarray], optional + Test images for validation during training. + test_masks : List[np.ndarray], optional + Test masks for validation. + base_model : str + Base model to fine-tune from. Options: "cpsam", "cyto3", "cyto2", "nuclei". + n_epochs : int + Number of training epochs. + learning_rate : float + Initial learning rate. + weight_decay : float + L2 regularization weight. + batch_size : int + Training batch size. + save_path : str | Path + Directory to save trained model. + model_name : str + Name for the saved model. + gpu : bool + Use GPU acceleration if available. + channels : List[int], optional + Channel configuration for Cellpose. Default [0, 0] for grayscale. + + Returns: + ------- + model : CellposeModel + Trained Cellpose model. + """ + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + + if channels is None: + channels = [0, 0] # Grayscale + + # Set up logging to see training progress + io.logger_setup() + + print(f"Initializing model from base: {base_model}") + print( + f"Training parameters: epochs={n_epochs}, lr={learning_rate}, batch={batch_size}" + ) + + # Initialize model with version-aware helper (validates model compatibility) + model = create_cellpose_model(base_model, gpu=gpu) + + print(f"Starting training with {len(train_images)} samples...") + + # In Cellpose 3.0+, use train.train_seg() instead of model.train() + # Note: channels is not a parameter for train_seg - images should be pre-formatted + model_path, train_losses, test_losses = train.train_seg( + model.net, + train_data=train_images, + train_labels=train_masks, + test_data=test_images, + test_labels=test_masks, + save_path=str(save_path), + n_epochs=n_epochs, + learning_rate=learning_rate, + weight_decay=weight_decay, + batch_size=batch_size, + model_name=model_name, + ) + + print(f"Training complete. Model saved to: {model_path}") + + # Load and return the trained model + trained_model = models.CellposeModel(gpu=gpu, pretrained_model=model_path) + return trained_model + + +def load_trained_model( + model_path: Union[str, Path], + gpu: bool = True, +) -> models.CellposeModel: + """Load a fine-tuned Cellpose model. + + Parameters + ---------- + model_path : str | Path + Path to saved model file. + gpu : bool + Use GPU acceleration. + + Returns: + ------- + model : CellposeModel + Loaded Cellpose model ready for inference. + """ + model = models.CellposeModel(gpu=gpu, pretrained_model=str(model_path)) + print(f"Loaded model from: {model_path}") + return model + + +def predict_masks( + model: models.CellposeModel, + images: List[np.ndarray], + diameter: Optional[float] = None, + flow_threshold: float = 0.4, + cellprob_threshold: float = 0.0, + channels: List[int] = None, +) -> List[np.ndarray]: + """Run inference on images using a Cellpose model. + + Parameters + ---------- + model : CellposeModel + Cellpose model (base or fine-tuned). + images : List[np.ndarray] + List of 2D images to segment. + diameter : float, optional + Expected object diameter. None for auto-estimation. + flow_threshold : float + Flow error threshold. + cellprob_threshold : float + Cell probability threshold. + channels : List[int], optional + Channel configuration. Default [0, 0] for grayscale. + + Returns: + ------- + masks : List[np.ndarray] + Predicted segmentation masks. + """ + if channels is None: + channels = [0, 0] + + masks, flows, styles = model.eval( + images, + diameter=diameter, + channels=channels, + flow_threshold=flow_threshold, + cellprob_threshold=cellprob_threshold, + ) + + return masks + + +def calculate_iou(pred_mask: np.ndarray, gt_mask: np.ndarray) -> float: + """Calculate Intersection over Union between predicted and ground truth masks. + + This computes the average IoU across all objects. + + Parameters + ---------- + pred_mask : np.ndarray + Predicted labeled mask. + gt_mask : np.ndarray + Ground truth labeled mask. + + Returns: + ------- + iou : float + Mean IoU score (0 to 1). + """ + pred_binary = pred_mask > 0 + gt_binary = gt_mask > 0 + + intersection = np.logical_and(pred_binary, gt_binary).sum() + union = np.logical_or(pred_binary, gt_binary).sum() + + if union == 0: + return 1.0 if intersection == 0 else 0.0 + + return intersection / union + + +def calculate_object_metrics( + pred_mask: np.ndarray, + gt_mask: np.ndarray, + iou_threshold: float = 0.5, +) -> Dict[str, float]: + """Calculate object-level metrics (precision, recall, F1). + + An object is considered a true positive if it overlaps with a ground truth + object with IoU >= threshold. + + Parameters + ---------- + pred_mask : np.ndarray + Predicted labeled mask. + gt_mask : np.ndarray + Ground truth labeled mask. + iou_threshold : float + IoU threshold for matching objects. + + Returns: + ------- + metrics : dict + Dictionary with 'precision', 'recall', 'f1', 'n_pred', 'n_gt', 'n_tp'. + """ + pred_labels = np.unique(pred_mask[pred_mask > 0]) + gt_labels = np.unique(gt_mask[gt_mask > 0]) + + n_pred = len(pred_labels) + n_gt = len(gt_labels) + + if n_pred == 0 and n_gt == 0: + return { + "precision": 1.0, + "recall": 1.0, + "f1": 1.0, + "n_pred": 0, + "n_gt": 0, + "n_tp": 0, + } + + if n_pred == 0: + return { + "precision": 0.0, + "recall": 0.0, + "f1": 0.0, + "n_pred": 0, + "n_gt": n_gt, + "n_tp": 0, + } + + if n_gt == 0: + return { + "precision": 0.0, + "recall": 0.0, + "f1": 0.0, + "n_pred": n_pred, + "n_gt": 0, + "n_tp": 0, + } + + # Match predictions to ground truth + matched_gt = set() + tp = 0 + + for pred_label in pred_labels: + pred_region = pred_mask == pred_label + best_iou = 0 + best_gt = None + + for gt_label in gt_labels: + if gt_label in matched_gt: + continue + gt_region = gt_mask == gt_label + + intersection = np.logical_and(pred_region, gt_region).sum() + union = np.logical_or(pred_region, gt_region).sum() + iou = intersection / union if union > 0 else 0 + + if iou > best_iou: + best_iou = iou + best_gt = gt_label + + if best_iou >= iou_threshold and best_gt is not None: + tp += 1 + matched_gt.add(best_gt) + + precision = tp / n_pred if n_pred > 0 else 0 + recall = tp / n_gt if n_gt > 0 else 0 + f1 = ( + 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 + ) + + return { + "precision": precision, + "recall": recall, + "f1": f1, + "n_pred": n_pred, + "n_gt": n_gt, + "n_tp": tp, + } + + +def evaluate_segmentation( + model: models.CellposeModel, + images: List[np.ndarray], + gt_masks: List[np.ndarray], + diameter: Optional[float] = None, + flow_threshold: float = 0.4, + cellprob_threshold: float = 0.0, + iou_threshold: float = 0.5, +) -> Dict[str, float]: + """Evaluate model performance on a test set. + + Parameters + ---------- + model : CellposeModel + Cellpose model to evaluate. + images : List[np.ndarray] + Test images. + gt_masks : List[np.ndarray] + Ground truth masks. + diameter : float, optional + Object diameter for inference. + flow_threshold : float + Flow threshold for inference. + cellprob_threshold : float + Cell probability threshold. + iou_threshold : float + IoU threshold for object matching. + + Returns: + ------- + metrics : dict + Aggregated metrics: mean_iou, mean_precision, mean_recall, mean_f1. + """ + pred_masks = predict_masks( + model, + images, + diameter=diameter, + flow_threshold=flow_threshold, + cellprob_threshold=cellprob_threshold, + ) + + ious = [] + precisions = [] + recalls = [] + f1s = [] + + for pred, gt in zip(pred_masks, gt_masks): + iou = calculate_iou(pred, gt) + metrics = calculate_object_metrics(pred, gt, iou_threshold) + + ious.append(iou) + precisions.append(metrics["precision"]) + recalls.append(metrics["recall"]) + f1s.append(metrics["f1"]) + + return { + "mean_iou": np.mean(ious), + "mean_precision": np.mean(precisions), + "mean_recall": np.mean(recalls), + "mean_f1": np.mean(f1s), + "per_image_iou": ious, + "per_image_precision": precisions, + "per_image_recall": recalls, + "per_image_f1": f1s, + } + + +def random_label_cmap(n_labels: int = 256, seed: int = 42) -> ListedColormap: + """Create a random colormap for labeled masks.""" + np.random.seed(seed) + colors = np.random.rand(n_labels, 3) + colors[0] = [0, 0, 0] # Background is black + return ListedColormap(colors) + + +def visualize_comparison( + image: np.ndarray, + pred_mask: np.ndarray, + gt_mask: np.ndarray, + title: str = "", + figsize: Tuple[int, int] = (15, 5), +) -> plt.Figure: + """Visualize side-by-side comparison of prediction vs ground truth. + + Parameters + ---------- + image : np.ndarray + Original image. + pred_mask : np.ndarray + Predicted segmentation mask. + gt_mask : np.ndarray + Ground truth mask. + title : str + Figure title. + figsize : Tuple[int, int] + Figure size. + + Returns: + ------- + fig : plt.Figure + Matplotlib figure. + """ + fig, axes = plt.subplots(1, 4, figsize=figsize) + + # Original image + axes[0].imshow(image, cmap="gray") + axes[0].set_title("Original Image") + axes[0].axis("off") + + # Ground truth + cmap = random_label_cmap(max(gt_mask.max(), pred_mask.max()) + 1) + axes[1].imshow(gt_mask, cmap=cmap, interpolation="nearest") + axes[1].set_title(f"Ground Truth ({gt_mask.max()} objects)") + axes[1].axis("off") + + # Prediction + axes[2].imshow(pred_mask, cmap=cmap, interpolation="nearest") + axes[2].set_title(f"Prediction ({pred_mask.max()} objects)") + axes[2].axis("off") + + # Overlay + overlay = np.zeros((*image.shape, 3)) + overlay[..., 0] = image # Red channel = image + overlay[..., 1] = (gt_mask > 0).astype(float) * 0.5 # Green = GT + overlay[..., 2] = (pred_mask > 0).astype(float) * 0.5 # Blue = prediction + axes[3].imshow(np.clip(overlay, 0, 1)) + axes[3].set_title("Overlay (G=GT, B=Pred)") + axes[3].axis("off") + + # Calculate metrics + iou = calculate_iou(pred_mask, gt_mask) + metrics = calculate_object_metrics(pred_mask, gt_mask) + + fig.suptitle( + f"{title}\nIoU: {iou:.3f} | Precision: {metrics['precision']:.3f} | " + f"Recall: {metrics['recall']:.3f} | F1: {metrics['f1']:.3f}", + fontsize=12, + ) + + plt.tight_layout() + return fig + + +def visualize_training_sample( + image: np.ndarray, + mask: np.ndarray, + title: str = "", + figsize: Tuple[int, int] = (10, 5), +) -> plt.Figure: + """Visualize a single training sample (image + mask). + + Parameters + ---------- + image : np.ndarray + Training image. + mask : np.ndarray + Corresponding mask. + title : str + Figure title. + figsize : Tuple[int, int] + Figure size. + + Returns: + ------- + fig : plt.Figure + Matplotlib figure. + """ + fig, axes = plt.subplots(1, 2, figsize=figsize) + + axes[0].imshow(image, cmap="gray") + axes[0].set_title("Image") + axes[0].axis("off") + + cmap = random_label_cmap(mask.max() + 1) + axes[1].imshow(mask, cmap=cmap, interpolation="nearest") + axes[1].set_title(f"Mask ({mask.max()} objects)") + axes[1].axis("off") + + if title: + fig.suptitle(title) + + plt.tight_layout() + return fig diff --git a/workflow/lib/shared/file_utils.py b/workflow/lib/shared/file_utils.py index 0e9a6b4c..0f6e82ca 100644 --- a/workflow/lib/shared/file_utils.py +++ b/workflow/lib/shared/file_utils.py @@ -16,6 +16,7 @@ "round": ["R-", str], "cell_class": ["CeCl-", str], "channel_combo": ["ChCo-", str], + "compartment_combo": ["CmCo-", str], "gene": ["G-", str], "sgrna": ["SG-", str], "channel": ["CH-", str], @@ -242,3 +243,32 @@ def validate_data_type(data_type): if data_type not in valid_types: raise ValueError(f"data_type must be one of {valid_types}, got '{data_type}'") return data_type + + +def read_tsv_safe(filepath, return_empty_on_error=True): + """Read a TSV file safely with error handling. + + Parameters + ---------- + filepath : str or Path + Path to the TSV file to read + return_empty_on_error : bool, optional + If True, return empty DataFrame on EmptyDataError. + If False, return None on EmptyDataError. + Default: True + + Returns: + ------- + pd.DataFrame or None + DataFrame if file read successfully or is empty (when return_empty_on_error=True). + None if file is empty and return_empty_on_error=False. + + Examples: + -------- + >>> df = read_tsv_safe('data.tsv') + >>> df = read_tsv_safe('data.tsv', return_empty_on_error=False) + """ + try: + return pd.read_csv(filepath, sep="\t") + except pd.errors.EmptyDataError: + return pd.DataFrame() if return_empty_on_error else None diff --git a/workflow/lib/shared/metrics.py b/workflow/lib/shared/metrics.py index 1dfb1c4f..b6ba7ae5 100644 --- a/workflow/lib/shared/metrics.py +++ b/workflow/lib/shared/metrics.py @@ -298,7 +298,7 @@ def get_aggregate_stats(config, n_rows=10000, include_batch_effects=False): include_batch_effects: Whether to calculate batch effect metrics (slow) Returns: - dict: Dictionary with statistics for each cell_class/channel_combo combination + dict: Dictionary with statistics for each cell_class/channel_combo/compartment_combo combination """ root_fp = Path(config["all"]["root_fp"]) aggregate_dir = root_fp / "aggregate" @@ -307,8 +307,10 @@ def get_aggregate_stats(config, n_rows=10000, include_batch_effects=False): aggregate_combo_fp = Path(config["aggregate"]["aggregate_combo_fp"]) aggregate_combos = pd.read_csv(aggregate_combo_fp, sep="\t") - # Get unique cell_class and channel_combo combinations - unique_combos = aggregate_combos[["cell_class", "channel_combo"]].drop_duplicates() + # Get unique cell_class, channel_combo, and compartment_combo combinations + unique_combos = aggregate_combos[ + ["cell_class", "channel_combo", "compartment_combo"] + ].drop_duplicates() # Get perturbation column name from config perturbation_col = config["aggregate"].get("perturbation_name_col", "gene_symbol_0") @@ -319,12 +321,14 @@ def get_aggregate_stats(config, n_rows=10000, include_batch_effects=False): for _, combo in unique_combos.iterrows(): cell_class = combo["cell_class"] channel_combo = combo["channel_combo"] + compartment_combo = combo["compartment_combo"] try: result = _get_single_aggregate_stats( config, cell_class, channel_combo, + compartment_combo, n_rows, root_fp, aggregate_dir, @@ -332,10 +336,14 @@ def get_aggregate_stats(config, n_rows=10000, include_batch_effects=False): control_key, include_batch_effects, ) - all_results[f"{cell_class}_{channel_combo}"] = result + all_results[f"{cell_class}_{channel_combo}_{compartment_combo}"] = result except Exception as e: - print(f"Error processing {cell_class}/{channel_combo}: {str(e)}") - all_results[f"{cell_class}_{channel_combo}"] = {"error": str(e)} + print( + f"Error processing {cell_class}/{channel_combo}/{compartment_combo}: {str(e)}" + ) + all_results[f"{cell_class}_{channel_combo}_{compartment_combo}"] = { + "error": str(e) + } return all_results @@ -344,6 +352,7 @@ def _get_single_aggregate_stats( config, cell_class, channel_combo, + compartment_combo, n_rows, root_fp, aggregate_dir, @@ -351,14 +360,14 @@ def _get_single_aggregate_stats( control_key, include_batch_effects, ): - """Helper function to get stats for a single cell_class/channel_combo combination.""" + """Helper function to get stats for a single cell_class/channel_combo/compartment_combo combination.""" from lib.shared.file_utils import load_parquet_subset # Load the aggregated TSV file aggregated_path = ( aggregate_dir / "tsvs" - / f"CeCl-{cell_class}_ChCo-{channel_combo}__aggregated.tsv" + / f"CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__aggregated.tsv" ) aggregated = pd.read_csv(aggregated_path, sep="\t") @@ -378,7 +387,7 @@ def _get_single_aggregate_stats( merge_data_dir = aggregate_dir / "parquets" merge_data_paths = list( merge_data_dir.glob( - f"*_CeCl-{cell_class}_ChCo-{channel_combo}__merge_data.parquet" + f"*_CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__merge_data.parquet" ) ) @@ -390,7 +399,9 @@ def _get_single_aggregate_stats( # Count filtered cells from parquet metadata (fast) filtered_dir = aggregate_dir / "parquets" filtered_paths = list( - filtered_dir.glob(f"*_CeCl-{cell_class}_ChCo-{channel_combo}__filtered.parquet") + filtered_dir.glob( + f"*_CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__filtered.parquet" + ) ) total_filtered_cells = 0 @@ -414,6 +425,7 @@ def _get_single_aggregate_stats( config, cell_class, channel_combo, + compartment_combo, n_rows, aggregate_dir, perturbation_col, @@ -428,6 +440,7 @@ def _calculate_batch_effects( config, cell_class, channel_combo, + compartment_combo, n_rows, aggregate_dir, perturbation_col, @@ -441,7 +454,9 @@ def _calculate_batch_effects( filtered_dir = aggregate_dir / "parquets" # Use direct path instead of recursive glob for speed filtered_paths = list( - filtered_dir.glob(f"*_CeCl-{cell_class}_ChCo-{channel_combo}__filtered.parquet") + filtered_dir.glob( + f"*_CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__filtered.parquet" + ) ) # Load filtered data - sample from a subset of files to avoid @@ -504,7 +519,7 @@ def _calculate_batch_effects( aligned_path = ( aggregate_dir / "parquets" - / f"CeCl-{cell_class}_ChCo-{channel_combo}__aligned.parquet" + / f"CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__aligned.parquet" ) # Use same sample size as pre-alignment for consistency @@ -565,10 +580,15 @@ def get_cluster_stats(config): for _, row in cluster_combos.iterrows(): cell_class = row["cell_class"] channel_combo = row["channel_combo"] + compartment_combo = row["compartment_combo"] leiden_resolution = row["leiden_resolution"] cluster_specific_dir = ( - cluster_dir / channel_combo / cell_class / str(leiden_resolution) + cluster_dir + / channel_combo + / compartment_combo + / cell_class + / str(leiden_resolution) ) # Path to metrics files @@ -598,6 +618,7 @@ def get_cluster_stats(config): result = { "cell_class": cell_class, "channel_combo": channel_combo, + "compartment_combo": compartment_combo, "leiden_resolution": leiden_resolution, "unique_clusters": unique_clusters, # CORUM metrics @@ -641,7 +662,7 @@ def get_cluster_stats(config): except Exception as e: print( - f"Error reading metrics for {cell_class}/{channel_combo}/{leiden_resolution}: {e}" + f"Error reading metrics for {cell_class}/{channel_combo}/{compartment_combo}/{leiden_resolution}: {e}" ) results_df = pd.DataFrame(results) @@ -813,7 +834,7 @@ def get_all_stats( if not stats["cluster"]["detailed_results"].empty: for _, row in stats["cluster"]["detailed_results"].iterrows(): print( - f"\n {row['cell_class']} - {row['channel_combo']} (resolution={row['leiden_resolution']}):" + f"\n {row['cell_class']} - {row['channel_combo']} - {row['compartment_combo']} (resolution={row['leiden_resolution']}):" ) print(f" - Clusters: {row['unique_clusters']}") print( diff --git a/workflow/lib/shared/rule_utils.py b/workflow/lib/shared/rule_utils.py index d6c38cd5..4339ec64 100644 --- a/workflow/lib/shared/rule_utils.py +++ b/workflow/lib/shared/rule_utils.py @@ -313,6 +313,7 @@ def get_bootstrap_inputs( gene_pvals_pattern: Union[str, Path], cell_class: str, channel_combo: str, + compartment_combo: str, ) -> List[str]: """Get all bootstrap inputs for completion flag. @@ -324,13 +325,16 @@ def get_bootstrap_inputs( gene_pvals_pattern (Union[str, Path]): Template string for gene p-value files. cell_class (str): Cell class for bootstrap analysis. channel_combo (str): Channel combination for bootstrap analysis. + compartment_combo (str): Compartment combination for bootstrap analysis. Returns: List[str]: List of all bootstrap output file paths for both constructs and genes. """ # Get all construct data files from checkpoint bootstrap_data_dir = checkpoint.get( - cell_class=cell_class, channel_combo=channel_combo + cell_class=cell_class, + channel_combo=channel_combo, + compartment_combo=compartment_combo, ).output[0] construct_files = glob.glob(f"{bootstrap_data_dir}/*__construct_data.tsv") @@ -362,12 +366,14 @@ def get_bootstrap_inputs( str(construct_nulls_pattern).format( cell_class=cell_class, channel_combo=channel_combo, + compartment_combo=compartment_combo, gene=gene, construct=construct, ), str(construct_pvals_pattern).format( cell_class=cell_class, channel_combo=channel_combo, + compartment_combo=compartment_combo, gene=gene, construct=construct, ), @@ -380,10 +386,16 @@ def get_bootstrap_inputs( outputs.extend( [ str(gene_nulls_pattern).format( - cell_class=cell_class, channel_combo=channel_combo, gene=gene + cell_class=cell_class, + channel_combo=channel_combo, + compartment_combo=compartment_combo, + gene=gene, ), str(gene_pvals_pattern).format( - cell_class=cell_class, channel_combo=channel_combo, gene=gene + cell_class=cell_class, + channel_combo=channel_combo, + compartment_combo=compartment_combo, + gene=gene, ), ] ) @@ -397,6 +409,7 @@ def get_bootstrap_construct_outputs( construct_pvals_pattern: Union[str, Path], cell_class: str, channel_combo: str, + compartment_combo: str, ) -> List[str]: """Get all construct bootstrap outputs for completion flag. @@ -410,6 +423,8 @@ def get_bootstrap_construct_outputs( cell_class (str): Cell class identifier for bootstrap analysis (e.g., 'live', 'dead'). channel_combo (str): Channel combination identifier for bootstrap analysis (e.g., 'dapi_tubulin', 'all_channels'). + compartment_combo (str): Compartment combination identifier for bootstrap analysis + (e.g., 'cell-nucleus-cytoplasm', 'nucleus'). Returns: List[str]: List of all construct bootstrap output file paths, including both @@ -418,7 +433,9 @@ def get_bootstrap_construct_outputs( """ # Get all construct data files from checkpoint bootstrap_data_dir = checkpoint.get( - cell_class=cell_class, channel_combo=channel_combo + cell_class=cell_class, + channel_combo=channel_combo, + compartment_combo=compartment_combo, ).output[0] construct_files = glob.glob(f"{bootstrap_data_dir}/*__construct_data.tsv") @@ -449,12 +466,14 @@ def get_bootstrap_construct_outputs( str(construct_nulls_pattern).format( cell_class=cell_class, channel_combo=channel_combo, + compartment_combo=compartment_combo, gene=gene, construct=construct, ), str(construct_pvals_pattern).format( cell_class=cell_class, channel_combo=channel_combo, + compartment_combo=compartment_combo, gene=gene, construct=construct, ), diff --git a/workflow/lib/shared/segment_cellpose.py b/workflow/lib/shared/segment_cellpose.py index bf72de9c..85cc0d36 100644 --- a/workflow/lib/shared/segment_cellpose.py +++ b/workflow/lib/shared/segment_cellpose.py @@ -54,8 +54,8 @@ CELLPOSE_4X = CELLPOSE_VERSION >= (4, 0) -def initialize_cellpose_model(model_type: str, gpu: bool = False) -> CellposeModel: - """Initialize a CellposeModel with version-aware configuration. +def create_cellpose_model(model_type: str, gpu: bool = False) -> CellposeModel: + r"""Create a CellposeModel with version-aware initialization. Handles differences between Cellpose 3.x and 4.x APIs and validates model compatibility with the installed Cellpose version. @@ -65,7 +65,7 @@ def initialize_cellpose_model(model_type: str, gpu: bool = False) -> CellposeMod or a path to a custom trained model (e.g., 'models/my_custom_model'). - Cellpose 3.x: Supports 'cyto3', 'nuclei', 'cyto2' - Cellpose 4.x: Only supports 'cpsam' - - Custom model paths (containing path separators) are supported in both versions + - Custom model paths (containing '/' or '\\') are supported in both versions gpu (bool, optional): Whether to use GPU for inference. Default is False. Returns: @@ -434,8 +434,8 @@ def segment_cellpose_rgb( # Create Cellpose models using version-aware helper # Nuclei model: "cpsam" for 4.x, "nuclei" for 3.x nuclei_model_type = "cpsam" if CELLPOSE_4X else "nuclei" - model_dapi = initialize_cellpose_model(nuclei_model_type, gpu=gpu) - model_cyto = initialize_cellpose_model(cellpose_model, gpu=gpu) + model_dapi = create_cellpose_model(nuclei_model_type, gpu=gpu) + model_cyto = create_cellpose_model(cellpose_model, gpu=gpu) # Set default kwargs if not provided if nuclei_kwargs is None: @@ -533,7 +533,7 @@ def segment_cellpose_nuclei_rgb( numpy.ndarray: Labeled segmentation mask of nuclei. """ # Create Cellpose model using version-aware helper - model = initialize_cellpose_model(cellpose_model, gpu=gpu) + model = create_cellpose_model(cellpose_model, gpu=gpu) # Segment nuclei using CellposeModel from the RGB image # Pass only blue channel (DAPI) for nuclei segmentation diff --git a/workflow/lib/shared/target_utils.py b/workflow/lib/shared/target_utils.py index 0ed6a537..d8d248c0 100644 --- a/workflow/lib/shared/target_utils.py +++ b/workflow/lib/shared/target_utils.py @@ -221,7 +221,7 @@ def get_merge_targets_by_approach(config): def map_wildcard_outputs(wildcard_combos_df, output_template, wildcards_to_map): """Map specified wildcards in a template string using values from a DataFrame. - Given a template path and a list of wildcards to map (e.g. ["cell_class", "channel_combo"]), + Given a template path and a list of wildcards to map (e.g. ["cell_class", "channel_combo", "compartment_combo"]), replaces only those placeholders with values from the DataFrame, leaving others untouched. Useful for creating lists of output paths where we need to fill in some wildcards but not others. diff --git a/workflow/rules/aggregate.smk b/workflow/rules/aggregate.smk index aa3d5ccf..5c1370d2 100644 --- a/workflow/rules/aggregate.smk +++ b/workflow/rules/aggregate.smk @@ -2,17 +2,37 @@ from lib.shared.target_utils import output_to_input, map_wildcard_outputs from lib.shared.rule_utils import get_montage_inputs, get_bootstrap_inputs, get_bootstrap_construct_outputs +# Aggregate secondary object features into cell-level data +SECOND_OBJ_DETECTION = config["phenotype"].get("second_obj_detection", True) + +if SECOND_OBJ_DETECTION: + rule aggregate_cells_second_objs: + input: + ancient(MERGE_OUTPUTS["final_merge"]), + ancient(PHENOTYPE_OUTPUTS["merge_phenotype_second_objs"][0]), + output: + AGGREGATE_OUTPUTS_MAPPED["aggregate_cells_second_objs"], + params: + agg_strategy=config.get("aggregate", {}).get( + "second_obj_agg_strategy", "none" + ), + script: + "../scripts/aggregate/aggregate_cells_second_objs.py" + + # Create datasets with cell classes and channel combos rule split_datasets: input: - # final merge data - ancient(MERGE_OUTPUTS["final_merge"]), + # cell data (with or without secondary object features) + ancient(AGGREGATE_OUTPUTS["aggregate_cells_second_objs"]) + if SECOND_OBJ_DETECTION + else ancient(MERGE_OUTPUTS["final_merge"]), priority: 100 output: map_wildcard_outputs( aggregate_wildcard_combos, AGGREGATE_OUTPUTS["split_datasets"][0], - ["cell_class", "channel_combo"], + ["cell_class", "channel_combo", "compartment_combo"], ), params: all_channels=config["phenotype"]["channel_names"], @@ -21,8 +41,7 @@ rule split_datasets: confidence_thresholds=config.get("classify", {}).get("confidence_thresholds"), class_title=config.get("classify", {}).get("class_title"), class_mapping=config.get("classify", {}).get("class_mapping"), - cell_classes=aggregate_wildcard_combos["cell_class"].unique(), - channel_combos=aggregate_wildcard_combos["channel_combo"].unique(), + aggregate_wildcard_combos=aggregate_wildcard_combos, script: "../scripts/aggregate/split_datasets.py" @@ -59,6 +78,7 @@ rule generate_feature_table: wildcards={ "cell_class": wildcards.cell_class, "channel_combo": wildcards.channel_combo, + "compartment_combo": wildcards.compartment_combo, }, expansion_values=["plate", "well"], metadata_combos=aggregate_wildcard_combos, @@ -72,10 +92,12 @@ rule generate_feature_table: perturbation_name_col=config["aggregate"]["perturbation_name_col"], perturbation_id_col=config["aggregate"]["perturbation_id_col"], control_key=config["aggregate"]["control_key"], + control_name_col=config["aggregate"].get("control_name_col"), batch_cols=config["aggregate"]["batch_cols"], num_align_batches=config["aggregate"]["num_align_batches"], feature_normalization=config["aggregate"].get("feature_normalization", "standard"), pseudogene_patterns=config.get("aggregate", {}).get("pseudogene_patterns", None), + drop_cols_threshold=config["aggregate"].get("drop_cols_threshold"), script: "../scripts/aggregate/generate_feature_table.py" @@ -87,6 +109,7 @@ rule align: wildcards={ "cell_class": wildcards.cell_class, "channel_combo": wildcards.channel_combo, + "compartment_combo": wildcards.compartment_combo, }, expansion_values=["plate", "well"], metadata_combos=aggregate_wildcard_combos, @@ -105,6 +128,7 @@ rule align: num_align_batches=config["aggregate"]["num_align_batches"], skip_perturbation_score=config["aggregate"]["skip_perturbation_score"], control_name_col=config["aggregate"].get("control_name_col"), + drop_cols_threshold=config["aggregate"].get("drop_cols_threshold"), script: "../scripts/aggregate/align.py" @@ -122,6 +146,7 @@ rule aggregate: agg_method=config["aggregate"]["agg_method"], ps_probability_threshold=config["aggregate"]["ps_probability_threshold"], ps_percentile_threshold=config["aggregate"]["ps_percentile_threshold"], + control_name_col=config["aggregate"].get("control_name_col"), script: "../scripts/aggregate/aggregate.py" @@ -136,6 +161,7 @@ rule eval_aggregate: wildcards={ "cell_class": wildcards.cell_class, "channel_combo": wildcards.channel_combo, + "compartment_combo": wildcards.compartment_combo, }, expansion_values=["plate", "well"], metadata_combos=aggregate_wildcard_combos, @@ -161,9 +187,8 @@ checkpoint prepare_montage_data: AGGREGATE_OUTPUTS["filter"], wildcards={ "cell_class": wildcards.cell_class, - "channel_combo": aggregate_wildcard_combos["channel_combo"].unique()[ - 0 - ], + "channel_combo": aggregate_wildcard_combos["channel_combo"].iloc[0], + "compartment_combo": aggregate_wildcard_combos["compartment_combo"].iloc[0], }, expansion_values=["plate", "well"], metadata_combos=aggregate_wildcard_combos, @@ -231,13 +256,19 @@ rule initiate_montage: checkpoint prepare_bootstrap_data: input: features_singlecell=lambda wildcards: str(AGGREGATE_OUTPUTS["generate_feature_table"][0]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), construct_table=lambda wildcards: str(AGGREGATE_OUTPUTS["generate_feature_table"][1]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), gene_table=lambda wildcards: str(AGGREGATE_OUTPUTS["generate_feature_table"][2]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), output: directory(BOOTSTRAP_OUTPUTS["bootstrap_data_dir"]), @@ -261,13 +292,19 @@ rule bootstrap_construct: input: construct_data=BOOTSTRAP_OUTPUTS["construct_data"], controls_arr=lambda wildcards: str(BOOTSTRAP_OUTPUTS["controls_arr"]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), construct_features_arr=lambda wildcards: str(BOOTSTRAP_OUTPUTS["construct_features_arr"]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), sample_sizes=lambda wildcards: str(BOOTSTRAP_OUTPUTS["sample_sizes"]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), output: BOOTSTRAP_OUTPUTS["bootstrap_construct_nulls"], @@ -287,17 +324,20 @@ rule construct_bootstrap_complete: BOOTSTRAP_OUTPUTS["bootstrap_construct_pvals"], wildcards.cell_class, wildcards.channel_combo, + wildcards.compartment_combo, ), output: - touch(AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__construct_bootstrap_complete.flag"), + touch(AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__construct_bootstrap_complete.flag"), # Aggregate construct results to gene level rule bootstrap_gene: input: - construct_flag=AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__construct_bootstrap_complete.flag", + construct_flag=AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__construct_bootstrap_complete.flag", gene_table=lambda wildcards: str(AGGREGATE_OUTPUTS["generate_feature_table"][2]).format( - cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo + cell_class=wildcards.cell_class, + channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, ), output: BOOTSTRAP_OUTPUTS["bootstrap_gene_nulls"], @@ -307,6 +347,7 @@ rule bootstrap_gene: construct_nulls_pattern=lambda wildcards: str(BOOTSTRAP_OUTPUTS["bootstrap_construct_nulls"]).format( cell_class=wildcards.cell_class, channel_combo=wildcards.channel_combo, + compartment_combo=wildcards.compartment_combo, gene=wildcards.gene, construct="{construct}" ), @@ -326,6 +367,7 @@ rule initiate_bootstrap: BOOTSTRAP_OUTPUTS["bootstrap_gene_pvals"], wildcards.cell_class, wildcards.channel_combo, + wildcards.compartment_combo, ), output: touch(BOOTSTRAP_OUTPUTS["bootstrap_flag"]), @@ -339,8 +381,14 @@ rule combine_bootstrap: BOOTSTRAP_OUTPUTS["combined_construct_results"], BOOTSTRAP_OUTPUTS["combined_gene_results"], params: - constructs_dir=lambda wildcards: str(AGGREGATE_FP / "bootstrap" / f"{wildcards.cell_class}__{wildcards.channel_combo}__constructs"), - genes_dir=lambda wildcards: str(AGGREGATE_FP / "bootstrap" / f"{wildcards.cell_class}__{wildcards.channel_combo}__genes"), + constructs_dir=lambda wildcards: str( + AGGREGATE_FP / "bootstrap" + / f"{wildcards.cell_class}__{wildcards.channel_combo}__{wildcards.compartment_combo}__constructs" + ), + genes_dir=lambda wildcards: str( + AGGREGATE_FP / "bootstrap" + / f"{wildcards.cell_class}__{wildcards.channel_combo}__{wildcards.compartment_combo}__genes" + ), script: "../scripts/aggregate/combine_bootstrap.py" diff --git a/workflow/rules/merge.smk b/workflow/rules/merge.smk index 317261f6..0aae40fd 100644 --- a/workflow/rules/merge.smk +++ b/workflow/rules/merge.smk @@ -89,7 +89,7 @@ if merge_approach == "stitch": phenotype_metadata=ancient(PREPROCESS_OUTPUTS["combine_metadata_phenotype"]), phenotype_stitch_config=MERGE_OUTPUTS["estimate_stitch_phenotype"][0], phenotype_tiles=lambda wildcards: output_to_input( - PHENOTYPE_OUTPUTS["align_phenotype"], + PHENOTYPE_OUTPUTS["align_phenotype"][0], wildcards=wildcards, expansion_values=["tile"], metadata_combos=phenotype_wildcard_combos, @@ -231,7 +231,7 @@ rule format_merge: else MERGE_OUTPUTS["fast_merge"][0] ), ancient(SBS_OUTPUTS["combine_cells"]), - ancient(PHENOTYPE_OUTPUTS["merge_phenotype"][1]), + ancient(PHENOTYPE_OUTPUTS["merge_phenotype_cp"][1]), output: MERGE_OUTPUTS_MAPPED["format_merge"][0], params: @@ -246,7 +246,7 @@ rule deduplicate_merge: input: MERGE_OUTPUTS["format_merge"][0], ancient(SBS_OUTPUTS["combine_cells"]), - ancient(PHENOTYPE_OUTPUTS["merge_phenotype"][1]), + ancient(PHENOTYPE_OUTPUTS["merge_phenotype_cp"][1]), output: deduplication_stats=MERGE_OUTPUTS_MAPPED["deduplicate_merge"][0], deduplicated_data=MERGE_OUTPUTS_MAPPED["deduplicate_merge"][1], @@ -263,7 +263,7 @@ rule deduplicate_merge: rule final_merge: input: MERGE_OUTPUTS["deduplicate_merge"][1], - ancient(PHENOTYPE_OUTPUTS["merge_phenotype"][0]), + ancient(PHENOTYPE_OUTPUTS["merge_phenotype_cp"][0]), output: MERGE_OUTPUTS_MAPPED["final_merge"][0], params: @@ -288,7 +288,7 @@ rule eval_merge: ancient_output=True, ), min_phenotype_cp_paths=lambda wildcards: output_to_input( - PHENOTYPE_OUTPUTS["merge_phenotype"][1], + PHENOTYPE_OUTPUTS["merge_phenotype_cp"][1], wildcards=wildcards, expansion_values=["well"], metadata_combos=phenotype_wildcard_combos, diff --git a/workflow/rules/phenotype.smk b/workflow/rules/phenotype.smk index 6a97acb8..5ff16ce2 100644 --- a/workflow/rules/phenotype.smk +++ b/workflow/rules/phenotype.smk @@ -18,7 +18,8 @@ rule align_phenotype: input: PHENOTYPE_OUTPUTS["apply_ic_field_phenotype"], output: - PHENOTYPE_OUTPUTS_MAPPED["align_phenotype"], + PHENOTYPE_OUTPUTS_MAPPED["align_phenotype"][0], # aligned image + PHENOTYPE_OUTPUTS_MAPPED["align_phenotype"][1], # alignment metrics TSV params: config=lambda wildcards: get_alignment_params(wildcards, config), script: @@ -28,7 +29,7 @@ rule align_phenotype: # Segments cells and nuclei using pre-defined methods rule segment_phenotype: input: - PHENOTYPE_OUTPUTS["align_phenotype"], + PHENOTYPE_OUTPUTS["align_phenotype"][0], output: PHENOTYPE_OUTPUTS_MAPPED["segment_phenotype"], params: @@ -57,6 +58,8 @@ rule extract_phenotype_info: input: # nuclei segmentation map PHENOTYPE_OUTPUTS["segment_phenotype"][0], + # alignment metrics TSV + PHENOTYPE_OUTPUTS["align_phenotype"][1], output: PHENOTYPE_OUTPUTS_MAPPED["extract_phenotype_info"], script: @@ -78,8 +81,70 @@ rule combine_phenotype_info: "../scripts/shared/combine_dfs.py" -# Extract full phenotype information from phenotype images -rule extract_phenotype: +# Identify secondary objects from aligned phenotype image and cell segmentation +if config["phenotype"].get("second_obj_detection", True): + rule identify_second_objs: + input: + # aligned phenotype image + PHENOTYPE_OUTPUTS["align_phenotype"][0], + # cell segmentation map + PHENOTYPE_OUTPUTS["segment_phenotype"][1], + # cytoplasm mask + PHENOTYPE_OUTPUTS["identify_cytoplasm"], + # phenotype info with nuclei centroids + PHENOTYPE_OUTPUTS["extract_phenotype_info"], + output: + # secondary object mask + PHENOTYPE_OUTPUTS_MAPPED["identify_second_objs"][0], + # cell secondary object table + PHENOTYPE_OUTPUTS_MAPPED["identify_second_objs"][1], + # updated cytoplasm masks + PHENOTYPE_OUTPUTS_MAPPED["identify_second_objs"][2], + params: + # Pass all secondary object parameters from config + second_obj_params=config["phenotype"], + script: + "../scripts/phenotype/identify_second_objs.py" + +# Extract secondary object phenotype features +if config["phenotype"].get("second_obj_detection", True): + rule extract_phenotype_second_objs: + input: + # aligned phenotype image + PHENOTYPE_OUTPUTS["align_phenotype"][0], + # secondary object mask + PHENOTYPE_OUTPUTS["identify_second_objs"][0], + # cell secondary object table + PHENOTYPE_OUTPUTS["identify_second_objs"][1], + output: + PHENOTYPE_OUTPUTS_MAPPED["extract_phenotype_second_objs"], + params: + foci_channel_index=config["phenotype"]["foci_channel_index"], + channel_names=config["phenotype"]["channel_names"], + script: + "../scripts/phenotype/extract_phenotype_second_objs.py" + + +# Combine secondary object phenotype results from different tiles +if config["phenotype"].get("second_obj_detection", True): + rule merge_phenotype_second_objs: + input: + lambda wildcards: output_to_input( + PHENOTYPE_OUTPUTS["extract_phenotype_second_objs"], + wildcards=wildcards, + expansion_values=["tile"], + metadata_combos=phenotype_wildcard_combos, + ), + params: + channel_names=config["phenotype"]["channel_names"], + output: + PHENOTYPE_OUTPUTS_MAPPED["merge_phenotype_second_objs"], + script: + "../scripts/phenotype/merge_phenotype_second_objs.py" + + +# Extract full phenotype information using CellProfiler from phenotype images +rule extract_phenotype_cp: input: # aligned phenotype image PHENOTYPE_OUTPUTS["align_phenotype"][0], @@ -89,8 +154,10 @@ rule extract_phenotype: PHENOTYPE_OUTPUTS["segment_phenotype"][1], # cytoplasm segmentation map PHENOTYPE_OUTPUTS["identify_cytoplasm"][0], + # alignment metrics TSV (offset_y, offset_x) + PHENOTYPE_OUTPUTS["align_phenotype"][1], output: - PHENOTYPE_OUTPUTS_MAPPED["extract_phenotype"], + PHENOTYPE_OUTPUTS_MAPPED["extract_phenotype_cp"], params: foci_channel_index=config["phenotype"]["foci_channel_index"], channel_names=config["phenotype"]["channel_names"], @@ -100,11 +167,27 @@ rule extract_phenotype: "../scripts/phenotype/extract_phenotype.py" +# Merge secondary object data with main phenotype data +if config["phenotype"].get("second_obj_detection", True): + rule merge_second_objs_phenotype_cp: + input: + # main phenotype data + PHENOTYPE_OUTPUTS["extract_phenotype_cp"], + # secondary object data + PHENOTYPE_OUTPUTS["identify_second_objs"][1], + output: + PHENOTYPE_OUTPUTS_MAPPED["merge_second_objs_phenotype_cp"], + script: + "../scripts/phenotype/merge_second_objs_phenotype_cp.py" + + # Combine phenotype results from different tiles rule merge_phenotype: input: lambda wildcards: output_to_input( - PHENOTYPE_OUTPUTS["extract_phenotype"], + PHENOTYPE_OUTPUTS["merge_second_objs_phenotype_cp"] + if config["phenotype"].get("second_obj_detection", True) + else PHENOTYPE_OUTPUTS["extract_phenotype_cp"], wildcards=wildcards, expansion_values=["tile"], metadata_combos=phenotype_wildcard_combos, @@ -113,7 +196,7 @@ rule merge_phenotype: channel_names=config["phenotype"]["channel_names"], segment_cells=config["phenotype"].get("segment_cells", True), output: - PHENOTYPE_OUTPUTS_MAPPED["merge_phenotype"], + PHENOTYPE_OUTPUTS_MAPPED["merge_phenotype_cp"], script: "../scripts/phenotype/merge_phenotype.py" @@ -152,7 +235,7 @@ rule eval_features: input: # use minimum phenotype features for evaluation cells_paths=lambda wildcards: output_to_input( - PHENOTYPE_OUTPUTS["merge_phenotype"][1], + PHENOTYPE_OUTPUTS["merge_phenotype_cp"][1], wildcards=wildcards, expansion_values=["well"], metadata_combos=phenotype_wildcard_combos, diff --git a/workflow/scripts/aggregate/aggregate.py b/workflow/scripts/aggregate/aggregate.py index 42cd5e83..03bef877 100644 --- a/workflow/scripts/aggregate/aggregate.py +++ b/workflow/scripts/aggregate/aggregate.py @@ -26,14 +26,29 @@ tvn_normalized = tvn_normalized.to_numpy() del cell_data +# When aggregating by a construct-level ID (e.g. cell_barcode_0) while controls +# are identified via a different column (e.g. gene_symbol_0), carry the +# control/gene column through so downstream clustering / benchmarking / +# annotation merges can match on the human-readable name. +pert_col = snakemake.params.perturbation_name_col +control_name_col = snakemake.params.get("control_name_col") +carry_cols = None +if ( + control_name_col + and control_name_col != pert_col + and control_name_col in metadata.columns +): + carry_cols = [control_name_col] + # Aggregate aggregated_embeddings, aggregated_metadata = aggregate( tvn_normalized, metadata, - snakemake.params.perturbation_name_col, + pert_col, method=snakemake.params.agg_method, ps_probability_threshold=snakemake.params.ps_probability_threshold, ps_percentile_threshold=snakemake.params.ps_percentile_threshold, + carry_cols=carry_cols, ) # Save aggregated data diff --git a/workflow/scripts/aggregate/aggregate_cells_second_objs.py b/workflow/scripts/aggregate/aggregate_cells_second_objs.py new file mode 100644 index 00000000..1c3bb4c2 --- /dev/null +++ b/workflow/scripts/aggregate/aggregate_cells_second_objs.py @@ -0,0 +1,36 @@ +import pandas as pd + +from lib.shared.file_utils import validate_dtypes +from lib.aggregate.second_obj_utils import aggregate_second_obj_data + +# Load cell-level data from final_merge +cells_df = validate_dtypes(pd.read_parquet(snakemake.input[0])) + +# Load per-object secondary object data +second_objs_df = validate_dtypes(pd.read_parquet(snakemake.input[1])) + +# Get aggregation strategy from config +agg_strategy = snakemake.params.agg_strategy + +print(f"Aggregating secondary objects with strategy: {agg_strategy}") +print(f" Cells: {len(cells_df)} rows") +print(f" Secondary objects: {len(second_objs_df)} rows") + +# Filter secondary objects to matching plate/well +plate = int( + snakemake.wildcards.plate +) # plate is int64 in both merge_final and phenotype parquets +well = str(snakemake.wildcards.well) # well is always string ("A1", etc.) +second_objs_filtered = second_objs_df[ + (second_objs_df["plate"] == plate) & (second_objs_df["well"] == well) +] + +print(f" Secondary objects after plate/well filter: {len(second_objs_filtered)} rows") + +# Aggregate +result = aggregate_second_obj_data(cells_df, second_objs_filtered, agg_strategy) + +print(f" Result: {len(result)} rows, {len(result.columns)} columns") + +# Save +result.to_parquet(snakemake.output[0]) diff --git a/workflow/scripts/aggregate/align.py b/workflow/scripts/aggregate/align.py index 2632af6c..a2c2de66 100644 --- a/workflow/scripts/aggregate/align.py +++ b/workflow/scripts/aggregate/align.py @@ -16,6 +16,7 @@ centerscale_by_batch, tvn_on_controls, ) +from lib.aggregate.filter import harmonize_pool_schema from lib.aggregate.perturbation_score import perturbation_score warnings.filterwarnings( @@ -51,13 +52,38 @@ random_indices = np.random.choice(total_rows, size=n_sample, replace=False) random_indices.sort() -# load sample df -sample_df = cell_dataset.scanner().take(random_indices) -sample_df = sample_df.to_pandas(use_threads=True, memory_pool=None) - # load sample df as pandas dataframe use_classifier = snakemake.params.get("use_classifier", False) metadata_cols = load_metadata_cols(snakemake.params.metadata_cols_fp, use_classifier) + +# Harmonize the pool's column set once: drop cols present in only some per-well +# files (typically driven by per-well filter decisions diverging when class +# composition differs across wells) and apply drop_cols_threshold at the pool +# level — matching the convention used by missing_values_filter. +kept_metadata_cols, kept_feature_cols, _pool_report = harmonize_pool_schema( + non_empty_paths, + metadata_cols, + drop_cols_threshold=snakemake.params.get("drop_cols_threshold"), +) +scan_cols = kept_metadata_cols + kept_feature_cols + +# load sample df +sample_df = ( + cell_dataset.scanner(columns=scan_cols) + .take(random_indices) + .to_pandas(use_threads=True, memory_pool=None) +) + +# Residual NaN handling: after pool-level col drops, any remaining NaN is +# sparse and row-localized. Drop those rows from the PCA training sample so +# PCA/StandardScaler see a clean matrix. Matches the drop_rows_threshold intent +# at a per-row granularity appropriate for a training sample. +_sample_pre = len(sample_df) +sample_df = sample_df.dropna(subset=kept_feature_cols).reset_index(drop=True) +if len(sample_df) < _sample_pre: + print( + f"[pool] dropped {_sample_pre - len(sample_df)} PCA-sample rows with residual NaN" + ) metadata, features = split_cell_data(sample_df, metadata_cols) metadata, features = prepare_alignment_data( metadata, @@ -87,11 +113,19 @@ print(f"Processing subset {i + 1}/{num_align_batches} with {len(indices)} cells") subset_df = ( - cell_dataset.scanner() + cell_dataset.scanner(columns=scan_cols) .take(pa.array(indices)) .to_pandas(use_threads=True, memory_pool=None) - .dropna(axis=1) ) + # Drop any residual per-row NaN (pool-level col drops above have already + # eliminated systematically-missing columns). + _pre = len(subset_df) + subset_df = subset_df.dropna(subset=kept_feature_cols).reset_index(drop=True) + if len(subset_df) < _pre: + print( + f"[pool] dropped {_pre - len(subset_df)} rows with residual NaN " + f"from batch {i + 1}" + ) # CALCULATE PERTURBATION SCORE subset_df["perturbation_score"] = np.nan @@ -103,6 +137,9 @@ metadata_cols, snakemake.params.perturbation_name_col, snakemake.params.control_key, + perturbation_id_col=snakemake.params.perturbation_id_col, + control_name_col=snakemake.params.get("control_name_col"), + batch_cols=snakemake.params.batch_cols, ) for col in subset_df.columns: diff --git a/workflow/scripts/aggregate/bootstrap_gene.py b/workflow/scripts/aggregate/bootstrap_gene.py index 03dcf7e2..7a4c6ef1 100644 --- a/workflow/scripts/aggregate/bootstrap_gene.py +++ b/workflow/scripts/aggregate/bootstrap_gene.py @@ -12,6 +12,7 @@ gene_id = snakemake.wildcards.gene cell_class = snakemake.wildcards.cell_class channel_combo = snakemake.wildcards.channel_combo +compartment_combo = snakemake.wildcards.compartment_combo num_sims = snakemake.params.num_sims bootstrap_features_fp = snakemake.params.get("bootstrap_features_fp", None) diff --git a/workflow/scripts/aggregate/eval_aggregate.py b/workflow/scripts/aggregate/eval_aggregate.py index 87e809b6..a999d7c7 100644 --- a/workflow/scripts/aggregate/eval_aggregate.py +++ b/workflow/scripts/aggregate/eval_aggregate.py @@ -64,16 +64,29 @@ # determine original and aligned columns random.seed(42) -# Prefer cell-level features, fall back to nuclear if none found -merge_feature_cols = [ - col for col in merge_data.columns if ("cell_" in col and col.endswith("_mean")) +# Try compartment prefixes in priority order and use the first that yields +# *__mean feature columns. Driven by the wildcard's compartment_combo +# first (whichever compartments this output is actually about), then a fallback +# chain so e.g. second_obj or cytoplasm combos still find intensity means. +compartment_combo = snakemake.wildcards.compartment_combo.split("-") +prefix_priority = compartment_combo + [ + c + for c in ("cell", "nucleus", "cytoplasm", "second_obj") + if c not in compartment_combo ] -if len(merge_feature_cols) == 0: +merge_feature_cols = [] +for prefix in prefix_priority: merge_feature_cols = [ col for col in merge_data.columns - if ("nucleus_" in col and col.endswith("_mean")) + if col.startswith(f"{prefix}_") and col.endswith("_mean") ] + if merge_feature_cols: + print( + f"[eval] using {len(merge_feature_cols)} '{prefix}_*_mean' feature columns" + ) + break + pc_cols = [col for col in aligned_data.columns if col.startswith("PC_")] aligned_feature_cols = random.sample( pc_cols, k=min(len(merge_feature_cols), len(pc_cols)) diff --git a/workflow/scripts/aggregate/generate_feature_table.py b/workflow/scripts/aggregate/generate_feature_table.py index 489fe756..b2a6afe2 100644 --- a/workflow/scripts/aggregate/generate_feature_table.py +++ b/workflow/scripts/aggregate/generate_feature_table.py @@ -11,11 +11,17 @@ from lib.aggregate.align import prepare_alignment_data, centerscale_on_controls from lib.aggregate.cell_data_utils import load_metadata_cols, split_cell_data from lib.aggregate.bootstrap import create_pseudogene_groups +from lib.aggregate.filter import harmonize_pool_schema # get snakemake parameters pert_col = snakemake.params.perturbation_name_col pert_id_col = snakemake.params.perturbation_id_col or pert_col control_key = snakemake.params.control_key +# Column used to identify controls via control_key. When aggregating by a +# construct ID (pert_col = cell_barcode_0) controls are still named in a +# separate column (e.g. gene_symbol_0); fall back to pert_col when unset so +# gene-level runs (control_name_col == pert_col) are unchanged. +control_name_col = snakemake.params.get("control_name_col") or pert_col num_batches = snakemake.params.get("num_align_batches", 1) # Load cell data using PyArrow dataset (lazy - no data loaded yet) @@ -34,14 +40,18 @@ cell_dataset = ds.dataset(non_empty_paths, format="parquet") -# Determine columns -cell_data_cols = cell_dataset.schema.names use_classifier = snakemake.params.get("use_classifier", False) metadata_cols = load_metadata_cols(snakemake.params.metadata_cols_fp, use_classifier) -feature_cols = [col for col in cell_dataset.schema.names if col not in metadata_cols] -# Filter metadata_cols to only include columns that exist in the parquet -existing_metadata_cols = [col for col in metadata_cols if col in cell_data_cols] +# Harmonize the pool's column set once: drop cols present in only some per-well +# files (typically driven by per-well filter decisions diverging when class +# composition differs across wells) and apply drop_cols_threshold at the pool +# level — matching the convention used by missing_values_filter. +existing_metadata_cols, feature_cols, _pool_report = harmonize_pool_schema( + non_empty_paths, + metadata_cols, + drop_cols_threshold=snakemake.params.get("drop_cols_threshold"), +) print( f"Number of metadata columns: {len(existing_metadata_cols)} | Number of feature columns: {len(feature_cols)}" @@ -69,6 +79,7 @@ # We'll collect per-construct data across batches construct_cell_counts = {} # {construct_id: count} construct_gene_map = {} # {construct_id: gene_name} +construct_control_map = {} # {construct_id: control_name_col value (for control id)} construct_feature_sums = {} # {construct_id: [sum of features]} construct_feature_counts = {} # {construct_id: count for averaging} # For median, we need all values - store them @@ -118,6 +129,7 @@ control_key, "batch_values", method=snakemake.params.feature_normalization, + control_col=control_name_col, ).astype(np.float32) # OUTPUT 1: Write center-scaled single-cell data incrementally @@ -142,10 +154,12 @@ mask = metadata[pert_id_col].values == construct_id construct_features = features[mask] gene_name = metadata.loc[mask, pert_col].iloc[0] + control_name = metadata.loc[mask, control_name_col].iloc[0] if construct_id not in construct_cell_counts: construct_cell_counts[construct_id] = 0 construct_gene_map[construct_id] = gene_name + construct_control_map[construct_id] = control_name construct_feature_values[construct_id] = [] construct_cell_counts[construct_id] += mask.sum() @@ -173,6 +187,7 @@ row = { pert_id_col: construct_id, pert_col: construct_gene_map[construct_id], + control_name_col: construct_control_map[construct_id], "cell_count": construct_cell_counts[construct_id], } for i, col in enumerate(feature_cols): @@ -186,7 +201,14 @@ construct_table = pd.DataFrame(construct_rows) # Reorder columns: sgRNA, gene, cell_count, features -construct_columns = [pert_id_col, pert_col, "cell_count"] + feature_cols +# Dedupe while preserving order: when perturbation_id_col == perturbation_name_col +# (a valid config when aggregation is already at the per-construct level), +# listing both would create a duplicate column and break downstream Series ops. +construct_columns = list( + dict.fromkeys( + [pert_id_col, pert_col, control_name_col, "cell_count"] + feature_cols + ) +) construct_table = construct_table[construct_columns] print(f"Construct table shape: {construct_table.shape}") @@ -196,7 +218,7 @@ # Filter out controls for gene-level aggregation non_control_constructs = construct_table[ - ~construct_table[pert_col].str.contains(control_key, na=False) + ~construct_table[control_name_col].str.contains(control_key, na=False) ] # Calculate gene-level sample sizes (sum of construct cell counts) @@ -218,7 +240,7 @@ # Add controls to gene table (controls are their own "genes") control_constructs = construct_table[ - construct_table[pert_col].str.contains(control_key, na=False) + construct_table[control_name_col].str.contains(control_key, na=False) ] control_gene_table = control_constructs[[pert_col, "cell_count"] + feature_cols].copy() diff --git a/workflow/scripts/aggregate/split_datasets.py b/workflow/scripts/aggregate/split_datasets.py index ccaa01f7..1e9c7e76 100644 --- a/workflow/scripts/aggregate/split_datasets.py +++ b/workflow/scripts/aggregate/split_datasets.py @@ -3,9 +3,11 @@ from lib.aggregate.cell_classification import CellClassifier from lib.aggregate.cell_data_utils import ( + COMPARTMENT_PREFIXES, load_metadata_cols, split_cell_data, channel_combo_subset, + compartment_combo_subset, ) @@ -148,8 +150,22 @@ def apply_confidence_thresholds(metadata, features, thresholds, class_col="class # Load all channels all_channels = snakemake.params.all_channels -# split cells by cell class -for cell_class in snakemake.params.cell_classes: +# Each row pairs a channel_combo with its compartment_combo; compartments are zipped not crossed. +aggregate_wildcard_combos = snakemake.params.aggregate_wildcard_combos +unique_specs = aggregate_wildcard_combos[ + ["cell_class", "channel_combo", "compartment_combo"] +].drop_duplicates() + +# Universe of compartment prefixes that may exist in the input data. Must come from +# the known prefix set, not from the TSV — compartments never requested in any combo +# would otherwise pass through undropped. +all_compartments_run = sorted(COMPARTMENT_PREFIXES.keys()) + +for _, row in unique_specs.iterrows(): + cell_class = row["cell_class"] + channel_combo = row["channel_combo"] + compartment_combo = row["compartment_combo"] + if cell_class == "all": cell_class_metadata = metadata cell_class_features = features @@ -158,22 +174,23 @@ def apply_confidence_thresholds(metadata, features, thresholds, class_col="class cell_class_metadata = metadata[cell_class_mask] cell_class_features = features[cell_class_mask] - # split features into channel combos - for channel_combo in snakemake.params.channel_combos: - channel_combo_list = channel_combo.split("_") - channel_combo_features = channel_combo_subset( - cell_class_features, channel_combo_list, all_channels - ) + channel_combo_list = channel_combo.split("_") + compartment_combo_list = compartment_combo.split("-") + + filtered = channel_combo_subset( + cell_class_features, channel_combo_list, all_channels + ) + filtered = compartment_combo_subset( + filtered, compartment_combo_list, all_compartments_run + ) + + cell_class_data = pd.concat([cell_class_metadata, filtered], axis=1).reset_index( + drop=True + ) - # concatenate metadata and features - cell_class_data = pd.concat( - [cell_class_metadata, channel_combo_features], axis=1 - ).reset_index(drop=True) - - # Save data - dataset_fp = [ - f - for f in snakemake.output - if f"CeCl-{cell_class}_ChCo-{channel_combo}__" in f - ][0] - cell_class_data.to_parquet(dataset_fp, index=False) + dataset_fp = [ + f + for f in snakemake.output + if f"CeCl-{cell_class}_ChCo-{channel_combo}_CmCo-{compartment_combo}__" in f + ][0] + cell_class_data.to_parquet(dataset_fp, index=False) diff --git a/workflow/scripts/phenotype/align_phenotype.py b/workflow/scripts/phenotype/align_phenotype.py index b6193c01..5091ee09 100644 --- a/workflow/scripts/phenotype/align_phenotype.py +++ b/workflow/scripts/phenotype/align_phenotype.py @@ -1,3 +1,4 @@ +import pandas as pd from tifffile import imread, imwrite from lib.phenotype.align_channels import align_phenotype_channels @@ -13,6 +14,9 @@ # Start with original image data aligned_data = image_data +# Dictionary to collect all alignment metrics +all_metrics = {} + # STEP 1: Apply custom offsets FIRST (if they exist) if align_config.get("custom_channel_offsets"): print("STEP 1: Applying custom channel offsets...") @@ -36,7 +40,7 @@ for i, step in enumerate(align_config["steps"], 1): print(f" Step {i}: Aligning channels...") print(f" Step parameters: {step}") - aligned_data = align_phenotype_channels( + aligned_data, metrics = align_phenotype_channels( aligned_data, target=step["target"], source=step["source"], @@ -46,11 +50,15 @@ "upsample_factor", align_config.get("upsample_factor", 2) ), window=step.get("window", align_config.get("window", 2)), + return_metrics=True, ) + # Add step-suffixed metrics + all_metrics[f"offset_y_step{i}"] = metrics["offset"][0] + all_metrics[f"offset_x_step{i}"] = metrics["offset"][1] else: # Handle single-step alignment print("Performing single-step alignment...") - aligned_data = align_phenotype_channels( + aligned_data, metrics = align_phenotype_channels( aligned_data, target=align_config["target"], source=align_config["source"], @@ -58,9 +66,30 @@ remove_channel=align_config["remove_channel"], upsample_factor=align_config.get("upsample_factor", 2), window=align_config.get("window", 2), + return_metrics=True, ) + # Add metrics without step suffix for single-step + all_metrics["offset_y"] = metrics["offset"][0] + all_metrics["offset_x"] = metrics["offset"][1] else: print("STEP 2: Skipping automatic alignment") + # No alignment - write placeholder metrics + all_metrics["offset_y"] = 0 + all_metrics["offset_x"] = 0 # Save the aligned/unaligned data as a .tiff file imwrite(snakemake.output[0], aligned_data) + +# Save alignment metrics to TSV (one row per tile) +metrics_df = pd.DataFrame( + [ + { + "plate": snakemake.wildcards.plate, + "well": snakemake.wildcards.well, + "tile": snakemake.wildcards.tile, + **all_metrics, + } + ] +) +metrics_df.to_csv(snakemake.output[1], index=False, sep="\t") +print(f"Alignment metrics saved to {snakemake.output[1]}") diff --git a/workflow/scripts/phenotype/extract_phenotype.py b/workflow/scripts/phenotype/extract_phenotype.py index 7a0c2416..064ff071 100644 --- a/workflow/scripts/phenotype/extract_phenotype.py +++ b/workflow/scripts/phenotype/extract_phenotype.py @@ -1,3 +1,4 @@ +import pandas as pd from tifffile import imread # load inputs @@ -47,5 +48,11 @@ f"Unknown cp_method: {cp_method}. Choose 'cp_measure' or 'cp_emulator'." ) +# Broadcast tile-level alignment offsets to each cell row +alignment_metrics = pd.read_csv(snakemake.input[4], sep="\t") +offset_cols = [c for c in alignment_metrics.columns if c.startswith("offset_")] +for col in offset_cols: + phenotype_cp[col] = alignment_metrics[col].iloc[0] + # save phenotype cp phenotype_cp.to_csv(snakemake.output[0], index=False, sep="\t") diff --git a/workflow/scripts/phenotype/extract_phenotype_second_objs.py b/workflow/scripts/phenotype/extract_phenotype_second_objs.py new file mode 100644 index 00000000..37766ca6 --- /dev/null +++ b/workflow/scripts/phenotype/extract_phenotype_second_objs.py @@ -0,0 +1,48 @@ +from tifffile import imread +import pandas as pd + +from lib.phenotype.extract_phenotype_second_objs import extract_phenotype_second_objs + +# Load inputs +data_phenotype = imread(snakemake.input[0]) +second_obj_masks = imread(snakemake.input[1]) + +# Load only the second_obj_cell_mapping table from the combined dataframe +combined_df = pd.read_csv(snakemake.input[2], sep="\t") +second_obj_cell_mapping_df = combined_df[ + combined_df["table_type"] == "second_obj_cell_mapping" +].copy() + +# Create a dictionary to rename columns by removing the 'second_obj_mapping_' prefix +rename_dict = {} +for col in second_obj_cell_mapping_df.columns: + if col.startswith("second_obj_mapping_"): + rename_dict[col] = col.replace("second_obj_mapping_", "") + +# Rename columns +second_obj_cell_mapping_df = second_obj_cell_mapping_df.rename(columns=rename_dict) + +# Get a list of all columns that start with 'cell_summary_' +cell_summary_cols = [ + col for col in second_obj_cell_mapping_df.columns if col.startswith("cell_summary_") +] + +# Drop the table_type column and all cell_summary columns +columns_to_drop = ["table_type"] + cell_summary_cols +second_obj_cell_mapping_df = second_obj_cell_mapping_df.drop(columns=columns_to_drop) + +# Print the final columns to verify +print("Final columns:", second_obj_cell_mapping_df.columns.tolist()) + +# Extract secondary object phenotype features +second_obj_phenotype = extract_phenotype_second_objs( + data_phenotype=data_phenotype, + second_objs=second_obj_masks, + wildcards=snakemake.wildcards, + second_obj_cell_mapping_df=second_obj_cell_mapping_df, + foci_channel=snakemake.params.foci_channel_index, + channel_names=snakemake.params.channel_names, +) + +# Save results +second_obj_phenotype.to_csv(snakemake.output[0], sep="\t", index=False) diff --git a/workflow/scripts/phenotype/identify_second_objs.py b/workflow/scripts/phenotype/identify_second_objs.py new file mode 100644 index 00000000..fc5f7807 --- /dev/null +++ b/workflow/scripts/phenotype/identify_second_objs.py @@ -0,0 +1,163 @@ +from tifffile import imread, imwrite +import pandas as pd + +# Load input files +data_phenotype = imread(snakemake.input[0]) +cells = imread(snakemake.input[1]) +cytoplasms = imread(snakemake.input[2]) +phenotype_info = pd.read_csv(snakemake.input[3], sep="\t") + +# Prepare nuclei centroids dictionary from phenotype info +nuclei_centroids_dict = None +if "i" in phenotype_info.columns and "j" in phenotype_info.columns: + nuclei_id_col = ( + "nuclei_id" if "nuclei_id" in phenotype_info.columns else phenotype_info.index + ) + nuclei_centroids_dict = { + row.get("nuclei_id", idx): (row["i"], row["j"]) + for idx, row in phenotype_info.iterrows() + } + +# Get parameters from config +params = snakemake.params.second_obj_params + +# Check which segmentation method to use +use_ml = params.get("use_ml_segmentation", False) + +# Common parameters shared by both methods +common_params = { + "image": data_phenotype, + "second_obj_channel_index": params["second_obj_channel_index"], + "cell_masks": cells, + "cytoplasm_masks": cytoplasms, + "second_obj_min_size": params.get("second_obj_min_size", 10), + "second_obj_max_size": params.get("second_obj_max_size", 200), + "size_filter_method": params.get("size_filter_method", "feret"), + "max_objects_per_cell": params.get("max_objects_per_cell", 120), + "overlap_threshold": params.get("overlap_threshold", 0.1), + "nuclei_centroids": nuclei_centroids_dict, + "max_total_objects": params.get("max_total_objects", 1000), +} + +if use_ml: + from lib.phenotype.segment_secondary_object import segment_second_objs_ml + + # Parameters already handled in common_params (by key name) + common_param_keys = { + "second_obj_channel_index", + "second_obj_min_size", + "second_obj_max_size", + "size_filter_method", + "max_objects_per_cell", + "overlap_threshold", + "max_total_objects", + } + # Parameters specific to threshold/CV method (should not be passed to ML) + cv_only_params = { + "threshold_smoothing_scale", + "threshold_method", + "use_morphological_opening", + "opening_disk_radius", + "fill_holes", + "declump_method", + "declump_mode", + "suppress_local_maxima", + "maxima_reduction_factor", + "use_shape_refinement", + "proportion_threshold", + } + + # General config parameters (not segmentation-specific) + config_level_params = { + "use_ml_segmentation", + "second_obj_detection", + "foci_channel_index", + "channel_names", + "dapi_index", + "cyto_index", + "align", + "segmentation_method", + "reconcile", + "cp_method", + "nuclei_diameter", + "cell_diameter", + "target", + "source", + "riders", + "remove_channel", + "upsample_factor", + "window", + } + + # Collect ML-specific parameters only + ml_params = { + k: v + for k, v in params.items() + if k not in common_param_keys + and k not in cv_only_params + and k not in config_level_params + } + + # Call ML segmentation with common params and ML-specific params + second_obj_masks, cell_second_obj_table, updated_cytoplasm_masks = ( + segment_second_objs_ml(**common_params, **ml_params) + ) +else: + from lib.phenotype.segment_secondary_object import segment_second_objs + + # CV-specific parameters with defaults + cv_params = { + "threshold_smoothing_scale": params.get("threshold_smoothing_scale", 1.3488), + "threshold_method": params.get("threshold_method", "otsu_two_peak"), + "use_morphological_opening": params.get("use_morphological_opening", True), + "opening_disk_radius": params.get("opening_disk_radius", 1), + "fill_holes": params.get("fill_holes", "both"), + "declump_method": params.get("declump_method", "shape"), + "declump_mode": params.get("declump_mode", "watershed"), + "suppress_local_maxima": params.get("suppress_local_maxima", 20), + "maxima_reduction_factor": params.get("maxima_reduction_factor", None), + "use_shape_refinement": params.get("use_shape_refinement", False), + "proportion_threshold": params.get("proportion_threshold", 0.4), + } + + # Call traditional segmentation + second_obj_masks, cell_second_obj_table, updated_cytoplasm_masks = ( + segment_second_objs(**common_params, **cv_params) + ) + +# Save outputs +# Save secondary object masks as TIFF +imwrite(snakemake.output[0], second_obj_masks) + +# Save cell-secondary object table as TSV +# It has two DataFrames, save both +cell_summary_df = cell_second_obj_table["cell_summary"] +second_obj_cell_mapping_df = cell_second_obj_table["second_obj_cell_mapping"] + +# Combine into one DataFrame with a 'table_type' column for filtering +cell_summary_df["table_type"] = "cell_summary" +second_obj_cell_mapping_df["table_type"] = "second_obj_cell_mapping" + +# Ensure no column conflicts by prefixing with table type +cell_summary_cols = { + col: f"cell_summary_{col}" for col in cell_summary_df.columns if col != "table_type" +} +second_obj_mapping_cols = { + col: f"second_obj_mapping_{col}" + for col in second_obj_cell_mapping_df.columns + if col != "table_type" +} + +cell_summary_df = cell_summary_df.rename(columns=cell_summary_cols) +second_obj_cell_mapping_df = second_obj_cell_mapping_df.rename( + columns=second_obj_mapping_cols +) + +# Combine and save +combined_df = pd.concat( + [cell_summary_df, second_obj_cell_mapping_df], ignore_index=True +) +combined_df.to_csv(snakemake.output[1], sep="\t", index=False) + +# Save updated cytoplasm masks as TIFF +imwrite(snakemake.output[2], updated_cytoplasm_masks) diff --git a/workflow/scripts/phenotype/merge_phenotype.py b/workflow/scripts/phenotype/merge_phenotype.py index 9fe6d3ab..9610f118 100644 --- a/workflow/scripts/phenotype/merge_phenotype.py +++ b/workflow/scripts/phenotype/merge_phenotype.py @@ -1,20 +1,14 @@ import pandas as pd from joblib import Parallel, delayed - -# Define function to read df tsv files -def get_file(f): - try: - return pd.read_csv(f, sep="\t") - except pd.errors.EmptyDataError: - pass - +from lib.shared.file_utils import read_tsv_safe # Load, concatenate, and save the phenotype CellProfiler data arr_reads = Parallel(n_jobs=snakemake.threads)( - delayed(get_file)(file) for file in snakemake.input + delayed(read_tsv_safe)(file) for file in snakemake.input ) -phenotype_cp = pd.concat(arr_reads) +valid_dfs = [df for df in arr_reads if not df.empty] +phenotype_cp = pd.concat(valid_dfs) if valid_dfs else pd.DataFrame() phenotype_cp.to_parquet(snakemake.output[0]) diff --git a/workflow/scripts/phenotype/merge_phenotype_second_objs.py b/workflow/scripts/phenotype/merge_phenotype_second_objs.py new file mode 100644 index 00000000..bb41f074 --- /dev/null +++ b/workflow/scripts/phenotype/merge_phenotype_second_objs.py @@ -0,0 +1,23 @@ +import pandas as pd +from joblib import Parallel, delayed + +from lib.shared.file_utils import read_tsv_safe + +# Load, concatenate, and save the secondary object phenotype data +arr_reads = Parallel(n_jobs=snakemake.threads)( + delayed(read_tsv_safe)(file) for file in snakemake.input +) + +# Combine all dataframes, filtering out empty dataframes +valid_dfs = [df for df in arr_reads if not df.empty] +if valid_dfs: + second_obj_phenotype = pd.concat(valid_dfs) + print( + f"Combined {len(valid_dfs)} files with a total of {len(second_obj_phenotype)} secondary object records" + ) +else: + print("Warning: No valid data files found!") + second_obj_phenotype = pd.DataFrame() + +# Save the combined secondary object phenotype data +second_obj_phenotype.to_parquet(snakemake.output[0]) diff --git a/workflow/scripts/phenotype/merge_second_objs_phenotype_cp.py b/workflow/scripts/phenotype/merge_second_objs_phenotype_cp.py new file mode 100644 index 00000000..6f30959f --- /dev/null +++ b/workflow/scripts/phenotype/merge_second_objs_phenotype_cp.py @@ -0,0 +1,63 @@ +import pandas as pd + +# Load the datasets +phenotype_data = pd.read_csv(snakemake.input[0], sep="\t") + +# Load the combined secondary object file and extract cell summary data +combined_second_obj_df = pd.read_csv(snakemake.input[1], sep="\t") +cell_summary_df = combined_second_obj_df[ + combined_second_obj_df["table_type"] == "cell_summary" +].copy() + +# Check if we have phenotype data and cell summary data +if len(phenotype_data) > 0 and len(cell_summary_df) > 0: + # Filter to only keep cell_summary columns (and table_type) + cell_summary_columns = [ + col + for col in cell_summary_df.columns + if col.startswith("cell_summary_") or col == "table_type" + ] + cell_summary_df = cell_summary_df[cell_summary_columns] + + # Remove the 'cell_summary_' prefix from column names + rename_dict = {} + for col in cell_summary_df.columns: + if col.startswith("cell_summary_"): + rename_dict[col] = col.replace("cell_summary_", "") + + cell_summary_df = cell_summary_df.rename(columns=rename_dict) + + # Drop the table_type column + cell_summary_df = cell_summary_df.drop(columns=["table_type"]) + + # Merge on cell_id (secondary objects) = label (phenotype) + merged_data = phenotype_data.merge( + cell_summary_df, left_on="label", right_on="cell_id", how="left" + ) + + # Drop the redundant cell_id column + merged_data = merged_data.drop("cell_id", axis=1) + + print( + f"Merged {len(phenotype_data)} phenotype records with {len(cell_summary_df)} secondary object records" + ) + +elif len(phenotype_data) > 0: + # No cell summary data available, just use phenotype data + merged_data = phenotype_data.copy() + print( + f"No secondary object data available - using phenotype data only ({len(phenotype_data)} records)" + ) + +else: + # Both datasets are empty - create an empty DataFrame + merged_data = pd.DataFrame() + print( + "Both phenotype and secondary object datasets are empty - creating empty output" + ) + +# Save the merged dataset +merged_data.to_csv(snakemake.output[0], sep="\t", index=False) +print( + f"Final dataset has {len(merged_data)} rows and {len(merged_data.columns)} columns" +) diff --git a/workflow/scripts/preprocess/combine_metadata.py b/workflow/scripts/preprocess/combine_metadata.py index 94155ab8..1eb972a6 100644 --- a/workflow/scripts/preprocess/combine_metadata.py +++ b/workflow/scripts/preprocess/combine_metadata.py @@ -2,20 +2,11 @@ import pandas as pd from joblib import Parallel, delayed -from lib.shared.file_utils import validate_dtypes - - -def get_file(f): - """Read a TSV file safely.""" - try: - return pd.read_csv(f, sep="\t") - except pd.errors.EmptyDataError: - return pd.DataFrame() - +from lib.shared.file_utils import validate_dtypes, read_tsv_safe # Load all metadata files all_dfs = Parallel(n_jobs=snakemake.threads)( - delayed(get_file)(file) for file in snakemake.input + delayed(read_tsv_safe)(file) for file in snakemake.input ) # Combine all dataframes diff --git a/workflow/scripts/shared/combine_dfs.py b/workflow/scripts/shared/combine_dfs.py index ce373631..a213564c 100644 --- a/workflow/scripts/shared/combine_dfs.py +++ b/workflow/scripts/shared/combine_dfs.py @@ -1,22 +1,16 @@ import pandas as pd from joblib import Parallel, delayed -from lib.shared.file_utils import validate_dtypes - - -# Define function to read df tsv files -def get_file(f): - try: - return pd.read_csv(f, sep="\t") - except pd.errors.EmptyDataError: - pass - +from lib.shared.file_utils import validate_dtypes, read_tsv_safe # Load and concatenate data all_dfs = Parallel(n_jobs=snakemake.threads)( - delayed(get_file)(file) for file in snakemake.input + delayed(read_tsv_safe)(file) for file in snakemake.input +) +valid_dfs = [df for df in all_dfs if not df.empty] +combined_df = ( + pd.concat(valid_dfs).reset_index(drop=True) if valid_dfs else pd.DataFrame() ) -combined_df = pd.concat(all_dfs).reset_index(drop=True) # Validate col types # Empty dfs can cause issues with dtype diff --git a/workflow/scripts/shared/extract_phenotype_minimal.py b/workflow/scripts/shared/extract_phenotype_minimal.py index b35e12db..422c6a26 100644 --- a/workflow/scripts/shared/extract_phenotype_minimal.py +++ b/workflow/scripts/shared/extract_phenotype_minimal.py @@ -1,3 +1,4 @@ +import pandas as pd from tifffile import imread from lib.shared.extract_phenotype_minimal import extract_phenotype_minimal @@ -12,5 +13,15 @@ wildcards=snakemake.wildcards, ) +# Add alignment metrics columns if provided (e.g., phenotype has them, SBS does not) +if len(snakemake.input) > 1: + alignment_metrics = pd.read_csv(snakemake.input[1], sep="\t") + # Excludes plate/well/tile as those are already in phenotype_minimal + metrics_cols = [ + c for c in alignment_metrics.columns if c not in ["plate", "well", "tile"] + ] + for col in metrics_cols: + phenotype_minimal[col] = alignment_metrics[col].iloc[0] + # save minimal phenotype data phenotype_minimal.to_csv(snakemake.output[0], index=False, sep="\t") diff --git a/workflow/targets/aggregate.smk b/workflow/targets/aggregate.smk index 7765c62c..2ac9af47 100644 --- a/workflow/targets/aggregate.smk +++ b/workflow/targets/aggregate.smk @@ -6,6 +6,15 @@ AGGREGATE_FP = ROOT_FP / "aggregate" # Define standard (non-montage) aggreagte outputs AGGREGATE_OUTPUTS = { + "aggregate_cells_second_objs": [ + AGGREGATE_FP + / "parquets" + / get_filename( + {"plate": "{plate}", "well": "{well}"}, + "aggregated_cells_second_objs", + "parquet", + ), + ], "split_datasets": [ AGGREGATE_FP / "parquets" @@ -15,6 +24,7 @@ AGGREGATE_OUTPUTS = { "well": "{well}", "cell_class": "{cell_class}", "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", }, "merge_data", "parquet", @@ -29,6 +39,7 @@ AGGREGATE_OUTPUTS = { "well": "{well}", "cell_class": "{cell_class}", "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", }, "filtered", "parquet", @@ -38,21 +49,33 @@ AGGREGATE_OUTPUTS = { AGGREGATE_FP / "parquets" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "features_singlecell", "parquet", ), AGGREGATE_FP / "tsvs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "features_constructs", "tsv", ), AGGREGATE_FP / "tsvs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "features_genes", "tsv", ), @@ -61,7 +84,11 @@ AGGREGATE_OUTPUTS = { AGGREGATE_FP / "parquets" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "aligned", "parquet", ), @@ -70,7 +97,11 @@ AGGREGATE_OUTPUTS = { AGGREGATE_FP / "tsvs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "aggregated", "tsv", ), @@ -79,21 +110,33 @@ AGGREGATE_OUTPUTS = { AGGREGATE_FP / "eval" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "na_stats", "tsv", ), AGGREGATE_FP / "eval" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "na_stats", "png", ), AGGREGATE_FP / "eval" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "feature_distributions", "png", ), @@ -101,6 +144,7 @@ AGGREGATE_OUTPUTS = { } AGGREGATE_OUTPUT_MAPPINGS = { + "aggregate_cells_second_objs": None, "split_datasets": None, "filter": None, "perturbation_score_filter": None, @@ -110,7 +154,25 @@ AGGREGATE_OUTPUT_MAPPINGS = { "generate_feature_table": None, } -AGGREGATE_OUTPUTS_MAPPED = map_outputs(AGGREGATE_OUTPUTS, AGGREGATE_OUTPUT_MAPPINGS) +# Determine which outputs to include based on config +AGGREGATE_SECOND_OBJ_DETECTION = config["phenotype"].get("second_obj_detection", True) + +if not AGGREGATE_SECOND_OBJ_DETECTION: + # Filter out secondary object rules when disabled + AGGREGATE_OUTPUTS_FILTERED = { + k: v for k, v in AGGREGATE_OUTPUTS.items() + if k not in ["aggregate_cells_second_objs"] + } + + AGGREGATE_OUTPUT_MAPPINGS_FILTERED = { + k: v for k, v in AGGREGATE_OUTPUT_MAPPINGS.items() + if k not in ["aggregate_cells_second_objs"] + } +else: + AGGREGATE_OUTPUTS_FILTERED = AGGREGATE_OUTPUTS + AGGREGATE_OUTPUT_MAPPINGS_FILTERED = AGGREGATE_OUTPUT_MAPPINGS + +AGGREGATE_OUTPUTS_MAPPED = map_outputs(AGGREGATE_OUTPUTS_FILTERED, AGGREGATE_OUTPUT_MAPPINGS_FILTERED) # TODO: Use all combos # aggregate_wildcard_combos = aggregate_wildcard_combos[ @@ -121,7 +183,7 @@ AGGREGATE_OUTPUTS_MAPPED = map_outputs(AGGREGATE_OUTPUTS, AGGREGATE_OUTPUT_MAPPI # ] AGGREGATE_TARGETS_ALL = outputs_to_targets( - AGGREGATE_OUTPUTS, aggregate_wildcard_combos, AGGREGATE_OUTPUT_MAPPINGS + AGGREGATE_OUTPUTS_FILTERED, aggregate_wildcard_combos, AGGREGATE_OUTPUT_MAPPINGS_FILTERED ) @@ -160,51 +222,76 @@ MONTAGE_OUTPUTS = { "montage_flag": AGGREGATE_FP / "montages" / "{cell_class}__montages_complete.flag", } cell_classes = aggregate_wildcard_combos["cell_class"].unique() -MONTAGE_TARGETS_ALL = [ - str(MONTAGE_OUTPUTS["montage_flag"]).format(cell_class=cell_class) - for cell_class in cell_classes -] +generate_montages = config.get("aggregate", {}).get("generate_montages", True) +MONTAGE_TARGETS_ALL = ( + [ + str(MONTAGE_OUTPUTS["montage_flag"]).format(cell_class=cell_class) + for cell_class in cell_classes + ] + if generate_montages + else [] +) # Define bootstrap outputs # These are special because we dynamically derive outputs BOOTSTRAP_OUTPUTS = { # Data preparation outputs - "bootstrap_data_dir": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__bootstrap_data", - "construct_data": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__bootstrap_data" / "{gene}__{construct}__construct_data.tsv", - + "bootstrap_data_dir": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__bootstrap_data", + "construct_data": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__bootstrap_data" / "{gene}__{construct}__construct_data.tsv", + # Input arrays "controls_arr": AGGREGATE_FP / "bootstrap" / "inputs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "controls_arr", "tsv" ), "construct_features_arr": AGGREGATE_FP / "bootstrap" / "inputs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "construct_features_arr", "tsv" ), "sample_sizes": AGGREGATE_FP / "bootstrap" / "inputs" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "sample_sizes", "tsv" ), # Construct-level outputs - "bootstrap_construct_nulls": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__constructs" / "{gene}__{construct}__nulls.npy", - "bootstrap_construct_pvals": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__constructs" / "{gene}__{construct}__pvals.tsv", - + "bootstrap_construct_nulls": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__constructs" / "{gene}__{construct}__nulls.npy", + "bootstrap_construct_pvals": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__constructs" / "{gene}__{construct}__pvals.tsv", + # Gene-level outputs - "bootstrap_gene_nulls": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__genes" / "{gene}__nulls.npy", - "bootstrap_gene_pvals": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__genes" / "{gene}__pvals.tsv", - + "bootstrap_gene_nulls": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__genes" / "{gene}__nulls.npy", + "bootstrap_gene_pvals": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__genes" / "{gene}__pvals.tsv", + # Completion flags - "bootstrap_flag": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__bootstrap_complete.flag", + "bootstrap_flag": AGGREGATE_FP / "bootstrap" / "{cell_class}__{channel_combo}__{compartment_combo}__bootstrap_complete.flag", # Combined results "combined_construct_results": AGGREGATE_FP / "bootstrap" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "all_construct_bootstrap_results", "tsv" ), "combined_gene_results": AGGREGATE_FP / "bootstrap" / get_filename( - {"cell_class": "{cell_class}", "channel_combo": "{channel_combo}"}, + { + "cell_class": "{cell_class}", + "channel_combo": "{channel_combo}", + "compartment_combo": "{compartment_combo}", + }, "all_gene_bootstrap_results", "tsv" ), } @@ -212,7 +299,11 @@ BOOTSTRAP_OUTPUTS = { # Bootstrap target combinations bootstrap_combos = config.get("aggregate", {}).get("bootstrap_combinations", []) BOOTSTRAP_TARGETS_ALL = [ - str(output_path).format(cell_class=combo["cell_class"], channel_combo=combo["channel_combo"]) + str(output_path).format( + cell_class=combo["cell_class"], + channel_combo=combo["channel_combo"], + compartment_combo=combo["compartment_combo"], + ) for combo in bootstrap_combos for output_path in [ BOOTSTRAP_OUTPUTS["combined_construct_results"], diff --git a/workflow/targets/cluster.smk b/workflow/targets/cluster.smk index 43d91cdc..8d2d6ebd 100644 --- a/workflow/targets/cluster.smk +++ b/workflow/targets/cluster.smk @@ -11,12 +11,14 @@ CLUSTER_OUTPUTS = { "clean_aggregate": [ CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / get_filename({}, "aggregate_cleaned", "tsv"), ], "phate_leiden_clustering": [ CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename( @@ -26,11 +28,13 @@ CLUSTER_OUTPUTS = { ), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({}, "cluster_sizes", "png"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({}, "clusters", "png"), @@ -38,41 +42,49 @@ CLUSTER_OUTPUTS = { "benchmark_clusters": [ CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Real"}, "integrated_results", "json"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Shuffled"}, "integrated_results", "json"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Real"}, "combined_table", "tsv"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Shuffled"}, "combined_table", "tsv"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Real"}, "global_metrics", "json"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Shuffled"}, "global_metrics", "json"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Real"}, "pie_chart", "png"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename( @@ -80,11 +92,13 @@ CLUSTER_OUTPUTS = { ), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename({"cluster_benchmark": "Real"}, "enrichment_bar_chart", "png"), CLUSTER_FP / "{channel_combo}" + / "{compartment_combo}" / "{cell_class}" / "{leiden_resolution}" / get_filename( diff --git a/workflow/targets/phenotype.smk b/workflow/targets/phenotype.smk index c6b43539..ab239070 100644 --- a/workflow/targets/phenotype.smk +++ b/workflow/targets/phenotype.smk @@ -26,6 +26,13 @@ PHENOTYPE_OUTPUTS = { / get_filename( {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, "aligned", "tiff" ), + PHENOTYPE_FP + / "tsvs" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "alignment_metrics", + "tsv", + ), ], "segment_phenotype": [ PHENOTYPE_FP @@ -71,7 +78,30 @@ PHENOTYPE_OUTPUTS = { {"plate": "{plate}", "well": "{well}"}, "phenotype_info", "parquet" ), ], - "extract_phenotype": [ + "identify_second_objs": [ + PHENOTYPE_FP + / "images" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "identified_second_objs", + "tiff", + ), + PHENOTYPE_FP + / "tsvs" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "cell_second_obj_table", + "tsv", + ), + PHENOTYPE_FP + / "images" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "updated_cytoplasms", + "tiff", + ), + ], + "extract_phenotype_cp": [ PHENOTYPE_FP / "tsvs" / get_filename( @@ -80,7 +110,32 @@ PHENOTYPE_OUTPUTS = { "tsv", ), ], - "merge_phenotype": [ + "extract_phenotype_second_objs": [ + PHENOTYPE_FP + / "tsvs" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "phenotype_second_objs", + "tsv", + ), + ], + "merge_phenotype_second_objs": [ + PHENOTYPE_FP + / "parquets" + / get_filename( + {"plate": "{plate}", "well": "{well}"}, "phenotype_second_objs", "parquet" + ), + ], + "merge_second_objs_phenotype_cp": [ + PHENOTYPE_FP + / "tsvs" + / get_filename( + {"plate": "{plate}", "well": "{well}", "tile": "{tile}"}, + "phenotype_with_second_objs", + "tsv", + ), + ], + "merge_phenotype_cp": [ PHENOTYPE_FP / "parquets" / get_filename( @@ -130,14 +185,46 @@ PHENOTYPE_OUTPUT_MAPPINGS = { "identify_cytoplasm": temp, "extract_phenotype_info": temp, "combine_phenotype_info": None, - "extract_phenotype": temp, - "merge_phenotype": None, + "identify_second_objs": None, + "extract_phenotype_cp": None, + "extract_phenotype_second_objs": None, + "merge_phenotype_second_objs": None, + "merge_second_objs_phenotype_cp": None, + "merge_phenotype_cp": None, "eval_segmentation_phenotype": None, "eval_features": None, } -PHENOTYPE_OUTPUTS_MAPPED = map_outputs(PHENOTYPE_OUTPUTS, PHENOTYPE_OUTPUT_MAPPINGS) +# Determine which outputs to include based on config +PHENOTYPE_SECOND_OBJ_DETECTION = config["phenotype"].get("second_obj_detection", True) + +if not PHENOTYPE_SECOND_OBJ_DETECTION: + # Filter out secondary object rules when disabled + PHENOTYPE_OUTPUTS_FILTERED = { + k: v for k, v in PHENOTYPE_OUTPUTS.items() + if k not in [ + "identify_second_objs", + "extract_phenotype_second_objs", + "merge_phenotype_second_objs", + "merge_second_objs_phenotype_cp", + ] + } + + PHENOTYPE_OUTPUT_MAPPINGS_FILTERED = { + k: v for k, v in PHENOTYPE_OUTPUT_MAPPINGS.items() + if k not in [ + "identify_second_objs", + "extract_phenotype_second_objs", + "merge_phenotype_second_objs", + "merge_second_objs_phenotype_cp", + ] + } +else: + PHENOTYPE_OUTPUTS_FILTERED = PHENOTYPE_OUTPUTS + PHENOTYPE_OUTPUT_MAPPINGS_FILTERED = PHENOTYPE_OUTPUT_MAPPINGS + +PHENOTYPE_OUTPUTS_MAPPED = map_outputs(PHENOTYPE_OUTPUTS_FILTERED, PHENOTYPE_OUTPUT_MAPPINGS_FILTERED) PHENOTYPE_TARGETS_ALL = outputs_to_targets( - PHENOTYPE_OUTPUTS, phenotype_wildcard_combos, PHENOTYPE_OUTPUT_MAPPINGS + PHENOTYPE_OUTPUTS_FILTERED, phenotype_wildcard_combos, PHENOTYPE_OUTPUT_MAPPINGS_FILTERED )