WIP - initial stardist test

pyproject.toml

@@ -12,6 +12,7 @@ requires = ["hatchling"]
 [tool.pixi.workspace]
 channels = ["conda-forge", "bioconda"]
 platforms = ["linux-64"]
+preview = ["pixi-build"]
 
 [tool.pixi.feature.runtime.dependencies]
 snakemake = ">=9.5.1,<10"
@@ -27,17 +28,16 @@ conda = ">=25.5.0,<26"
 snakemake-executor-plugin-slurm = ">=1.4.0,<2"
 snakemake-storage-plugin-fs = ">=1.1.2,<2"
 snakemake-storage-plugin-http = ">=0.3.0,<0.4.0"
-greedyreg = "==0.0.20250712"
 antspyx = "==0.6.1"
 markdown-it-py = "<4"
 docutils = "<0.22"
 
 [tool.pixi.pypi-dependencies]
 spimquant = { path = ".", editable = true }
 ngff-zarr = { version = ">=0.19.1", extras = ["cli","dask-image","itk","tensorstore","validate"] }
-zarrnii = ">=0.14.0-alpha1,<0.15.0"
+#zarrnii = ">=0.14.0-alpha1,<0.15.0"
 #zarrnii = { git = "https://github.com/khanlab/zarrnii", rev = "v0.13.0-alpha1",  editable = true }
-#zarrnii = { path = "./zarrnii", editable = true }
+zarrnii = { path = "./zarrnii", editable = true }
 
 [tool.pixi.feature.dev.dependencies]
 black   = ">=24.0.0, <25.0.0"
@@ -51,10 +51,23 @@ jupyterlab = ">=4.4.7,<5"
 napari = ">=0.6.6,<0.7"
 napari-ome-zarr = ">=0.6.1,<0.7"
 pyqt = ">=5.15.11,<6"
+jupyterlab = ">=4.4.10,<5"
+seaborn = ">=0.13.2,<0.14"
+
+[tool.pixi.feature.experimental.dependencies]
 micro_sam = ">=1.6.2,<2"
+cellpose = ">=4.0.7,<5"
+
+
+[tool.pixi.feature.stardist.dependencies]
+tensorflow = "*"
+
+[tool.pixi.feature.stardist.pypi-dependencies]
+stardist = { path = "/local/scratch/stardist", editable = true }
+
 
 [tool.pixi.environments.default]
-features = ["runtime"]
+features = ["runtime","stardist"]
 
 [tool.pixi.environments.dev]
 features = ["runtime", "dev", "vis"]
@@ -63,10 +76,14 @@ features = ["runtime", "dev", "vis"]
 [tool.pixi.environments.dev-only]
 features = ["dev"]
 
+[tool.pixi.environments.experimental]
+features = ["runtime","vis","experimental"]
+
+
+
+
 [tool.pixi.tasks]
 quality_check = "isort --profile black spimquant/*.py -c && black spimquant --check && snakefmt spimquant --check"
 quality_fix = "isort --profile black spimquant/*.py && black spimquant  && snakefmt spimquant"
 
-[tool.pixi.dependencies]
-jupyterlab = ">=4.4.10,<5"
-pip = ">=25.3,<26"
+

spimquant/config/snakebids.yml

@@ -302,4 +302,8 @@ seg_metrics:
   - count
   - volume
   - nvoxels
+  - equivdiameter
 
+stardist:
+  model_urls:
+    3D_demo: https://github.com/stardist/stardist-models/releases/download/v0.1/python_3D_demo.zip

spimquant/workflow/Snakefile

@@ -299,7 +299,7 @@ include: "rules/masking.smk"
 include: "rules/templatereg.smk"
 include: "rules/segmentation.smk"
 include: "rules/patches.smk"
-
+include: "rules/stardist.smk"
 
 if config["register_to_mri"]:
 

spimquant/workflow/envs/greedy.yaml

@@ -0,0 +1,5 @@
+name: greedy
+channels:
+  - conda-forge
+dependencies:
+  - greedyreg=0.0.20250712

spimquant/workflow/envs/stardist.yaml

@@ -0,0 +1,13 @@
+name: stardist
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.9
+  - stardist=0.9.1
+  - dask
+  - zarr
+  - pip
+  - pip:
+    - zarrnii
+

spimquant/workflow/rules/masking.smk

@@ -189,6 +189,7 @@ rule init_affine_reg:
     resources:
         mem_mb=16000,
         runtime=5,
+    conda: "../envs/greedy.yaml"
     shell:
         "greedy -threads {threads} -d 3 -i {input.template} {input.subject} "
         " -a -dof 12 -ia-image-centers -m NMI -o {output.xfm_ras} -n {params.iters} && "
@@ -226,6 +227,7 @@ rule affine_transform_template_mask_to_subject:
     resources:
         mem_mb=16000,
         runtime=5,
+    conda: "../envs/greedy.yaml"
     shell:
         " greedy -threads {threads} -d 3 -rf {input.ref} "
         " -ri NN"

spimquant/workflow/rules/preproc_mri.smk

@@ -124,6 +124,7 @@ rule rigid_greedy_reg_mri_to_template:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/greedy.yaml"
     shell:
         "greedy -threads {threads} -d 3 -i {input.template} {input.subject} "
         " -a -dof 6 -ia-image-centers -m {params.metric} -o {output.xfm_ras} && "
