[ICASSP 2025 Oral] DiffKillR

Killing and Recreating Diffeomorphisms for Cell Annotation in Dense Microscopy Images

Krishnaswamy Lab, Yale University

This is the authors' PyTorch implementation of DiffKillR, ICASSP 2025 (Oral).

The official version is maintained in the Lab GitHub repo.

Citation

``` @inproceedings{liu2025diffkillr, title={Diffkillr: Killing and recreating diffeomorphisms for cell annotation in dense microscopy images}, author={Liu, Chen and Liao, Danqi and Parada-Mayorga, Alejandro and Ribeiro, Alejandro and DiStasio, Marcello and Krishnaswamy, Smita}, booktitle={ICASSP 2025-2025 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)}, year={2025}, organization={IEEE} }

@article{liu2024diffkillr, title={Diffkillr: Killing and recreating diffeomorphisms for cell annotation in dense microscopy images}, author={Liu, Chen and Liao, Danqi and Parada-Mayorga, Alejandro and Ribeiro, Alejandro and DiStasio, Marcello and Krishnaswamy, Smita}, journal={arXiv preprint arXiv:2410.03058}, year={2024} } ```

Usage

Preprocess datasets. cd src/preprocessing python preprocess_MoNuSeg.py python preprocess_A28.py python preprocess_A28_axis.py

Orientation Prediction

Train and test DiffeoInvariantNet. (Remove --use-wandb if you don't want to use Weights and Biases.) cd src/ python main_DiffeoInvariantNet.py --dataset-name A28 --dataset-path '$ROOT/data/A28-87_CP_lvl1_HandE_1_Merged_RAW_ch00_patch_96x96/' --DiffeoInvariantNet-model AutoEncoder --use-wandb --wandb-username yale-cl2482

Train and test DiffeoMappingNet. cd src/ python main_DiffeoMappingNet.py --dataset-name A28 --dataset-path '$ROOT/data/A28-87_CP_lvl1_HandE_1_Merged_RAW_ch00_patch_96x96/' --DiffeoMappingNet-model VoxelMorph --use-wandb --wandb-username yale-cl2482

Segmentation on MoNuSeg

Train and test DiffeoInvariantNet. cd src/ python main_DiffeoInvariantNet.py --dataset-name MoNuSeg --dataset-path '$ROOT/data/MoNuSeg/MoNuSegByCancer_patch_96x96/' --organ Breast --percentage 10 --random-seed 1

Train and test DiffeoMappingNet. cd src/ python main_DiffeoMappingNet.py --dataset-name MoNuSeg --dataset-path '$ROOT/data/MoNuSeg/MoNuSegByCancer_patch_96x96/' --organ Breast --percentage 10 --random-seed 1

Run Inference python main_inference_segmentation.py --organ Breast --use-gt-loc --random-seed 1 python main_inference_segmentation.py --organ Breast --random-seed 1

Comparison

1. First train/infer the models.

1. Then, stitch the images and run evaluation. cd comparison/eval/ python stitch_patches.py python evaluate_monuseg.py python evaluate_glysac.py

Preparation

To use SAM.

```

under comparison/SAM/checkpoints/

wget https://dl.fbaipublicfiles.com/segmentanything/samvith4b8939.pth ```

To Use SAM2.

```

under comparison/SAM2/checkpoints/

wget https://dl.fbaipublicfiles.com/segmentanything2/072824/sam2hieralarge.pt ```

To use MedSAM.

```

under comparison/MedSAM/checkpoints/

download from https://drive.google.com/file/d/1ARiB5RkSsWmAB_8mqWnwDF8ZKTtFwsjl/view ```

To use SAM-Med2D.

```

under comparison/SAM_Med2D/checkpoints/

download from https://drive.google.com/file/d/1ARiB5RkSsWmAB_8mqWnwDF8ZKTtFwsjl/view ```

Environment

We developed the codebase in a miniconda environment. Tested on Python 3.9.13 + PyTorch 1.12.1. How we created the conda environment: ```

Optional: Update to libmamba solver.

conda update -n base conda conda install -n base conda-libmamba-solver conda config --set solver libmamba

conda create --name cellseg pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch -c nvidia -c anaconda -c conda-forge

conda activate cellseg conda install -c anaconda scikit-image scikit-learn pillow matplotlib seaborn tqdm

conda install -c conda-forge libstdcxx-ng=12

python -m pip install antspyx

python -m pip install dipy

python -m pip install opencv-python python -m pip install python-dotenv python -m pip install simple-lama-inpainting

MoNuSeg

python -m pip install xmltodict

PSM

python -m pip install tensorboardX python -m pip install shapely python -m pip install ml_collections python -m pip install ttach

LACSS

python -m pip install --upgrade "jax[cuda11pip]" -f https://storage.googleapis.com/jax-releases/jaxcuda_releases.html python -m pip install --upgrade flax python -m pip install --upgrade lacss

python -m pip install ml_dtypes==0.2.0

StarDist

python -m pip install stardist python -m pip install tensorflow

For SAM

python -m pip install git+https://github.com/facebookresearch/segment-anything.git

For SAM2

python -m pip install git+https://github.com/facebookresearch/segment-anything-2.git

For MedSAM

python -m pip install git+https://github.com/bowang-lab/MedSAM.git

For SAM-Med2D

python -m pip install albumentations python -m pip install scikit-learn==1.1.3 # need to downgrade to 1.1.3

Export CuDNN

echo 'CUDNNPATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.file)"))' >> $CONDAPREFIX/etc/conda/activate.d/env_vars.sh

echo 'export LDLIBRARYPATH=$CONDAPREFIX/lib/:$CUDNNPATH/lib:$LDLIBRARYPATH' >> $CONDAPREFIX/etc/conda/activate.d/envvars.sh

source $CONDAPREFIX/etc/conda/activate.d/envvars.sh

```