diff --git a/pyproject.toml b/pyproject.toml index 9af0fbe..3057bb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,8 @@ pyarrow = ">=18.0.0,<19" [tool.pixi.pypi-dependencies] spimquant = { path = ".", editable = true } -zarrnii = ">=0.21.1,<0.22.0" +#zarrnii = ">=0.21.1,<0.22.0" +zarrnii = { git = "https://github.com/khanlab/zarrnii", branch = "copilot/add-scale-and-offset-options" } #zarrnii = { git = "https://github.com/khanlab/zarrnii", rev = "resize-local-mean", editable = true } #zarrnii = { path = "/nfs/khan/trainees/akhan488/apps/zarrnii", editable = true } vesselfm = { git = "https://github.com/khanlab/vesselfm", rev = "504baab" } diff --git a/spimquant/workflow/rules/segmentation.smk b/spimquant/workflow/rules/segmentation.smk index 9c06453..50635f3 100644 --- a/spimquant/workflow/rules/segmentation.smk +++ b/spimquant/workflow/rules/segmentation.smk @@ -129,15 +129,16 @@ rule n4_pre_quant: rule calc_n4_rescaling: """ calculate the linear intensity rescaling, - scale and offset, that n4 uses to rescale + scale and offset, that n4 uses to rescale intensities back to the input min and max. - Does this by calculating the input min and max (within + Does this by calculating the input min and max (within the masked region), then calculating the min and max in the image - input image divided by the bias field (also only within the masked - region). Then use these min/max values to find the scale and offset + input image divided by the bias field (also only within the masked + region). Then use these min/max values to find the scale and offset (e.g. x*scale + offset) to apply to the divided image, so that the min and max become the input min and max. """ - input: + + input: uncorr=bids( root=root, datatype="micr", @@ -154,6 +155,7 @@ rule calc_n4_rescaling: desc="brain", suffix="mask.nii.gz", **inputs["spim"].wildcards, + ), biasfield=bids( root=root, datatype="micr", @@ -173,8 +175,11 @@ rule calc_n4_rescaling: suffix="scaleoffset.txt", **inputs["spim"].wildcards, ), + resources: + mem_mb=32000, + runtime=15, script: - "../scripts/calc_n4_rescaling.py" #TODO: make this script + "../scripts/calc_n4_rescaling.py" rule encode_mask_in_bias_field: """Use -1 to encode the mask in the bias field""" diff --git a/spimquant/workflow/scripts/calc_n4_rescaling.py b/spimquant/workflow/scripts/calc_n4_rescaling.py new file mode 100644 index 0000000..4335551 --- /dev/null +++ b/spimquant/workflow/scripts/calc_n4_rescaling.py @@ -0,0 +1,42 @@ +import nibabel as nib +import numpy as np + +uncorr_img = nib.load(snakemake.input.uncorr) +mask_img = nib.load(snakemake.input.mask) +biasfield_img = nib.load(snakemake.input.biasfield) + +uncorr = uncorr_img.get_fdata(dtype=np.float32) +mask = mask_img.get_fdata(dtype=np.float32) > 0 +biasfield = biasfield_img.get_fdata(dtype=np.float32) + +if not mask.any(): + raise ValueError( + f"Brain mask {snakemake.input.mask} contains no masked voxels; " + "cannot compute scale and offset." + ) + +# Avoid division by zero in bias field +epsilon = np.finfo(np.float32).eps +biasfield_safe = np.maximum(biasfield, epsilon) + +# Corrected image = uncorrected / bias field (mimics what N4BiasFieldApply does) +corrected = uncorr / biasfield_safe + +# Compute min/max within the masked region +input_min = float(uncorr[mask].min()) +input_max = float(uncorr[mask].max()) +corrected_min = float(corrected[mask].min()) +corrected_max = float(corrected[mask].max()) + +# Linear mapping: corrected * scale + offset -> [input_min, input_max] +corrected_range = corrected_max - corrected_min +if corrected_range == 0: + scale = 1.0 + offset = 0.0 +else: + scale = (input_max - input_min) / corrected_range + offset = input_min - scale * corrected_min + +with open(snakemake.output.scale_offset_params, "w") as f: + f.write(f"scale={scale}\n") + f.write(f"offset={offset}\n") diff --git a/spimquant/workflow/scripts/n4_biasfield.py b/spimquant/workflow/scripts/n4_biasfield.py index a719fe1..cedf8e6 100644 --- a/spimquant/workflow/scripts/n4_biasfield.py +++ b/spimquant/workflow/scripts/n4_biasfield.py @@ -6,6 +6,23 @@ is_imaris = str(snakemake.input.spim).lower().endswith(".ims") + # Read scale and offset from the pre-computed rescaling parameters file + scale_offset = {} + with open(snakemake.input.scale_offset_params) as f: + for line in f: + line = line.strip() + if not line or "=" not in line: + continue + key, val = line.split("=", 1) + scale_offset[key.strip()] = float(val.strip()) + if "scale" not in scale_offset or "offset" not in scale_offset: + raise ValueError( + f"scale_offset_params file {snakemake.input.scale_offset_params} " + "must contain 'scale' and 'offset' entries" + ) + scale = scale_offset["scale"] + offset = scale_offset["offset"] + with get_dask_client( snakemake.config["dask_scheduler"], snakemake.threads, @@ -32,11 +49,9 @@ # scaled_proc_kwargs controls how apply_scaled_processing is performed - # Apply bias field correction + # Apply bias field correction with linear rescaling to restore input intensity range znimg_corrected = znimg.apply_scaled_processing( - N4BiasFieldApply(log_space=True, - scale=#TODO get scale from txt file, - offset=#TODO get scale from txt file), + N4BiasFieldApply(log_space=True, scale=scale, offset=offset), **scaled_proc_kwargs, )