@@ -274,6 +275,7 @@ rule apply_mri_brain_mask:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/c3d.yaml"
     shell:
         "c3d {input.nii} {input.mask} -multiply -resample 300% -o {output.nii}"
 
@@ -383,6 +385,7 @@ rule rigid_greedy_reg_mri_to_spim:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/greedy.yaml"
     shell:
         "greedy -threads {threads} -d 3 -i {input.spim} {input.mri} "
         " -a -dof {params.dof} -ia-image-centers -m {params.metric_rigid} -o {output.xfm_ras} && "
@@ -507,6 +510,7 @@ rule warp_mri_to_template_via_spim:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/greedy.yaml"
     shell:
         " greedy -threads {threads} -d 3 -rf {input.ref} "
         "  -rm {input.mri} {output.warped} "
@@ -587,6 +591,7 @@ rule warp_mri_brainmask_to_spim:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/greedy.yaml"
     shell:
         " greedy -threads {threads} -d 3 -rf {input.ref} -ri NN"
         "  -rm {input.mask} {output.mask} "

spimquant/workflow/rules/segmentation.smk

@@ -387,6 +387,7 @@ rule deform_negative_mask_to_subject_nii:
     resources:
         mem_mb=16000,
         runtime=15,
+    conda: "../envs/greedy.yaml"
     shell:
         " greedy -threads {threads} -d 3 -rf {input.ref} "
         " -ri NN "

spimquant/workflow/rules/stardist.smk

@@ -0,0 +1,59 @@
+rule get_model:
+    input:
+        model_zip=lambda wildcards: storage(config['stardist']['model_urls'][wildcards.model_name])
+    output:
+        model_dir=directory('resources/stardist_models/{model_name}')
+    shell:
+        'unzip -d {output.model_dir} {input.model_zip}'
+
+rule test_stardist:
+    input:
+        model_dir=directory('resources/stardist_models/3D_demo')
+    output:
+        out_dir=directory('test_stardist')
+    script: '../scripts/test_stardist.py'
+
+rule segment_stardist:
+    input:
+        spim=inputs["spim"].path,
+    params:
+        zarrnii_kwargs=lambda wildcards: {"orientation": config["orientation"],channel_labels=[wildcards.stain]},
+    output:
+        regionprops=temp(
+            directory(
+                bids(
+                    root=work,
+                    datatype="micr",
+                    stain="{stain}",
+                    level="{level}",
+                    desc="correctedn4",
+                    suffix="SPIM.ome.zarr",
+                    **inputs["spim"].wildcards,
+                )
+            )
+        ),
+        biasfield=temp(
+            directory(
+                bids(
+                    root=work,
+                    datatype="micr",
+                    stain="{stain}",
+                    level="{level}",
+                    desc="n4",
+                    suffix="biasfield.ome.zarr",
+                    **inputs["spim"].wildcards,
+                )
+            )
+        ),
+    group:
+        "subj"
+    threads: 128
+    resources:
+        mem_mb=500000,
+        disk_mb=2097152,
+        runtime=60,
+    script:
+        "../scripts/n4_biasfield.py"
+
+
+

spimquant/workflow/rules/templatereg.smk

@@ -162,6 +162,7 @@ rule affine_reg:
     resources:
         mem_mb=16000,
         runtime=5,
+    conda: "../envs/greedy.yaml"
     shell:
         "greedy -threads {threads} -d 3 -i {input.template} {input.subject} "
         " -a -dof 12 -ia-image-centers -m NMI -o {output.xfm_ras} -n {params.iters} && "
@@ -265,6 +266,7 @@ rule deform_reg:
     resources:
         mem_mb=16000,
         runtime=5,
+    conda: "../envs/greedy.yaml"
     shell:
         "greedy -threads {threads} -d 3 -i {input.template} {input.subject} "
         " -it {input.xfm_ras} -m {params.metric} "

spimquant/workflow/scripts/test_stardist.py

@@ -0,0 +1,43 @@
+if __name__ == "__main__":
+
+    from dask.distributed import Client, LocalCluster
+
+    cluster = LocalCluster(
+        n_workers=int(snakemake.threads / 2),  # or 32, depending on workload
+        threads_per_worker=2,  # isolate GIL
+        memory_limit="auto",  # or tune to your RAM
+        dashboard_address=":8788",
+    )
+    client = Client(cluster)
+    print(cluster.dashboard_link)
+
+
+    from zarrnii import ZarrNii
+
+    import zarrnii
+    znimg = ZarrNii.from_ome_zarr('/nfs/trident3/lightsheet/prado/mouse_app_lecanemab_ki3/bids/sub-AS134F1/micr/sub-AS134F1_sample-brain_acq-imaris4x_SPIM.ome.zarr')
+    print(znimg)
+
+    import nibabel as nib
+    import numpy as np
+    from stardist.models import StarDist3D
+    from csbdeep.utils import normalize
+    from stardist.data import test_image_nuclei_3d
+    from pathlib import Path
+
+
+    model = StarDist3D(None, name=snakemake.input.model_dir)
+
+    img = test_image_nuclei_3d()
+
+
+
+    labels,_ = model.predict_instances(normalize(img))
+
+
+    out_dir_path = Path(snakemake.output.out_dir)
+    out_dir_path.mkdir()
+
+    nib.Nifti1Image(img,affine=np.eye(4,4)).to_filename(out_dir_path / 'img.nii')
+    nib.Nifti1Image(labels,affine=np.eye(4,4)).to_filename(out_dir_path / 'labels.nii')
+