Complete N4 bias field rescaling pipeline

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" }

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"""

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")

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,
         )