diff --git a/biahub/virtual_stain.py b/biahub/virtual_stain.py index bba2fd4a..1e74771d 100644 --- a/biahub/virtual_stain.py +++ b/biahub/virtual_stain.py @@ -1,3 +1,4 @@ +import logging import os import shutil import subprocess @@ -7,13 +8,16 @@ import click import numpy as np import submitit +import yaml from iohub.ngff import open_ome_zarr from iohub.ngff.utils import create_empty_plate from biahub.cli.monitor import monitor_jobs +from biahub.cli.option_eat_all import OptionEatAll from biahub.cli.parsing import ( - input_position_dirpaths, + config_filepath, + init_only, local, monitor, num_processes, @@ -22,7 +26,19 @@ sbatch_filepath_preprocess, sbatch_to_submitit, ) -from biahub.cli.utils import get_submitit_cluster +from biahub.cli.utils import estimate_resources, get_submitit_cluster + +logger = logging.getLogger(__name__) + + +def _optional_validate_paths( + ctx: click.Context, opt: click.Option, value: tuple | None +) -> list[Path] | None: + if value is None or len(value) == 0: + return None + from natsort import natsorted + + return [p for p in map(Path, natsorted(value)) if p.is_dir()] def run_viscy_preprocess( @@ -32,22 +48,6 @@ def run_viscy_preprocess( path_viscy_env: Path = None, verbose: bool = False, ): - """ - Run VisCy preprocess on a single FOV. - - Parameters - ---------- - data_path : str - Path to the FOV data. - num_workers : int - Number of workers to use. - config_file : str - Path to the VisCy config file. - path_viscy_env : Path - Path to the VisCy environment. - verbose : bool - Whether to print verbose output. - """ cmd = ( "module load anaconda && " f"conda activate {path_viscy_env} && " @@ -73,27 +73,7 @@ def run_viscy_predict( path_viscy_env: Path = None, verbose: bool = False, ): - """ - Run VisCy predict on a single FOV. - - Parameters - ---------- - data_path : str - Path to the FOV data. - config_file : str - Path to the VisCy config file. - output_store : str - Path to the output store. - log_dir : str - Path to the log directory. - path_viscy_env : Path - Path to the VisCy environment. - verbose : bool - Whether to print verbose output. - - """ os.chdir(log_dir) - # Compose the shell command cmd = ( "module load anaconda && " f"conda activate {path_viscy_env} && " @@ -116,20 +96,6 @@ def combine_fov_zarrs_to_plate( output_dirpath: Path, cleanup: bool = True, ): - """ - Combine VisCy-predicted FOV Zarrs (in temp) into a single HCS plate Zarr by moving. - - Parameters - ---------- - fovs : list of Path - Original FOV paths (used to extract B/1/000000). - temp_dir : Path - Directory containing the individual .zarr folders (each named like B_1_000000.zarr). - output_dirpath : Path - The plate-level HCS Zarr to merge into. - cleanup : bool - Whether to delete the moved files afterwards. Default True. - """ for fov in fovs: row, col, pos = fov.parts[-3:] nested_fov_path = temp_dir / f"{row}_{col}_{pos}.zarr" / row / col / pos @@ -146,7 +112,6 @@ def combine_fov_zarrs_to_plate( print(f"Moving {nested_fov_path} → {dest_path}") shutil.move(str(nested_fov_path), str(dest_path)) - # Optionally remove the full temp zarr folder if it's now empty if cleanup: try: shutil.rmtree(temp_dir) @@ -156,6 +121,76 @@ def combine_fov_zarrs_to_plate( print(f"Combined all FOVs into {output_dirpath}") +def _init_output_plate( + input_position_dirpaths: list[Path], + output_zarr: Path, + config_filepath: Path, +) -> tuple[int, int, int, int, int]: + """Create the empty virtual-stain output plate. + + Reads target channels from the predict config, creates the output store + with ``_prediction`` channel names, and copies per-position + metadata from the input plate. + + Returns the (T, C_pred, Z, Y, X) output shape. + """ + with open(config_filepath) as f: + cfg = yaml.safe_load(f) + + target_channels = cfg["data"]["init_args"]["target_channel"] + prediction_channels = [f"{ch}_prediction" for ch in target_channels] + + position_keys = [Path(p).parts[-3:] for p in input_position_dirpaths] + with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as ds: + shape = ds.data.shape + scale = ds.scale + T, _, Z, Y, X = shape + + output_shape = (T, len(prediction_channels), Z, Y, X) + + create_empty_plate( + store_path=output_zarr, + position_keys=position_keys, + channel_names=prediction_channels, + shape=output_shape, + scale=scale, + version="0.5", + dtype=np.float32, + metadata_sources=Path(input_position_dirpaths[0]).parents[2], + ) + click.echo( + f"Created {output_zarr} ({len(position_keys)} positions, " + f"channels={prediction_channels})" + ) + return output_shape + + +def _copy_position( + temp_zarr: Path, + output_zarr: Path, + position: str, +) -> None: + """Copy viscy prediction from a temp FOV zarr into the output plate position. + + The temp zarr is a per-position HCS plate produced by VisCy's + HCSPredictionWriter. This reads the data and attributes from the nested + position, writes them into the output plate, then removes the temp zarr. + """ + temp_position = temp_zarr / position + output_position = output_zarr / position + + with open_ome_zarr(str(temp_position), mode="r") as src: + src_data = np.asarray(src[0][:]) + src_attrs = dict(src.zattrs) + + with open_ome_zarr(str(output_position), mode="r+") as dst: + dst[0][:] = src_data + dst.zattrs.update(src_attrs) + + shutil.rmtree(temp_zarr) + click.echo(f"Virtual stain copied: {position}") + + def virtual_stain( input_position_dirpaths: list[str], output_dirpath: str, @@ -170,34 +205,6 @@ def virtual_stain( monitor: bool = True, verbose: bool = True, ): - """ - Run VisCy virtual stain on a plate. - - Parameters - ---------- - input_position_dirpaths : List[str] - List of paths to the input position directories. - output_dirpath : str - Path to the output directory. - predict_config_filepath : str - Path to the VisCy predict config file. - preprocess_config_filepath : str - Path to the VisCy preprocess config file. - path_viscy_env : str - Path to the VisCy environment. - sbatch_filepath_preprocess : str - Path to the VisCy preprocess sbatch file. - sbatch_filepath_predict : str - Path to the VisCy predict sbatch file. - run_mode : str - Which VisCy stage(s) to run. - num_processes : int - Number of processes to use. - local : bool - Whether to run locally. - monitor : bool - - """ output_dirpath = Path(output_dirpath) slurm_out_path = output_dirpath.parent / "slurm_output" @@ -236,9 +243,7 @@ def virtual_stain( ) job_ids_preprocess.append(job) - job_ids = [ - job.job_id for job in job_ids_preprocess - ] # Access job IDs after batch submission + job_ids = [job.job_id for job in job_ids_preprocess] log_path = Path(slurm_out_path / "preprocess" / "submitit_jobs_ids.log") log_path.parent.mkdir(parents=True, exist_ok=True) @@ -299,9 +304,7 @@ def virtual_stain( ) job_ids_predict.append(job) - job_ids = [ - job.job_id for job in job_ids_predict - ] # Access job IDs after batch submission + job_ids = [job.job_id for job in job_ids_predict] log_path = Path(slurm_out_path / "predict" / "submitit_jobs_ids.log") log_path.parent.mkdir(parents=True, exist_ok=True) @@ -353,9 +356,7 @@ def virtual_stain( ) job_ids_combine.append(job) - job_ids = [ - job.job_id for job in job_ids_combine - ] # Access job IDs after batch submission + job_ids = [job.job_id for job in job_ids_combine] log_path = Path(slurm_out_path / "combine" / "submitit_jobs_ids.log") log_path.parent.mkdir(parents=True, exist_ok=True) @@ -370,8 +371,39 @@ def virtual_stain( @click.command("virtual-stain") -@input_position_dirpaths() +@click.option( + "--input-position-dirpaths", + "-i", + required=False, + cls=OptionEatAll, + type=tuple, + callback=_optional_validate_paths, + help="Paths to input positions (required for --init and full runs).", +) @output_dirpath() +@config_filepath() +@init_only() +@click.option( + "--copy", + "copy_mode", + is_flag=True, + default=False, + help="Copy viscy prediction from temp zarr into output plate position.", +) +@click.option( + "--temp-zarr", + "-t", + default=None, + type=click.Path(), + help="Path to temp FOV zarr (required for --copy).", +) +@click.option( + "--position", + "-p", + default=None, + type=str, + help="Position key like B/3/000000 (required for --copy).", +) @sbatch_filepath_preprocess() @sbatch_filepath_predict() @num_processes() @@ -380,7 +412,7 @@ def virtual_stain( @click.option("--verbose", is_flag=True, default=False, help="Verbose output.") @click.option( "--path-viscy-env", - required=True, + default=None, help="Conda environment with VisCy installed.", ) @click.option( @@ -388,11 +420,6 @@ def virtual_stain( type=str, help="Path to the VisCy preprocess config file.", ) -@click.option( - "--predict-config-filepath", - type=str, - help="Path to the VisCy predict config file.", -) @click.option( "--run-mode", type=click.Choice(["all", "preprocess", "predict"]), @@ -400,10 +427,14 @@ def virtual_stain( help="Which VisCy stage(s) to run.", ) def virtual_stain_cli( - input_position_dirpaths: list[str], - output_dirpath: str, - predict_config_filepath: str, - path_viscy_env: str, + input_position_dirpaths: list[Path] | None, + output_dirpath: Path, + config_filepath: Path, + init_only: bool = False, + copy_mode: bool = False, + temp_zarr: str | None = None, + position: str | None = None, + path_viscy_env: str | None = None, preprocess_config_filepath: str = None, run_mode: str = "all", num_processes: int = 32, @@ -413,20 +444,52 @@ def virtual_stain_cli( monitor: bool = True, verbose: bool = True, ): - """Run VisCy virtual staining on a zarr plate from dedicated python environment. - - >>> biahub virtual-stain \ - --input-position-dirpaths path.zarr/*/*/* \ - --output-dirpath output.zarr \ - --predict-config-filepath predict.yml \ - --preprocess-config-filepath preprocess.yml \ - --path-viscy-env /path/to/viscy/env \ - --run-mode all + r"""Run VisCy virtual staining on a zarr plate. + + \b + Initialize the output plate only (Nextflow init step): + >>> biahub virtual-stain --init -i ./input.zarr/*/*/* -c ./predict.yml -o ./output.zarr + + \b + Copy a single position from temp zarr (Nextflow per-position copy step): + >>> biahub virtual-stain --copy -t ./temp/B_3_000000.zarr -o ./output.zarr -p B/3/000000 -c ./predict.yml + + \b + Full SLURM run (preprocess + predict + combine): + >>> biahub virtual-stain -i ./input.zarr/*/*/* -o ./output.zarr \ + -c ./predict.yml --path-viscy-env /path/to/viscy/env --run-mode all """ + if copy_mode: + if temp_zarr is None: + raise click.UsageError("--temp-zarr / -t is required when using --copy.") + if position is None: + raise click.UsageError("--position / -p is required when using --copy.") + _copy_position(Path(temp_zarr), output_dirpath, position) + return + + if not input_position_dirpaths: + raise click.UsageError( + "--input-position-dirpaths / -i is required for --init and full runs." + ) + + if init_only: + output_shape = _init_output_plate( + input_position_dirpaths, output_dirpath, config_filepath + ) + + num_cpus, mem_per_cpu = estimate_resources( + shape=output_shape, ram_multiplier=16, max_num_cpus=16 + ) + click.echo(f"RESOURCES:{num_cpus} {num_cpus * mem_per_cpu}") + return + + if path_viscy_env is None: + raise click.UsageError("--path-viscy-env is required for full virtual stain runs.") + virtual_stain( input_position_dirpaths=input_position_dirpaths, - output_dirpath=output_dirpath, - predict_config_filepath=predict_config_filepath, + output_dirpath=str(output_dirpath), + predict_config_filepath=str(config_filepath), preprocess_config_filepath=preprocess_config_filepath, sbatch_filepath_preprocess=sbatch_filepath_preprocess, sbatch_filepath_predict=sbatch_filepath_predict, diff --git a/nextflow/modules/virtual_stain.nf b/nextflow/modules/virtual_stain.nf new file mode 100644 index 00000000..fc51fd8f --- /dev/null +++ b/nextflow/modules/virtual_stain.nf @@ -0,0 +1,154 @@ +// Virtual stain subworkflow: init → preprocess → fan-out (predict + copy) × N positions. +// +// This subworkflow is PATH-AGNOSTIC. Callers pass the input zarr, output zarr, +// and config explicitly. Temp paths for predict output are derived from the +// output_zarr's parent directory. +// +// Virtual staining uses VisCy (external tool) for the GPU work, so the +// preprocess and predict steps call viscy directly rather than biahub CLI. +// Only init and copy are biahub commands: +// +// 1. init_virtual_stain: creates empty output plate with prediction channels, +// cleans temp dir, emits RESOURCES: +// 2. run_virtual_stain_preprocess: calls `viscy preprocess` on the whole plate +// 3. run_virtual_stain: calls `viscy predict` per position, then +// `biahub virtual-stain --copy` to merge the temp FOV zarr into the plate +// +// The predict step writes to a temp per-position zarr, and the --copy step +// moves data from that temp zarr into the output plate. + +include { parse_resources; biahub_cmd; slurm_logs; slurm_log_dir } from './common' + +def viscy_cmd() { + return params.viscy_project ? + "uv run --project ${params.viscy_project} viscy" : + "uv run --from 'viscy @ git+https://github.com/mehta-lab/VisCy@v0.3.4' viscy" +} + + +process init_virtual_stain { + label 'cpu_local' + + input: + val input_zarr + val output_zarr + val config + val output_dir + val trigger + + output: + stdout + + script: + """ + mkdir -p "${slurm_log_dir('virtual_stain')}" + rm -rf "${output_dir}/temp" + ${biahub_cmd()} virtual-stain --init \ + -i "${input_zarr}"/*/*/* \ + -o "${output_zarr}" \ + -c "${config}" + """ +} + +process run_virtual_stain_preprocess { + label 'cpu' + clusterOptions { slurm_logs('virtual_stain') } + cpus 16 + memory { "${64 * task.attempt} GB" } + time '1h' + maxRetries 1 + errorStrategy 'retry' + + input: + val input_zarr + val trigger + + output: + val true + + script: + """ + ${viscy_cmd()} preprocess \ + --data_path "${input_zarr}" \ + --channel_names -1 \ + --num_workers ${task.cpus} \ + --block_size 32 + """ +} + +process run_virtual_stain { + tag "${position}" + label 'gpu' + clusterOptions { "--gres=gpu:1 " + slurm_logs('virtual_stain') } + maxForks 30 + cpus { meta.cpus } + memory { "${meta.mem_gb} GB" } + time { task.attempt == 1 ? '8h' : '12h' } + maxRetries 2 + errorStrategy 'retry' + + input: + tuple val(position), val(meta) + val input_zarr + val output_zarr + val config + val output_dir + + output: + val position + + script: + def temp_zarr = "${output_dir}/temp/${position.replaceAll('/', '_')}.zarr" + """ + rm -rf "${temp_zarr}" + + ${viscy_cmd()} predict \ + -c "${config}" \ + --data.init_args.data_path "${input_zarr}/${position}" \ + --data.init_args.num_workers 0 \ + --trainer.callbacks+=viscy_utils.callbacks.prediction_writer.HCSPredictionWriter \ + --trainer.callbacks.output_store "${temp_zarr}" \ + --trainer.default_root_dir "${output_dir}/logs" + + ${biahub_cmd()} virtual-stain --copy \ + -t "${temp_zarr}" \ + -o "${output_zarr}" \ + -p "${position}" \ + -c "${config}" + """ +} + + +// take: +// positions collected channel of position keys (e.g. ['A/1/0', 'B/1/0']) +// input_zarr path to the input plate.zarr (reconstruct output) +// output_zarr path to the virtual stain output plate.zarr +// config path to the predict settings YAML +// prev_done gating channel — virtual stain starts once this emits +workflow virtual_stain_wf { + take: + positions + input_zarr + output_zarr + config + prev_done + + main: + def output_dir = new File(output_zarr).parent + + resources = init_virtual_stain(input_zarr, output_zarr, config, output_dir, prev_done.map { 'done' }) + .map { parse_resources(it) } + vs_preprocess = run_virtual_stain_preprocess(input_zarr, prev_done.map { 'done' }) + + ready = resources.combine(vs_preprocess) + + pos_meta = positions + .flatMap { it } + .combine(ready) + .map { pos, meta, preprocess_done -> [pos, meta] } + + vs_done = run_virtual_stain(pos_meta, input_zarr, output_zarr, config, output_dir) | collect + + emit: + done = vs_done +} diff --git a/tests/test_cli/test_virtual_stain_cli.py b/tests/test_cli/test_virtual_stain_cli.py new file mode 100644 index 00000000..59a1fc1c --- /dev/null +++ b/tests/test_cli/test_virtual_stain_cli.py @@ -0,0 +1,137 @@ +import numpy as np +import pytest +import yaml + +from click.testing import CliRunner +from iohub.ngff import open_ome_zarr +from iohub.ngff.utils import create_empty_plate + +from biahub.cli.main import cli + + +@pytest.fixture() +def predict_config(tmp_path): + cfg = { + "data": { + "init_args": { + "source_channel": "Phase3D", + "target_channel": ["nuclei", "membrane"], + "z_window_size": 15, + "batch_size": 1, + "num_workers": 0, + } + } + } + config_path = tmp_path / "predict.yml" + config_path.write_text(yaml.dump(cfg, default_flow_style=False)) + return config_path + + +@pytest.fixture() +def vs_input_plate(tmp_path): + plate_path = tmp_path / "input.zarr" + position_list = (("B", "3", "000000"), ("B", "3", "000001")) + + plate = open_ome_zarr( + plate_path, + layout="hcs", + mode="w", + channel_names=["Phase3D"], + ) + + for row, col, fov in position_list: + position = plate.create_position(row, col, fov) + position["0"] = np.random.uniform(1.0, 100.0, size=(1, 1, 15, 32, 32)).astype( + np.float32 + ) + + plate.close() + return plate_path + + +def test_virtual_stain_cli_init_only(tmp_path, vs_input_plate, predict_config): + output_path = tmp_path / "output.zarr" + + runner = CliRunner() + result = runner.invoke( + cli, + [ + "virtual-stain", + "--init", + "-i", + str(vs_input_plate) + "/B/3/000000", + str(vs_input_plate) + "/B/3/000001", + "-c", + str(predict_config), + "-o", + str(output_path), + ], + ) + + assert result.exit_code == 0, result.output + assert output_path.exists() + assert "RESOURCES:" in result.output + assert "nuclei_prediction" in result.output + assert "membrane_prediction" in result.output + + with open_ome_zarr(str(output_path / "B" / "3" / "000000"), mode="r") as ds: + assert ds.channel_names == ["nuclei_prediction", "membrane_prediction"] + assert ds.data.shape == (1, 2, 15, 32, 32) + + with open_ome_zarr(str(output_path / "B" / "3" / "000001"), mode="r") as ds: + assert ds.channel_names == ["nuclei_prediction", "membrane_prediction"] + + +def test_virtual_stain_cli_copy(tmp_path, predict_config): + output_path = tmp_path / "output.zarr" + + create_empty_plate( + store_path=output_path, + position_keys=[("B", "3", "000000")], + channel_names=["nuclei_prediction", "membrane_prediction"], + shape=(1, 2, 5, 8, 8), + scale=(1.0, 1.0, 1.0, 0.5, 0.5), + version="0.5", + dtype=np.float32, + ) + + temp_zarr = tmp_path / "temp" / "B_3_000000.zarr" + create_empty_plate( + store_path=temp_zarr, + position_keys=[("B", "3", "000000")], + channel_names=["nuclei_prediction", "membrane_prediction"], + shape=(1, 2, 5, 8, 8), + scale=(1.0, 1.0, 1.0, 0.5, 0.5), + version="0.5", + dtype=np.float32, + ) + + test_data = np.random.uniform(0, 1, (1, 2, 5, 8, 8)).astype(np.float32) + with open_ome_zarr(str(temp_zarr / "B" / "3" / "000000"), mode="r+") as ds: + ds[0][:] = test_data + + runner = CliRunner() + result = runner.invoke( + cli, + [ + "virtual-stain", + "--copy", + "-t", + str(temp_zarr), + "-o", + str(output_path), + "-p", + "B/3/000000", + "-c", + str(predict_config), + ], + ) + + assert result.exit_code == 0, result.output + assert "Virtual stain copied: B/3/000000" in result.output + + with open_ome_zarr(str(output_path / "B" / "3" / "000000"), mode="r") as ds: + copied = ds[0][:] + np.testing.assert_array_equal(copied, test_data) + + assert not temp_zarr.exists(), "temp zarr should be cleaned up"