diff --git a/.gitignore b/.gitignore index 6647f0e5..a23b1067 100644 --- a/.gitignore +++ b/.gitignore @@ -19,4 +19,5 @@ scripts/cooper/revision/evaluation_results/ scripts/cooper/revision/export_tif_to_h5.py scripts/cooper/revision/copy_path.py scripts/cooper/revision/view.py -scripts/cooper/revision/count_AZ.py \ No newline at end of file +scripts/cooper/revision/count_AZ.py +tmp_*/ diff --git a/CLAUDE.md b/CLAUDE.md index 85780dc5..eed2c393 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,7 +51,7 @@ synapse_net.visualize_vesicle_pools -h The package is organized by capability. The big picture spans several modules: ### Models & the model registry (`synapse_net/inference/inference.py`) -Models are identified by a **`model_type` string** (`vesicles_3d`, `vesicles_2d`, `vesicles_cryo`, `active_zone`, `compartments`, `mitochondria`/`mitochondria2`, `cristae`/`cristae2`/`cristae3`, `ribbon`, plus CLI-only `vesicles_*` variants). `_get_model_registry()` maps each name to a sha256 + GWDG ownCloud download URL and fetches via `pooch`. `get_model(model_type)` downloads (if needed) and `torch.load`s the checkpoint. +Models are identified by a **`model_type` string** (`vesicles_3d`, `vesicles_2d`, `vesicles_cryo`, `active_zone`, `compartments`, `mitochondria`/`mitochondria2`, `cristae`/`cristae2`/`cristae3`/`cristae4`/`cristae5`, `ribbon`, plus CLI-only `vesicles_*` variants). `_get_model_registry()` maps each name to a sha256 + GWDG ownCloud download URL and fetches via `pooch`. `get_model(model_type)` downloads (if needed) and `torch.load`s the checkpoint. **Voxel-size scaling is central.** Each model was trained at a specific resolution (`get_model_training_resolution`, in nm/axis). `compute_scale_from_voxel_size(voxel_size, model_type)` produces a zyx scale factor so input data is resized to match training resolution before inference and resized back after. Always pair a model with the right `scale` rather than feeding raw-resolution data. @@ -70,7 +70,12 @@ Per-structure **post-processing** lives in `inference/postprocessing/` (vesicles Built on `torch_em` U-Nets (`UNet2d`, `AnisotropicUNet`). Two regimes: - `supervised_training.py` — needs data + manual labels. Can init weights from a pretrained `model_type`. - `domain_adaptation.py` — unsupervised student-teacher (mean-teacher) adaptation to a new condition without labels; only works if the source model already partially detects the structure. -- `semisupervised_training.py`, `transform.py` — supporting pieces. +- `semisupervised_training.py`, `transform.py`, `loss.py`, `split.py` — supporting pieces. + +Per-structure recipes sit on top of `supervised_training.py`, each fixing the hyperparameters and data layout of one published model and exposing a `main()` behind a console script: +- `mitochondria.py` — mitochondria in electron tomography (`mitochondria2`). +- `cristae.py` — cristae in electron tomography (`cristae5`); two input channels, masked loss, membrane weighting. +- `mitochondria_vol_em.py` — mitochondria in volume EM (FIB-SEM, 25 nm z / 5 nm xy); two anisotropic U-Net levels, no normalization layers, white-filler removal in the raw transform. **Not** in the model registry. ### Tools (`synapse_net/tools/`) - `cli.py` — argparse wrappers behind the console scripts. diff --git a/doc/start_page.md b/doc/start_page.md index a04baf4d..1413f6ab 100644 --- a/doc/start_page.md +++ b/doc/start_page.md @@ -162,7 +162,7 @@ For more options supported by the IMOD exports, please run `synapse_net.export_t > Note: to use these commands you have to install IMOD. -SynapseNet also provides two CLI comamnds for training models, one for supervised network training (see [Supervised Training](#supervised-training) for details) and one for domain adaptation (see [Domain Adaptation](#domain-adaptation) for details). +SynapseNet also provides five CLI comamnds for training models: one for supervised network training (see [Supervised Training](#supervised-training) for details), one for training a mitochondria model for electron tomography (see [Mitochondria Training](#mitochondria-training) for details), one for training a cristae model (see [Cristae Training](#cristae-training) for details), one for training a mitochondria model for volume EM (see [Volume EM Mitochondria Training](#volume-em-mitochondria-training) for details) and one for domain adaptation (see [Domain Adaptation](#domain-adaptation) for details). ## Python Library @@ -211,6 +211,117 @@ synapse_net.run_supervised_training -h for more information and instructions on how to use the command. +### Mitochondria Training + +SynapseNet provides a dedicated function for training a model for mitochondria segmentation in electron tomograms. +It reproduces the recipe that was used for our `mitochondria2` model and is implemented in +`synapse_net.training.mitochondria`. Compared to the generic supervised training it fixes the hyperparameters of that +model, expects the tomograms and annotations in the same hdf5 file, and normalizes the tomograms with the 1st and 99th +percentile instead of standardizing them. + +We also provide a command line function to run it: `synapse_net.run_mitochondria_training`. +It expects one hdf5 file per tomogram, which contains the tomogram and the mitochondria annotations: +```bash +synapse_net.run_mitochondria_training \ + -n my-mito-model \ # The name of the model checkpoint. + -i /path/to/tomograms \ # The folder with the hdf5 files, which is searched recursively. + --raw_key raw \ # The internal path of the tomogram, 'raw' by default. + --label_key labels/mitochondria \ # The internal path of the annotations, 'labels/mitochondria' by default. + --patch_shape 32 256 256 \ # The patch shape in ZYX. + --batch_size 8 \ # The batch size for training. + --n_iterations 150000 \ # The maximal number of iterations to train for. +``` +The data is split into training and validation data randomly, with a fixed seed. Pass `--split_file` to use an explicit +split instead. The default batch size and patch shape require a GPU with a lot of memory; reduce them and pass +`--mixed_precision` to train on a smaller GPU. + +Note that the normalization is part of the model: to segment with the resulting model you have to pass +`preprocess=torch_em.transform.raw.normalize_percentile` to `synapse_net.inference.mitochondria.segment_mitochondria`. + +Run +```bash +synapse_net.run_mitochondria_training -h +``` +for more information and instructions on how to use the command. + + +### Cristae Training + +SynapseNet provides a dedicated function for training a model for cristae segmentation in electron tomograms. +It reproduces the recipe that was used for our `cristae5` model and is implemented in `synapse_net.training.cristae`. + +Cristae training differs from the other training functions in two ways. +First, the network takes **two input channels**: the tomogram and a semantic mitochondria state, where 0 is background, +1 is a mitochondrion that carries cristae annotations and 2 is a mitochondrion that does not. +Both are expected in a single hdf5 dataset of shape `(2, z, y, x)`. +Second, the voxels of the mitochondria without annotations are **excluded from the loss**, so that the network is not +penalized for predicting cristae where no annotation exists, and the loss is weighted towards the mitochondria +membrane, which improves the detection of cristae junctions. + +```bash +synapse_net.run_cristae_training \ + -n my-cristae-model \ # The name of the model checkpoint. + -i /path/to/tomograms \ # One or more folders with the hdf5 files, searched recursively. + --raw_key raw_mitos_combined \ # The internal path of the tomogram and mitochondria state. + --label_key labels/cristae \ # The internal path of the cristae annotations. + --patch_shape 32 256 256 \ # The patch shape in ZYX. + --batch_size 24 \ # The batch size for training. + --n_iterations 100000 \ # The maximal number of iterations to train for. +``` +The membrane weighting is controlled with `--membrane_w_pos` and `--membrane_w_neg`; pass 1.0 for both to train +without it. The data is split into training and validation data randomly, with a fixed seed; pass `--split_file` to +use an explicit split instead. The default batch size and patch shape require a GPU with a lot of memory. + +Run +```bash +synapse_net.run_cristae_training -h +``` +for more information and instructions on how to use the command. + + +### Volume EM Mitochondria Training + +SynapseNet provides a dedicated function for training a model for mitochondria segmentation in volume electron +microscopy. It reproduces the recipe that was used for our volume EM mitochondria model, which was trained on FIB-SEM +data with a voxel size of 25 nm in z and 5 nm in xy, and is implemented in `synapse_net.training.mitochondria_vol_em`. + +Volume EM training differs from the [tomography training](#mitochondria-training) in two ways. +First, the data is strongly anisotropic (5:1), so the U-Net downsamples only in xy for its first **two** levels instead +of one, and the network is trained **without normalization layers**. +Second, the blocks were cut out of a larger FIB-SEM volume and carry **white filler borders** where the cutout extends +past the imaged region. They take up a large part of some blocks, so they are removed before the normalization, which +would otherwise be skewed by them. + +```bash +synapse_net.run_vol_em_mitochondria_training \ + -n my-vol-em-mito-model \ # The name of the model checkpoint. + -i /path/to/blocks \ # One or more folders with the hdf5 files, searched recursively. + --raw_key raw \ # The internal path of the image data, 'raw' by default. + --label_key labels/mitochondria \ # The internal path of the annotations, 'labels/mitochondria' by default. + --patch_shape 32 512 512 \ # The patch shape in ZYX. + --batch_size 4 \ # The batch size for training. + --n_iterations 50000 \ # The maximal number of iterations to train for. +``` +Pass `--no_white_patch_fix` for data that was not cut out of a larger volume, so that it does not have the filler +borders. The data is split into training and validation data randomly, with a fixed seed; pass `--split_file` to use an +explicit split instead. The default batch size and patch shape require a GPU with a lot of memory; reduce them and pass +`--mixed_precision` to train on a smaller GPU. + +Note that the preprocessing is part of the model and has to be reproduced at inference time in two steps: remove the +filler from the whole volume with `synapse_net.training.transform.remove_white_patches`, then pass +`preprocess=torch_em.transform.raw.normalize_percentile` to `synapse_net.inference.mitochondria.segment_mitochondria`. +The filler removal cannot go into `preprocess`, because the input volume is standardized before that runs and the +filler is identified by its literal value. Leaving the filler in is not a small deviation: on a block that is 14% +filler it cost 0.04 Dice (0.79 instead of 0.84), because the filler skews the percentile normalization of every tile it +overlaps. See `scripts/volume_em/inference/run_mitochondria_vol_em_segmentation.py` for a script that does both. + +Run +```bash +synapse_net.run_vol_em_mitochondria_training -h +``` +for more information and instructions on how to use the command. + + ### Domain Adaptation SynapseNet provides functionality for (unsupervised) domain adaptation. diff --git a/scripts/cooper/training/split-cristae_membw_pos3neg2_2026-07-31.json b/scripts/cooper/training/split-cristae_membw_pos3neg2_2026-07-31.json new file mode 100644 index 00000000..e067712b --- /dev/null +++ b/scripts/cooper/training/split-cristae_membw_pos3neg2_2026-07-31.json @@ -0,0 +1,178 @@ +{ + "roots": { + "cooper_s2": "/scratch-grete/projects/nim00007/data/mitochondria/cooper/raw_mito_combined_s2", + "cooper_cristae": "/mnt/lustre-grete/usr/u12103/mitochondria/cooper/cristae", + "wichmann": "/mnt/lustre-grete/usr/u12103/cristae_data/wichmann", + "wichmann_needs_corrections": "/mnt/lustre-grete/usr/u12103/cristae_data/wichmann_needs_corrections" + }, + "train": [ + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_455L_KO_M.Stim_B2_11_35933_combined.h5", + "wichmann/inital_data/WT22_eb10_model2_combined.h5", + "wichmann/inital_data/WT39_eb2_model_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_438X_KO_M.Stim_M3_7_35935_combined.h5", + "cooper_cristae/for_Luca_2026/36859_J1_STEM750_66K_SP_14_rec_2kb1dawbp_crop_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_455L_KO_M.Stim_B4_9_35933_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_M4_4_35934_combined.h5", + "wichmann/inital_data/M1_eb3_model_combined.h5", + "wichmann/inital_data/WT39_eb10_model_combined.h5", + "wichmann/inital_data/WT20_eb4_model2_combined.h5", + "wichmann/inital_data/WT20_syn6_model2_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_M4_16_35934_combined.h5", + "wichmann/inital_data/WT20_syn7_model2_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_449T_KO_M.Stim_M1_9_35934_combined.h5", + "wichmann/inital_data/WT20_eb9_model2_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429A_WT_M.Stim_D3_2model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_449T_KO_M.Stim_M1_1_35934_combined.h5", + "wichmann/inital_data/WT11_syn6_model2_combined.h5", + "wichmann/inital_data/M7_eb7_model_combined.h5", + "wichmann/inital_data/WT40_eb5_model_combined.h5", + "wichmann/inital_data/WT22_syn7_model2_combined.h5", + "wichmann/inital_data/WT39_eb4_model_combined.h5", + "wichmann/inital_data/WT22_eb6_AZ1_model2_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429D_WT_Rest_H5_3_35461_model_combined.h5", + "cooper_s2/37371_O5_66K_TS_SP_35_rec_2Kb1dawbp_crop3_combined.h5", + "wichmann/inital_data/KO9_eb10_model_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_L3_1_35934_combined.h5", + "wichmann/inital_data/M7_eb2_model_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429D_WT_Rest_H5_4_35461_model_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_M4_7_35934_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B3_6_35933_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_455L_KO_M.Stim_B4_4_35933_combined.h5", + "wichmann/inital_data/WT40_eb8_model_combined.h5", + "wichmann/inital_data/WT13_eb3_model2_combined.h5", + "wichmann/inital_data/WT20_syn3_model2_combined.h5", + "wichmann/inital_data/WT21_eb7_model2_combined.h5", + "wichmann/inital_data/M2_eb10_model_combined.h5", + "wichmann/inital_data/M6_eb2_model_combined.h5", + "wichmann/inital_data/WT11_syn3_model2_combined.h5", + "cooper_cristae/forgotten_cristae_h5_combined/36859_J1_66K_TS_PS_01_rec_2kb1dawbp_crop_combined.h5", + "wichmann/inital_data/WT20_eb7_AZ1_model2_combined.h5", + "wichmann/inital_data/M7_eb10_model_combined.h5", + "wichmann/inital_data/M2_eb1_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B2_2_35933_combined.h5", + "wichmann/inital_data/KO9_eb1_model_combined.h5", + "wichmann/inital_data/WT40_eb9_model_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_455L_KO_M.Stim_B4_6_35933_combined.h5", + "wichmann/inital_data/WT22_syn5_model2_combined.h5", + "wichmann/inital_data/M3_eb4_model_combined.h5", + "wichmann/inital_data/WT21_eb9_model2_combined.h5", + "wichmann/inital_data/WT16_syn2_model2_combined.h5", + "cooper_cristae/2026/36194_B4_66K_TS_SC_07_rec_crop_combined.h5", + "cooper_cristae/forgotten_cristae_h5_combined/36859_J1_66K_TS_CA3_PS_52_rec_2Kb1dawbp_crop_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429D_WT_Rest_G3_4_model_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429D_WT_Rest_H5_5_35461_model_combined.h5", + "wichmann/inital_data/WT11_eb5_model2_combined.h5", + "cooper_s2/36859_J1_STEM750_66K_SP_03_rec_2kb1dawbp_crop_combined.h5", + "wichmann/inital_data/WT20_eb1_AZ2_12K_model2_combined.h5", + "cooper_cristae/forgotten_cristae_h5_combined/36859_J1_66K_TS_CA3_PS_23_rec_2Kb1dawbp_crop_combined.h5", + "cooper_s2/36859_J2_66K_TS_R04_MF04_rec_2Kb1dawbp_crop3_combined.h5", + "wichmann/inital_data/WT20_eb11_model2_combined.h5", + "wichmann/inital_data/WT19_syn3_model2_combined.h5", + "wichmann/inital_data/M1_eb8_model_combined.h5", + "wichmann/inital_data/WT13_syn10_model2_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B2_5_35933_combined.h5", + "wichmann/inital_data/WT21_eb3_model2_combined.h5", + "wichmann/inital_data/WT19_syn6_model2_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN13_480A_KO_M.Stim_B3_13_37819_combined.h5", + "wichmann/inital_data/M1_eb7_model_combined.h5", + "wichmann/inital_data/M8_eb8_model_combined.h5", + "wichmann/inital_data/WT40_eb11_model_combined.h5", + "cooper_cristae/for_Luca_2026/36194_B4_66K_TS_SC_12_rec_crop_combined.h5", + "cooper_cristae/fidi_cristae_h5combined_2026-07/36859_J1_STEM750_66K_SP_03_rec_2kb1dawbp_crop3_manual-corr_combined.h5", + "wichmann/inital_data/KO9_eb4_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B4_3_35933_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_449T_KO_M.Stim_M1_6_35934_combined.h5", + "wichmann/inital_data/WT20_eb8_AZ2_model2_combined.h5", + "wichmann/inital_data/M1_eb5_model_combined.h5", + "wichmann/inital_data/WT40_eb3_model_combined.h5", + "wichmann/inital_data/WT11_eb1_model2_combined.h5", + "wichmann/inital_data/WT21_syn3_model2_combined.h5", + "wichmann/inital_data/WT13_syn6_model2_combined.h5", + "wichmann/inital_data/WT13_syn1_model2_combined.h5", + "wichmann/inital_data/WT20_syn2_model2_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_438X_KO_M.Stim_N5_2_35935_combined.h5", + "wichmann/inital_data/WT13_syn11_model2_combined.h5", + "wichmann/inital_data/WT21_eb4_model2_combined.h5", + "wichmann/inital_data/WT13_syn5_model2_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429D_WT_Rest_H5_1_35461_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN13_480A_KO_M.Stim_B3_7_37819_combined.h5", + "wichmann/inital_data/WT22_syn10_model2_combined.h5", + "wichmann/inital_data/WT22_syn9_model2_combined.h5", + "wichmann/inital_data/M5_eb3_model_combined.h5", + "wichmann/inital_data/WT13_eb5_model2_combined.h5", + "wichmann/inital_data/KO8_eb2_model_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429C_WT_M.Stim_G3_4_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN13_480A_KO_M.Stim_B3_8_37819_combined.h5", + "cooper_cristae/fidi_cristae_h5combined_2026-07/36194_Q4_66K_TS_CA1_PS_43_rec_2Kb1dawbp_crop3_manual-corr_combined.h5", + "wichmann/inital_data/WT21_syn5_model2_combined.h5", + "cooper_cristae/forgotten_cristae_h5_combined/36859_J1_STEM750_66K_SP_03_rec_2kb1dawbp_crop_combined.h5", + "wichmann/inital_data/WT22_eb9_model2_combined.h5", + "wichmann/inital_data/WT19_syn9_model2_combined.h5", + "wichmann/inital_data/WT20_eb2_AZ2_12K_model2_combined.h5", + "cooper_cristae/2026/36859_J1_STEM750_66K_SP_06_rec_2kb1dawbp_crop_combined.h5", + "wichmann/inital_data/WT39_eb7_model_combined.h5", + "wichmann/inital_data/M7_eb9_model_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429C_WT_M.Stim_G3_5_model_combined.h5", + "wichmann/inital_data/WT19_syn1_model2_combined.h5", + "wichmann/inital_data/WT13_syn7_model2_combined.h5", + "cooper_cristae/fidi_cristae_h5combined_2026-07/36859_J1_66K_TS_PS_01_rec_2kb1dawbp_crop3_manual-corr_combined.h5", + "wichmann/inital_data/KO9_eb13_model_combined.h5", + "cooper_cristae/fidi_cristae_h5combined_2026-07/36859_J2_66K_TS_R04_MF04_rec_2Kb1dawbp_crop3_manual-corr_combined.h5", + "cooper_cristae/forgotten_cristae_h5_combined/36859_J2_66K_TS_R04_MF04_rec_2Kb1dawbp_crop_combined.h5", + "wichmann/inital_data/WT20_syn5_model2_combined.h5", + "wichmann/inital_data/WT22_syn6_model2_combined.h5", + "wichmann/inital_data/WT11_eb7_model2_combined.h5", + "cooper_cristae/for_Luca_2026/36859_J2_66K_TS_R04_MF03_rec_2Kb1dawbp_crop_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_449T_KO_M.Stim_M1_7_35934_combined.h5", + "wichmann/inital_data/M8_eb12_model_combined.h5", + "wichmann/inital_data/WT41_eb6_model_combined.h5", + "wichmann/inital_data/M2_eb3_model_combined.h5", + "cooper_cristae/for_Luca_2026/36859_J2_66K_TS_R04_MF05_rec_2Kb1dawbp_crop_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_M4_15_35934_combined.h5", + "wichmann/inital_data/KO9_eb9_model_combined.h5", + "cooper_s2/36859_J1_66K_TS_PS_01_rec_2kb1dawbp_crop_combined.h5", + "wichmann/inital_data/KO9_eb12_model_combined.h5", + "wichmann/inital_data/M7_eb15_model_combined.h5", + "wichmann/inital_data/WT20_eb7_AZ2_model2_combined.h5", + "wichmann/inital_data/WT13_eb4_model2_combined.h5", + "wichmann/inital_data/M2_eb6_model_combined.h5", + "wichmann/inital_data/M7_eb6_model_combined.h5", + "wichmann/inital_data/M7_eb12_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B4_7_35933_combined.h5", + "wichmann/2026-05-26-dataset/2026-05-26_corrected_combined/Otof_AVCN07_449T_KO_M.Stim_L3_9_35934_combined.h5", + "wichmann/inital_data/WT11_syn1_model2_combined.h5" + ], + "val": [ + "wichmann/inital_data/M2_eb2_AZ2_model_combined.h5", + "wichmann/inital_data/M2_eb14_model_combined.h5", + "cooper_cristae/2025/raw_mitos_combined_s2/37371_O4_66K_TS_SC_52_rec_2Kb1dawbp_cropF_s2_combined.h5", + "cooper_cristae/2025/raw_mitos_combined_s2/37371_O5_66K_TS_SP_34-01_rec_2Kb1dawbp_cropF_s2_combined.h5", + "wichmann/inital_data/M10_eb9_model_combined.h5", + "cooper_cristae/for_Luca_2026/36859_J2_66K_TS_R04_PS06_rec_2Kb1dawbp_crop_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_438X_KO_M.Stim_M5_1_35935_combined.h5", + "wichmann/inital_data/WT11_eb2_model2_combined.h5", + "wichmann/inital_data/M2_eb8_model_combined.h5", + "wichmann/inital_data/KO9_eb14_model_combined.h5", + "wichmann/inital_data/WT13_syn4_model2_combined.h5", + "wichmann/inital_data/WT13_syn9_model2_combined.h5", + "cooper_cristae/2026/36194_B4_66K_TS_R01A_SC_15_rec_crop_combined.h5", + "wichmann/inital_data/M1_eb1_model_combined.h5" + ], + "test": [ + "wichmann/inital_data/WT22_eb5_model2_combined.h5", + "wichmann/inital_data/WT40_eb10_model_combined.h5", + "wichmann/inital_data/M5_eb1_model_combined.h5", + "cooper_cristae/2026/36859_J1_66K_TS_PS_03_rec_2kb1dawbp_crop_combined.h5", + "cooper_cristae/2026/36194_B4_66K_TS_SC_22_rec_crop_combined.h5", + "wichmann/inital_data/WT21_eb5_model2_combined.h5", + "wichmann/inital_data/KO8_eb4_model_combined.h5", + "wichmann_needs_corrections/2026-05-26-dataset/Otof_AVCN07_455L_KO_M.Stim_B3_2_35933_combined.h5", + "wichmann/inital_data/M8_eb6_model_combined.h5", + "wichmann/inital_data/WT20_eb5_model2_combined.h5", + "wichmann/inital_data/M1_eb6_model_combined.h5", + "wichmann/inital_data/M2_eb5_model_combined.h5", + "wichmann/inital_data/Otof_AVCN03_429A_WT_M.Stim_D3_4model_combined.h5", + "cooper_cristae/2026/36194_B4_66K_TS_R01A_SC_01_rec_crop_combined.h5", + "wichmann/inital_data/WT21_syn4_model2_combined.h5" + ] +} diff --git a/scripts/cooper/training/split-mito_tomo_s4_refined.json b/scripts/cooper/training/split-mito_tomo_s4_refined.json new file mode 100644 index 00000000..d24aa67e --- /dev/null +++ b/scripts/cooper/training/split-mito_tomo_s4_refined.json @@ -0,0 +1,98 @@ +{ + "train": [ + "fidi_s4/2025/ctrl/37371_O5_66K_TS_SP_34-01_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/ctrl/37371_O5_66K_TS_SP_35_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/ctrl/37371_O5_66K_TS_SP_38-01_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/ctrl/37371_O5_66K_TS_SP_67_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/ctrl/mod_37371_O5_66K_TS_SP_34-02_re_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_50_rec_2Kb1dawbp_crop_F_1_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_50_rec_2Kb1dawbp_crop_F_2_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_50_rec_2Kb1dawbp_crop_F_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_54_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_57_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_59_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/orig/3.2_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_MF_18_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_MF_19_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_MF_20_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_MF_21_rec_2Kb1dawbp_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_23_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_25_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_32_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_45_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_46_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_52_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_PS_01_rec_2kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J1_STEM750_66K_SP_03_rec_2kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/36859_J2_66K_TS_R04_MF04_rec_2Kb1dawbp_crop3_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_35_rec_2Kb1dawbp_crop3_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_35_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_50_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_51_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_52_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_67_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/M13_CTRL_09201_S2_02_DIV31_mtk_02_downscaled_s2.h5", + "fidi_s4/orig/M13_CTRL_09201_S2_06_DIV31_mtk_06_downscaled_s2.h5", + "fidi_s4/orig/WT_Unt_SC_09175_B5_01_DIV16_mtk_01_downscaled_s2.h5", + "fidi_s4/orig/WT_Unt_SC_09175_B5_03_DIV16_mtk_05_downscaled_s2.h5", + "fidi_s4/orig/WT_Unt_SC_09175_E4_04_DIV14_mtk_04_downscaled_s2.h5", + "fidi_s4/orig/mod_37371_O5_66K_TS_SP_34-02_re_2Kb1dawbp_crop_downscaled_s2.h5", + "wichmann_refined_s2/KO8_eb4_model.h5", + "wichmann_refined_s2/KO9_eb10_model.h5", + "wichmann_refined_s2/KO9_eb11_model.h5", + "wichmann_refined_s2/KO9_eb12_model.h5", + "wichmann_refined_s2/KO9_eb14_model.h5", + "wichmann_refined_s2/KO9_eb1_model.h5", + "wichmann_refined_s2/KO9_eb4_model.h5", + "wichmann_refined_s2/KO9_eb9_model.h5", + "wichmann_refined_s2/M10_eb12_model.h5", + "wichmann_refined_s2/M10_eb8_model.h5", + "wichmann_refined_s2/M10_eb9_model.h5", + "wichmann_refined_s2/M1_eb3_model.h5", + "wichmann_refined_s2/M1_eb5_model.h5", + "wichmann_refined_s2/M1_eb6_model.h5", + "wichmann_refined_s2/M1_eb8_model.h5", + "wichmann_refined_s2/M1_eb9_model.h5", + "wichmann_refined_s2/M2_eb14_model.h5", + "wichmann_refined_s2/M2_eb1_model.h5", + "wichmann_refined_s2/M2_eb2_AZ2_model.h5", + "wichmann_refined_s2/M2_eb5_model.h5", + "wichmann_refined_s2/M2_eb6_model.h5", + "wichmann_refined_s2/M2_eb7_model.h5", + "wichmann_refined_s2/M2_eb8_model.h5", + "wichmann_refined_s2/M2_eb9_model.h5", + "wichmann_refined_s2/M3_eb4_model.h5", + "wichmann_refined_s2/M5_eb1_model.h5", + "wichmann_refined_s2/M5_eb3_model.h5", + "wichmann_refined_s2/M7_eb11_model.h5", + "wichmann_refined_s2/M7_eb12_model.h5", + "wichmann_refined_s2/M7_eb15_model.h5", + "wichmann_refined_s2/M7_eb2_model.h5", + "wichmann_refined_s2/M7_eb3_model.h5", + "wichmann_refined_s2/M7_eb4_model.h5", + "wichmann_refined_s2/M7_eb5_model.h5", + "wichmann_refined_s2/WT20_eb9_model2.h5", + "wichmann_refined_s2/WT20_syn1_model2.h5", + "wichmann_refined_s2/WT20_syn5_model2.h5", + "wichmann_refined_s2/WT20_syn7_model2.h5", + "wichmann_refined_s2/WT21_eb3_model2.h5", + "wichmann_refined_s2/WT21_eb7_model2.h5", + "wichmann_refined_s2/WT21_syn3_model2.h5", + "wichmann_refined_s2/WT21_syn4_model2.h5" + ], + "val": [ + "fidi_s4/2025/m13dko/37371_O4_66K_TS_SC_52_rec_2Kb1dawbp_cropF_s2_s2.h5", + "fidi_s4/orig/36859_J1_66K_TS_CA3_PS_42_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/37371_O5_66K_TS_SP_34-01_rec_2Kb1dawbp_crop_downscaled_s2.h5", + "fidi_s4/orig/M13_CTRL_09201_S2_04_DIV31_mtk_04_downscaled_s2.h5", + "wichmann_refined_s2/KO9_eb13_model.h5", + "wichmann_refined_s2/M1_eb1_model.h5", + "wichmann_refined_s2/M2_eb10_model.h5", + "wichmann_refined_s2/M2_eb3_model.h5", + "wichmann_refined_s2/M6_eb2_model.h5", + "wichmann_refined_s2/M7_eb10_model.h5", + "wichmann_refined_s2/WT20_syn3_model2.h5", + "wichmann_refined_s2/WT21_eb5_model2.h5", + "wichmann_refined_s2/WT21_eb9_model2.h5" + ] +} diff --git a/scripts/cooper/training/train_cristae.py b/scripts/cooper/training/train_cristae.py new file mode 100644 index 00000000..dd1ef605 --- /dev/null +++ b/scripts/cooper/training/train_cristae.py @@ -0,0 +1,89 @@ +"""Train the cristae model for electron tomography. + +This reproduces the training of SynapseNet's 'cristae5' model. The training data consists of the +cooper and wichmann tomograms at a resolution of roughly 1.74 nm, with cristae annotations and a +semantic mitochondria state that marks which mitochondria carry those annotations. + +The split file next to this script holds the exact train / val split of the published run. Pass +'--random_split' to split the data randomly instead; the three volumes in EXCLUDE are then left out, +because their raw data or their annotations are unreliable, and so are the volumes that the split +file holds out for testing, which the run would otherwise train on. + +Note that the published checkpoint did not converge: it stopped at iteration 39,406 of 100,000 after +running into the job time limit twice. Rerun this script with '--resume' to continue it; N_ITERATIONS +is the total number of iterations, so it will train the remaining ones. +""" + +import argparse +import os + +from synapse_net.training.cristae import cristae_training, get_cristae_paths, get_cristae_test_paths + +TRAIN_ROOTS = { + "cooper_s2": "/scratch-grete/projects/nim00007/data/mitochondria/cooper/raw_mito_combined_s2", + "cooper_cristae": "/mnt/lustre-grete/usr/u12103/mitochondria/cooper/cristae", + "wichmann": "/mnt/lustre-grete/usr/u12103/cristae_data/wichmann", + "wichmann_needs_corrections": "/mnt/lustre-grete/usr/u12103/cristae_data/wichmann_needs_corrections", +} +OUTPUT_ROOT = "/mnt/lustre-grete/usr/u12103/cristae" +SPLIT_FILE = os.path.join(os.path.dirname(__file__), "split-cristae_membw_pos3neg2_2026-07-31.json") + +EXCLUDE = [ + "Otof_AVCN03_429C_WT_M.Stim_G3_1_model_combined", # The raw data of this volume looks wrong. + "WT20_eb8_AZ1_model_combined", # Poor cristae annotations. + "WT22_eb8_model_combined", # Poor cristae annotations. +] + +PATCH_SHAPE = (32, 256, 256) +BATCH_SIZE = 24 +LEARNING_RATE = 1.7e-4 +N_ITERATIONS = 100000 +EARLY_STOPPING = 25 + +# The membrane proximity weighting. This was validated both on the average precision and directly on +# the detection of cristae junctions, where it raised the fraction of mitochondria with a detected +# junction from 4.2% to 54.2% within 8 nm of the membrane. +MEMBRANE_W_POS = 3.0 +MEMBRANE_W_NEG = 2.0 +MEMBRANE_BAND_NM = 12.0 + + +def main(): + parser = argparse.ArgumentParser(description="Train the cristae model for electron tomography.") + parser.add_argument("-n", "--name", default="cristae-net32-lr1.7e-4-bs24-ps32x256x256-membw-pos3-neg2-band12-2026-07-31", # noqa + help="The name of the model to be trained.") + parser.add_argument("-o", "--output_root", default=OUTPUT_ROOT, help="The folder for the checkpoint and logs.") + parser.add_argument("--random_split", action="store_true", + help="Split the data randomly instead of using the split of the published run.") + parser.add_argument("--resume", action="store_true", + help="Continue the previous run with the same name, restoring its optimizer and " + "iteration count.") + parser.add_argument("--check", action="store_true", help="Check the dataloaders instead of running training.") + args = parser.parse_args() + + exclude = EXCLUDE if not args.random_split else EXCLUDE + get_cristae_test_paths(SPLIT_FILE, TRAIN_ROOTS) + train_paths, val_paths = get_cristae_paths( + TRAIN_ROOTS, split_file=None if args.random_split else SPLIT_FILE, exclude=exclude, + ) + print("Training on", len(train_paths), "tomograms and validating on", len(val_paths), "tomograms.") + + cristae_training( + name=args.name, + train_paths=train_paths, + val_paths=val_paths, + save_root=args.output_root, + patch_shape=PATCH_SHAPE, + batch_size=BATCH_SIZE, + lr=LEARNING_RATE, + n_iterations=N_ITERATIONS, + early_stopping=EARLY_STOPPING, + membrane_w_pos=MEMBRANE_W_POS, + membrane_w_neg=MEMBRANE_W_NEG, + membrane_band_nm=MEMBRANE_BAND_NM, + resume=args.resume, + check=args.check, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/cooper/training/train_mitochondria_tomo.py b/scripts/cooper/training/train_mitochondria_tomo.py new file mode 100644 index 00000000..a6198116 --- /dev/null +++ b/scripts/cooper/training/train_mitochondria_tomo.py @@ -0,0 +1,62 @@ +"""Train the mitochondria model for electron tomography. + +This reproduces the training of SynapseNet's 'mitochondria2' model. The training data consists of +the fidi (downscaled by a factor of 4) and wichmann (downscaled by a factor of 2) tomograms, which +are both at a resolution of roughly 2.87 nm, with refined mitochondria annotations. + +The split file next to this script holds the exact train / val split of the published run. +Delete the '--split_file' argument below to split the data randomly instead. +""" + +import argparse +import os + +from synapse_net.training.mitochondria import get_mitochondria_paths, mitochondria_training + +TRAIN_ROOT = "/mnt/lustre-grete/usr/u12103/mitochondria/mito-tomo-all" +OUTPUT_ROOT = "/mnt/lustre-grete/usr/u12103/mitochondria/tomo" +SPLIT_FILE = os.path.join(os.path.dirname(__file__), "split-mito_tomo_s4_refined.json") + +PATCH_SHAPE = (32, 256, 256) +BATCH_SIZE = 8 +LEARNING_RATE = 1e-4 +N_ITERATIONS = 150000 +EARLY_STOPPING = 20 + + +def main(): + parser = argparse.ArgumentParser(description="Train the mitochondria model for electron tomography.") + parser.add_argument("-n", "--name", default="mitotomo-net32-lr1e-4-bs8-ps32x256x256-s4-refined-final", + help="The name of the model to be trained.") + parser.add_argument("-i", "--train_root", default=TRAIN_ROOT, help="The folder with the training data.") + parser.add_argument("-o", "--output_root", default=OUTPUT_ROOT, help="The folder for the checkpoint and logs.") + parser.add_argument("--random_split", action="store_true", + help="Split the data randomly instead of using the split of the published run.") + parser.add_argument("--resume", action="store_true", + help="Continue the previous run with the same name, restoring its optimizer and " + "iteration count.") + parser.add_argument("--check", action="store_true", help="Check the dataloaders instead of running training.") + args = parser.parse_args() + + train_paths, val_paths = get_mitochondria_paths( + args.train_root, split_file=None if args.random_split else SPLIT_FILE, + ) + print("Training on", len(train_paths), "tomograms and validating on", len(val_paths), "tomograms.") + + mitochondria_training( + name=args.name, + train_paths=train_paths, + val_paths=val_paths, + save_root=args.output_root, + patch_shape=PATCH_SHAPE, + batch_size=BATCH_SIZE, + lr=LEARNING_RATE, + n_iterations=N_ITERATIONS, + early_stopping=EARLY_STOPPING, + resume=args.resume, + check=args.check, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/volume_em/README.md b/scripts/volume_em/README.md new file mode 100644 index 00000000..642ce58a --- /dev/null +++ b/scripts/volume_em/README.md @@ -0,0 +1,368 @@ +# Volume EM mitochondria + +Training, inference and evaluation for the volume EM (FIB-SEM) mitochondria model, at a voxel size of +25 nm in z and 5 nm in xy. The recipe itself lives in `synapse_net.training.mitochondria_vol_em`; the +scripts here are the cluster-specific parts and the reproduction of the published model. + +``` +training/ + train_mitochondria_vol_em.py the published recipe, with its constants + split-mito_vol_em_aniso2lvl_final.json the pinned split of the published run, 14 train / 2 val + compare_training_setup.py checks that the port builds the published training setup + repair_checkpoint.py makes the published checkpoint loadable again + repro_aniso2lvl_seed{42,43,44}.yaml job manifests for the reproduction runs +inference/ + run_mitochondria_vol_em_segmentation.py segmentation with a local checkpoint +evaluation/ + compare_models.py model comparison on the held-out test blocks + voxel_sweep.py block lookup and resampling shared by the baseline sweeps + segment_mitonet.py MitoNet baseline at a chosen voxel size, for compare_models.py -s + mitonet_iso30.yaml job manifest for it at 30 nm isotropic on the test blocks + mitonet_sweep_val.yaml job manifest for the MitoNet sweep on the validation blocks + mitonet_sweep_score.yaml job manifest for scoring it + segment_microsam.py micro-sam baseline at a chosen voxel size, for compare_models.py -s + microsam_sweep_val.yaml job manifest for the micro-sam sweep on the validation blocks + microsam_sweep_score.yaml job manifest for scoring it + microsam_test_selected.yaml job manifest for the test run of the selected setting + score_voxel_size_sweep.py scores a voxel-size sweep of one method, selects per approach +``` + +## The published model + +`/mnt/lustre-grete/usr/u15205/volume-em/models/checkpoints/volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final`, +best validation DiceLoss **0.292714**. + +Two things about it are not visible from its configuration: + +- **It cannot be loaded.** torch-em pickles the training dataset into the checkpoint, and the script + that trained it defined its raw transform as a local function, so the checkpoint references + `__main__._raw_transform_with_fix_white_patches`. Use `training/repair_checkpoint.py` to write a + loadable copy; the weights are copied through bit-identically and the original is not modified. +- **Its counters describe only half of its history.** Two jobs were submitted 24 minutes apart and ran + concurrently for ~23 h against the same checkpoint folder. The second warm-started weights-only from + the first's `best.pt` as it stood at epoch 15, so its optimizer and counters restarted. Both ended on + early stopping, neither on the configured 50,000 iterations nor on the wall clock. The surviving + `best.pt` is the second job's: ~16,500 gradient steps in total, with an optimizer reset at 2,000. + +The recipe as documented is a single clean run, and that is what this module trains. Running twice +with the same name and `save_root` is now refused rather than silently racing. + +## Reproduction + +Three clean runs, seeds 42/43/44, into a separate `save_root`: + +```bash +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ + scripts/volume_em/training/repro_aniso2lvl_seed42.yaml # and 43, 44 +``` + +Progress goes to **stderr**, so watch the `.err` file that `sbatch_runner.py` writes, not the `.log`: +the `.log` only holds the four startup lines, and looks stalled if you watch it instead. + +```bash +tr '\r' '\n' < logs/sbatch_jobs/_.err | tail -3 +``` + +`synapse_net` is installed editable, so a queued job picks up whatever is in the working tree at the +moment it *starts*, not at the moment it was submitted. Do not switch branches in this checkout while +runs are queued or running, or they will train something other than what you submitted. `git log -1` +of the checkout at job start is worth recording alongside the results. + +Before spending the GPU time, the cheap check is the one that actually verifies the port: + +```bash +python scripts/volume_em/training/compare_training_setup.py +``` + +It builds the trainer through the ported recipe, stops before training, and diffs it against the +setup stored in the published checkpoint — model kwargs, optimizer, scheduler, loss, both loader +kwargs, the sampler parameters, the per-file dataset lengths and the file lists. It also checks the +two places where the port deviates structurally: the raw transform is compared numerically against a +transcription of the published one, and the `PadIfNecessary` that the port adds to the joint transform +is confirmed never to fire. + +### What to expect + +- **Stopping.** `early_stopping=20` is the operative criterion, not `n_iterations`. Every historical + run stopped after 16k–19k iterations / 26–29 h with 8–10 h of its limit unused. A forced 50,000 + iterations would take ~71 h at the measured 5.13 s/it, which exceeds the 48 h partition maximum. +- **Run-to-run spread.** Three comparable trainings on this split reached 0.292714, 0.296378 and + 0.298842, so a spread of about 0.006 is normal and a reproduction inside that band is + indistinguishable from noise. +- **torch-em version.** The reproduction runs against the torch-em checkout that produced the + published model (`f247988b`, 0.8.3), so the library is not a variable. The setup check also passes + under a clean environment built from `environment.yaml` (Python 3.14, torch 2.13, torch_em 0.10.6): + every recipe field is identical there too, and the only difference is a `mixed_precision_dtype` + field that newer torch-em adds and that this recipe never reads, since it trains in single + precision. The check reports such additions explicitly rather than ignoring unknown fields. +- **Seeding.** `--seed` makes runs repeatable in every way that matters — two runs with the same seed + agree to ~1e-7 in the validation metric — but not bit-identically, because cudnn picks kernels + adaptively. `--deterministic` does *not* fix that here (torch-em sets `warn_only=True`, so + non-deterministic kernels silently fall back); it only costs throughput, so it is not used. + +## Comparison + +```bash +python scripts/volume_em/evaluation/compare_models.py \ + -m paper=/mnt/lustre-grete/usr/u15205/volume-em/models-repro/reference/volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final \ + -m seed42=/mnt/lustre-grete/usr/u15205/volume-em/models-repro/checkpoints/volume-em-mito-aniso2lvl-repro-seed42 \ + -m seed43=... -m seed44=... \ + -o scripts/volume_em/RESULTS.md +``` + +Every model is segmented with identical preprocessing and identical post-processing; tuning the +post-processing per model would fold a post-processing difference into what is meant to be a +comparison of the trained networks. Segmentations are cached, so a model can be added later without +redoing the others. + +### Post-processing, and why `min_size` is not tuned here + +The published post-processing parameters come from a grid search against a different, out-of-core +watershed, and under this one they over-segment: on the test blocks the model predicts 48 and 57 +instances against 30 and 43 annotated ones, so precision (0.585) is far below recall (0.847) while the +semantic dice is a healthy 0.849. The extra objects are small fragments. + +Raising `min_size` from 1,000 to 20,000 looks like a free fix — F1 goes from 0.69 to 0.85, the counts +become exactly 30/30 and 43/43, and recall does not move at all, so only false positives are removed. +**It is not a fix, and the default is deliberately left alone.** That threshold was read off the two +test blocks, whose mitochondria happen to be large (5th percentile 21,100 and 35,263 voxels). In the +14 training blocks, of 717 annotated mitochondria: + +| below | objects | share | +|---|---|---| +| 1,000 voxels | 2 | 0.3 % | +| 5,000 voxels | 18 | 2.5 % | +| 10,000 voxels | 45 | 6.3 % | +| 20,000 voxels | 129 | **18.0 %** | + +So a 20,000-voxel filter would throw away nearly a fifth of the objects the annotators marked. It +scores well on these two blocks and would generalize badly. The inherited `min_size=1000` is instead +consistent with the annotation: 99.7 % of annotated mitochondria are larger than it. + +The real lever is the seed and boundary parameters, which is where the fragmentation comes from, and +re-tuning those needs a validation set that is not these two blocks. `compare_models.py +--size_filter_sweep` reproduces the table above as a diagnostic; read it as a measure of how much of +the error is fragments, not as a tuning result. + +### MitoNet baseline + +MitoNet (empanada, `MitoNet_v1`) is a 2D generalist whose 3D mode, the ortho-plane consensus, infers on +xy, xz and yz planes and so assumes isotropic voxels. `evaluation/segment_mitonet.py` resamples each +test block from 25/5/5 nm to 30 nm isotropic (107 x 267 x 267; linear with anti-aliasing), runs +ortho-plane inference, and resizes the result back to the native grid with nearest-neighbor +interpolation. The ground truth is never resampled. Both the consensus and the xy stack alone come +out of the same run: + +```bash +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/mitonet_iso30.yaml +python scripts/volume_em/evaluation/compare_models.py \ + -m paper=/mnt/lustre-grete/usr/u15205/volume-em/models-repro/reference/volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final \ + -s mitonet-iso30-ortho=/mnt/lustre-grete/usr/u15205/volume-em/repro-comparison/mitonet-iso30-ortho \ + -s mitonet-iso30-xy=/mnt/lustre-grete/usr/u15205/volume-em/repro-comparison/mitonet-iso30-xy \ + --size_filter_sweep 1000 10000 20000 30000 \ + --title "Volume EM mitochondria: MitoNet at 30 nm isotropic vs. the published model" \ + -o scripts/volume_em/RESULTS_mitonet_iso30.md +``` + +It runs in its own environment, `/mnt/lustre-grete/usr/u15205/envs/empanada`, built from the synapse +repo's `env_empanada.yaml` (empanada-napari 1.2.1, torch 2.8). + +- **Size filter.** `min_size` is 23 voxels at 30 nm, the same physical volume as synapse-net's 1,000 + native voxels (a 30 nm voxel is 43.2 native ones). empanada's default of 500 would be ~21,600 + native voxels, the kind of threshold the section above explains does not generalize, and it would + flatter MitoNet on these two blocks. `min_extent` is empanada's 4. +- **Everything else is the `Engine3d` default**, as in the earlier MitoNet runs (`confidence_thr` 0.3, + median filter over 5 slices). The napari widget's defaults differ (0.5, 3 slices, `min_extent` 5). +- **The raw is fed as-is.** There is no white-filler removal, which is part of synapse-net's + preprocessing, not MitoNet's. The 4009 block has 11.4 % filler. +- **Result: MitoNet fails at 30 nm**, with F1 0 on both blocks in both modes + (`RESULTS_mitonet_iso30.md`). On 4009 it predicts no objects at all. On 4007 its 16 (consensus) + or 20 (xy) objects barely touch the mitochondria (semantic dice 0.01), and they are small. + - **It is not a resampling artefact.** MitoNet's 2D engine run directly on the resampled slices + also has a dice of 0, and no flip or transpose of the result improves the overlap. + - **The cause is the scale, not the contrast.** The objects MitoNet predicts at 30 nm are dark + (mean raw 97.5), while the mitochondria in these blocks are brighter than their surroundings + (144 vs. 134 in 4007). That first suggested an inverted contrast. The voxel-size sweep below + refutes it: the inverted contrast scores F1 ≈ 0 at every voxel size, while the raw contrast + works once xy is at 8–15 nm. +- **The February 2026 MitoNet numbers are not comparable.** They are in + `/mnt/lustre-grete/usr/u12103/mitochondria/volume-em/test_split_empanada_vs25-*` (F1 ≤ 0.07), and + they differ in three ways: + - they ran on anisotropic input, in the xy stack only; + - they saved the raw panoptic stack as uint8, whose `1000 + n` IDs wrap (ID 1024 became background); + - they were scored with border filtering (`eval_mitos_touching_borders.py -b 2 -z`) against + downsampled labels. + +### MitoNet voxel-size sweep + +30 nm may simply be the wrong scale: a median mitochondrion here is 305–370 nm across (measured on +the validation blocks), so only ~10–12 px at 30 nm. The sweep looks for the voxel size, and the +contrast polarity, at which MitoNet works best. + +**It is run on the two validation blocks of the pinned split, not on the test blocks.** MitoNet never +saw them, and they mirror the test pair: one 4007 and one 4009 block, `(128, 1600, 1600)`, 37 and 28 +mitochondria, the 4009 one with 13.8 % filler. Only the setting selected for each approach is then +run once on the test split. + +Every setting runs on the raw and on the inverted contrast (`255 - raw`): + +| group | voxel size (z/y/x nm) | modes | +|---|---|---| +| A: z native, xy swept | 25 / {5, 7.5, 8, 10, 12.5, 15, 20, 25, 30} | xy stack | +| B: isotropic | {10, 12.5, 15, 20, 25, 30} in every axis | ortho-plane and xy stack | +| C: single settings | 30 / 8 / 8, which MitoNet_v1 was probably trained at | xy stack | + +Why these bounds: +- **A below 5 nm** would upsample the native data. +- **B below 10 nm** would interpolate z by more than 2.5× and reach 0.5–1.6 G voxels. +- **Above 30 nm** the smallest mitochondria fall under ~8 px. + +A setting at the edge of its grid is flagged by the scorer, as a cue to extend the grid by one step. +The sweep holds fixed: +- `min_size`, scaled to the same physical volume as synapse-net's 1,000 native voxels at every + voxel size; +- `min_extent` at 4 voxels; +- the `Engine3d` and consensus defaults. + +25 nm isotropic is shared by A and B. The xy stack comes out identical whether it runs alone or +next to the ortho-plane consensus, which was checked. + +```bash +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/mitonet_sweep_val.yaml +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/mitonet_sweep_score.yaml \ + --dependency afterok: +``` + +The scorer selects per group by the F1 averaged over the validation blocks, ties broken by msa. +It says plainly when even the best setting is below F1 0.05, rather than presenting a failure as a +selection. It writes `RESULTS_mitonet_sweep_val.md` and `best.json` in the sweep root, whose +`segment_args` are the arguments for the one test run. `compare_models.py --split_file +... --split_key val` scores models on the validation blocks the same way; their cached +segmentations go into a `val/` folder so they cannot be confused with the test ones. + +**Result.** The val tables are in `RESULTS_mitonet_sweep_val.md`, and the test run of the three +selected settings is in `RESULTS_mitonet_voxel_size.md`. + +- **Contrast.** The inverted contrast scores F1 ≈ 0 at every voxel size, so the raw contrast is + the right one. +- **Scale.** With the raw contrast, MitoNet works only at fine scales. It peaks at 8–15 nm in xy + (median mitochondrion 20–45 px across) and collapses to 0 at 25–30 nm, in every group. +- **Noise.** The top settings are within a few hundredths of each other, which is within noise + for 65 objects. + +| selected on val | val F1 | test F1 | test precision / recall | test dice | +|---|---|---|---|---| +| A: z 25, xy 10 nm, xy stack | 0.182 | 0.094 | 0.070 / 0.143 | 0.368 | +| B: 12.5 nm isotropic, ortho-plane | 0.235 | 0.101 | 0.083 / 0.131 | 0.328 | +| C: z 30, xy 8 nm, xy stack | 0.223 | 0.121 | 0.087 / 0.200 | 0.407 | +| synapse-net, published model | – | 0.691 | 0.585 / 0.847 | 0.849 | + +At its best voxel size MitoNet goes from F1 0 to about 0.1 on the test blocks, still far below +synapse-net. It predicts 55–99 objects against 30 and 43, most of which do not match. The test +scores are about half the val scores, so the val selection should be read as "8–15 nm" rather +than as one exact voxel size. + +### micro-sam voxel-size sweep + +`evaluation/segment_microsam.py` runs the same grid with micro-sam's automatic instance segmentation +(`vit_b_em_organelles`, AIS decoder), selected on the same validation blocks by the same rule, with +the same physically matched size filter. It runs on the raw contrast only. It uses its own +environment, `/mnt/lustre-grete/usr/u15205/envs/micro-sam`, built from the micro-sam checkout's +`environment.yaml` plus an editable install of the checkout: micro-sam 1.8.14 (`de4231f`), +python-elf 0.9.2, bioimage-cpp 0.9.0. + +- **The micro-sam version matters.** micro-sam 1.7.7 builds nifty graphs for the merge along z, + while elf 0.9 takes bioimage-cpp graphs. Its `environment.yaml` only asks for `python-elf >=0.7.1`, + so a fresh environment for 1.7.7 gets elf 0.9.2 and fails in the merge (`'UndirectedGraph' object + has no attribute 'number_of_nodes'`). micro-sam 1.8.x requires elf 0.9 and works. +- The model is unchanged: the cached `vit_b_em_organelles` weights and decoder match the registry + of 1.8.14. + +- **The encoder scale is controlled, not only the resampling.** + - SAM resizes every image, and every tile, so that its longest side is 1024 px, up or down. + Passed as-is, a slice of these 8 µm blocks would reach the encoder at ~7.8 nm/px at every + voxel size, only more or less blurred. + - So every image or tile that holds data is made exactly 1024 px. A slice of up to 1024 px is + reflect-padded to 1024 × 1024. The two larger ones, xy 5 and 7.5 nm, are padded onto a canvas + on which micro-sam's tiling (768 + a halo of 128) gives every tile holding data its full 1024 px. + That was checked for every grid size. + - The padding is cropped away before scoring. +- **One mode.** micro-sam segments the xy slices and merges them along z; it has no ortho-plane + mode. +- **Settings.** Everything else is micro-sam's own: the AIS defaults, and `gap_closing` 2 and + `min_z_extent` 2, the napari annotator's defaults for automatic 3D segmentation. +- **The February 2026 `test_split_microsam_vs25-*` runs are not comparable.** They were never + scored, and they are stored at the reduced grids. Because of the rescaling above, their "10" and + "20 nm" runs saw ~7.8 nm/px. + +```bash +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/microsam_sweep_val.yaml +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/microsam_sweep_score.yaml \ + --dependency afterok: +python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py scripts/volume_em/evaluation/microsam_test_selected.yaml +``` + +**Result on the validation blocks** (`RESULTS_microsam_sweep_val.md`): + +- **Finer is better, down to the native resolution.** In group A the F1 rises steadily as xy gets + finer: 0 at 12.5–30 nm, 0.03 at 10 nm, 0.04 at 7.5–8 nm, and **0.089 at the native 5 nm**. + - The first run selected 5 nm at the edge of its grid, so the grid was extended by a step. + At 3.75 nm, which upsamples the native data, the F1 falls back to 0.057, so 5 nm is a real peak. + - At 5 nm there is no resampling at all, only the tiled 1024 px canvas. +- **B and C have no working setting.** The best isotropic setting is 10 nm at F1 0.017, and 30/8/8 + reaches 0.029, both below the floor of 0.05, so neither group has a selection. + - B is also rising towards its finest setting. It could only be extended below 10 nm by + interpolating z by more than 2.5×, and A shows that even the native data gives 0.089. +- **micro-sam over-segments at every scale that finds anything.** At 5 nm it predicts 346 objects + against 65 (precision 0.05, recall 0.29). Its semantic dice of 0.46 shows it finds the + mitochondria, but splits them into fragments that fail the IoU 0.5 match. +- **Compared with MitoNet** on the same blocks and grid, micro-sam is worse everywhere: MitoNet's + best settings reach F1 0.18–0.23. + +**Result on the test blocks** (`RESULTS_baselines_voxel_size.md`). The selected setting was run +once, and is scored next to the MitoNet settings selected the same way: + +| setting (selected on val) | val F1 | test F1 | test precision / recall | test dice | objects 4007 / 4009 | +|---|---|---|---|---|---| +| micro-sam A: z 25, xy 5 nm (native) | 0.089 | 0.073 | 0.046 / 0.185 | 0.379 | 132 / 170 | +| MitoNet A: z 25, xy 10 nm | 0.182 | 0.094 | 0.070 / 0.143 | 0.368 | 66 / 85 | +| MitoNet B: 12.5 nm isotropic, ortho-plane | 0.235 | 0.101 | 0.083 / 0.131 | 0.328 | 55 / 63 | +| MitoNet C: z 30, xy 8 nm | 0.223 | 0.121 | 0.087 / 0.200 | 0.407 | 69 / 99 | +| synapse-net, published model | – | 0.691 | 0.585 / 0.847 | 0.849 | 48 / 57 | + +The ground truth has 30 / 43 objects. + +- **Both generalists fail on this data at every voxel size.** Each detects part of the mitochondria + (dice 0.33–0.41) but splits it into fragments, which is why precision is so low. +- **micro-sam is worse than MitoNet**, with the gap largest on 4007 (F1 0.025 against 0.06–0.10). +- **The ranking holds on test.** It matches the one on val, which suggests the val selection + carried over. + +### Test data + +The only held-out ground truth is two blocks under +`/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split/`: + +| block | shape | instances | filler | +|---|---|---|---| +| `4007/4007_block_z001024_001152_y001936_003872_x003312_004968.h5` | (128, 1600, 1600) | 30 | none | +| `4009/4009_raw_z128-256_y1600-3200_x0-1600_6.h5` | (128, 1600, 1600) | 43 | 11.4 % | + +Neither overlaps the training data: the 4007 block continues the training z-stack where it ends, the +4009 block is the one grid tile that training does not use, and the two training cutouts whose +filenames carry no coordinates were checked by content (best normalized cross-correlation 0.33, +against 0.79 for merely adjacent slices and 1.0 for a true match). The blocks differ in whether they +contain filler at all, which makes the pair a useful control, so report them separately as well as +averaged. + +The `test_split_s1*`, `_s2_*` and `_s3_*` folders are the *same two blocks* downsampled, not extra +test volume, and most of them store `raw` as float64, which the filler removal correctly rejects. + +## Known issue, not addressed here + +Every inference and evaluation config in the `synapse` repo points at +`volume-em-mito-net32-lr1e-4-bs4-ps32x512x512-final`, a different architecture (one anisotropic level +rather than two). On 2026-05-28 a job auto-warm-started from it, ran 125 steps at a reset learning +rate of 1e-4 and was killed; its validation metric went 0.320197 → 0.394636 and the good February +weights are gone. Numbers produced through that path after 2026-05-28 12:37 come from a degraded +model. The recipe here does not expose `--scale_factors`, so it cannot currently retrain that variant. diff --git a/scripts/volume_em/RESULTS.md b/scripts/volume_em/RESULTS.md new file mode 100644 index 00000000..1ae38b59 --- /dev/null +++ b/scripts/volume_em/RESULTS.md @@ -0,0 +1,48 @@ +# Volume EM mitochondria: reproduction vs. the published model + +**Status: baseline only.** Only one model is in this report, so there is nothing to compare +against yet. Re-run with the reproduction runs added (`-m seed42=... -m seed43=...`) once they +finish; the segmentations already here are cached and are not recomputed. + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +`dice` is the semantic foreground dice. `val_metric` is the best validation DiceLoss of the run, +read from its checkpoint, where lower is better. + +All models were segmented with identical preprocessing and identical post-processing +(seed_distance 1, boundary_threshold 0.12, area_threshold 200, min_size 1000). + +## Averaged over the test blocks + +| model | f1 | precision | recall | msa | sbd | dice | val_metric | iteration | +|-------|--------|-----------|--------|--------|--------|--------|------------|-----------| +| paper | 0.6905 | 0.5850 | 0.8469 | 0.3197 | 0.5722 | 0.8490 | 0.2927 | 14500 | + +## Per test block + +| model | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|-------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| paper | 4007 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 0.5188 | 0.8742 | 48 | 30 | +| paper | 4009 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 0.6256 | 0.8237 | 57 | 43 | + +## Size filter sweep + +Applied to the finished segmentations, which is not the same as running the watershed with +that `min_size`, so read it as a diagnostic for how much of the error is sub-threshold +fragments. + +**Do not read a `min_size` off this table.** It would be tuning on the test data, and it +would generalize badly: these two test blocks happen to hold only large mitochondria, while +18% of the 717 mitochondria annotated in the training blocks are smaller than 20,000 voxels. +The default of 1,000 is the value consistent with the annotation, which has 99.7% of its +objects above it. + +| model | dataset | min_size | f1 | precision | recall | msa | n_pred | n_true | +|-------|---------|----------|--------|-----------|--------|--------|--------|--------| +| paper | 4007 | 1000 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 48 | 30 | +| paper | 4007 | 10000 | 0.8065 | 0.7812 | 0.8333 | 0.4515 | 32 | 30 | +| paper | 4007 | 20000 | 0.8333 | 0.8333 | 0.8333 | 0.4761 | 30 | 30 | +| paper | 4007 | 30000 | 0.8136 | 0.8276 | 0.8000 | 0.4573 | 29 | 30 | +| paper | 4009 | 1000 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 57 | 43 | +| paper | 4009 | 10000 | 0.8132 | 0.7708 | 0.8605 | 0.3697 | 48 | 43 | +| paper | 4009 | 20000 | 0.8605 | 0.8605 | 0.8605 | 0.4048 | 43 | 43 | +| paper | 4009 | 30000 | 0.8571 | 0.8780 | 0.8372 | 0.4134 | 41 | 43 | diff --git a/scripts/volume_em/RESULTS_baselines_voxel_size.md b/scripts/volume_em/RESULTS_baselines_voxel_size.md new file mode 100644 index 00000000..3a371e57 --- /dev/null +++ b/scripts/volume_em/RESULTS_baselines_voxel_size.md @@ -0,0 +1,41 @@ +# Volume EM mitochondria: the baselines at the voxel sizes selected on the validation blocks + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +`dice` is the semantic foreground dice. `val_metric` is the best validation DiceLoss of the run, +read from its checkpoint, where lower is better. + +All models were segmented with identical preprocessing and identical post-processing +(seed_distance 1, boundary_threshold 0.12, area_threshold 200, min_size 1000). + +Segmentations passed with `-s` were produced elsewhere and are only scored here, so the +sentence above does not apply to them. They were made with: + +- `microsam-A-z25-xy5-xy`: ais_min_size 0, blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, boundary_distance_threshold 0.5, canvas_shape 128/3200/3200, center_distance_threshold 0.5, distance_smoothing 1.6, foreground_smoothing 1.0, foreground_threshold 0.5, gap_closing 2, halo 128/128, inference_voxel_size_nm 25.0/5.0/5.0, invert False, micro_sam 1.8.14, min_size 1000, min_z_extent 2, mode xy slices merged in z, model vit_b_em_organelles, n_lost_to_native 0, sam_input 1024, segmentation_mode ais, target_voxel_size_nm 25.0/5.0/5.0, tile_shape 768/768 +- `mitonet-A-z25-xy10-xy`: blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 25.0/10.0/10.0, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 250, mode xy stack, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, target_voxel_size_nm 25.0/10.0/10.0 +- `mitonet-B-iso12.5-ortho`: allow_one_view False, blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, cluster_iou_thr 0.75, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 12.5/12.5/12.5, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 320, mode ortho-plane consensus of xy, xz, yz, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, pixel_vote_thr 2, target_voxel_size_nm 12.5/12.5/12.5 +- `mitonet-C-z30-xy8-xy`: blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 29.91/8.0/8.0, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 326, mode xy stack, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, target_voxel_size_nm 30.0/8.0/8.0 + +## Averaged over the test blocks + +| model | f1 | precision | recall | msa | sbd | dice | val_metric | iteration | +|-------------------------|--------|-----------|--------|--------|--------|--------|------------|-----------| +| microsam-A-z25-xy5-xy | 0.0734 | 0.0458 | 0.1845 | 0.0152 | 0.1117 | 0.3786 | - | - | +| mitonet-A-z25-xy10-xy | 0.0938 | 0.0698 | 0.1430 | 0.0263 | 0.1665 | 0.3675 | - | - | +| mitonet-B-iso12.5-ortho | 0.1013 | 0.0828 | 0.1314 | 0.0186 | 0.1831 | 0.3281 | - | - | +| mitonet-C-z30-xy8-xy | 0.1209 | 0.0867 | 0.1996 | 0.0267 | 0.1718 | 0.4073 | - | - | +| paper | 0.6905 | 0.5850 | 0.8469 | 0.3197 | 0.5722 | 0.8490 | 0.2927 | 14500 | + +## Per test block + +| model | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|-------------------------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| microsam-A-z25-xy5-xy | 4007 | 0.0247 | 0.0152 | 0.0667 | 0.0037 | 0.0657 | 0.2162 | 132 | 30 | +| mitonet-A-z25-xy10-xy | 4007 | 0.0625 | 0.0455 | 0.1000 | 0.0172 | 0.1413 | 0.3691 | 66 | 30 | +| mitonet-B-iso12.5-ortho | 4007 | 0.0706 | 0.0545 | 0.1000 | 0.0134 | 0.1519 | 0.3225 | 55 | 30 | +| mitonet-C-z30-xy8-xy | 4007 | 0.1010 | 0.0725 | 0.1667 | 0.0231 | 0.1587 | 0.4163 | 69 | 30 | +| paper | 4007 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 0.5188 | 0.8742 | 48 | 30 | +| microsam-A-z25-xy5-xy | 4009 | 0.1221 | 0.0765 | 0.3023 | 0.0266 | 0.1577 | 0.5410 | 170 | 43 | +| mitonet-A-z25-xy10-xy | 4009 | 0.1250 | 0.0941 | 0.1860 | 0.0355 | 0.1918 | 0.3658 | 85 | 43 | +| mitonet-B-iso12.5-ortho | 4009 | 0.1321 | 0.1111 | 0.1628 | 0.0239 | 0.2143 | 0.3338 | 63 | 43 | +| mitonet-C-z30-xy8-xy | 4009 | 0.1408 | 0.1010 | 0.2326 | 0.0304 | 0.1848 | 0.3982 | 99 | 43 | +| paper | 4009 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 0.6256 | 0.8237 | 57 | 43 | diff --git a/scripts/volume_em/RESULTS_microsam_sweep_val.md b/scripts/volume_em/RESULTS_microsam_sweep_val.md new file mode 100644 index 00000000..baf3e0f3 --- /dev/null +++ b/scripts/volume_em/RESULTS_microsam_sweep_val.md @@ -0,0 +1,87 @@ +# Volume EM mitochondria: vit_b_em_organelles voxel-size sweep on the validation blocks + +vit_b_em_organelles on the `val` blocks of `split-mito_vol_em_aniso2lvl_final.json` (4007, 4009), which it never saw. +Each block is resampled to the voxel size, segmented, and resized back to the native +25/5/5 nm grid, where it is scored against the untouched ground truth. `min_size` is scaled +to the same physical volume as synapse-net's 1,000 native voxels at every voxel size. + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +Metrics are averaged over the blocks, `n_pred` and `n_true` are summed. + +The best setting of each approach is selected here, by F1 with ties broken by msa, and then run +once on the test split. Nothing is selected on the test blocks. + +## A: z at 25 nm, xy resampled (xy stack) + +**Selected: `microsam-z25-xy5-xy`** — raw contrast, xy mode, F1 0.0891, msa 0.0241. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 3.75 | raw | xy | 0.0572 | 0.0331 | 0.2099 | 0.0124 | 0.0632 | 0.3038 | 426 | 65 | +| 5 | raw | xy | 0.0891 | 0.0528 | 0.2867 | 0.0241 | 0.0946 | 0.4558 | 346 | 65 | +| 7.5 | raw | xy | 0.0425 | 0.0262 | 0.1120 | 0.0075 | 0.0765 | 0.3136 | 278 | 65 | +| 8 | raw | xy | 0.0394 | 0.0253 | 0.0898 | 0.0081 | 0.0651 | 0.2359 | 231 | 65 | +| 10 | raw | xy | 0.0317 | 0.0234 | 0.0492 | 0.0022 | 0.0380 | 0.1185 | 146 | 65 | +| 12.5 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0218 | 0.0469 | 92 | 65 | +| 15 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0054 | 0.0082 | 64 | 65 | +| 20 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0022 | 0.0025 | 47 | 65 | +| 25 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0048 | 0.0065 | 35 | 65 | +| 30 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0042 | 0.0027 | 33 | 65 | + +## B: isotropic + +**No setting works:** the best is `microsam-iso10-xy` at F1 0.0172, below 0.05, so none is selected. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 10 | raw | xy | 0.0172 | 0.0106 | 0.0449 | 0.0009 | 0.0231 | 0.1229 | 275 | 65 | +| 12.5 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0158 | 0.0537 | 152 | 65 | +| 15 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0081 | 0.0144 | 94 | 65 | +| 20 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0025 | 0.0023 | 53 | 65 | +| 25 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0048 | 0.0065 | 35 | 65 | +| 30 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0008 | 0.0017 | 28 | 65 | + +## C: single settings (voxel size z/y/x in nm) + +**No setting works:** the best is `microsam-z30-xy8-xy` at F1 0.0285, below 0.05, so none is selected. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 30/8/8 | raw | xy | 0.0285 | 0.0184 | 0.0627 | 0.0065 | 0.0681 | 0.2468 | 218 | 65 | + +## Per block + +| setting | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|------------------------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| microsam-iso10-xy | 4007 | 0.0205 | 0.0127 | 0.0541 | 0.0010 | 0.0192 | 0.0784 | 158 | 37 | +| microsam-iso10-xy | 4009 | 0.0138 | 0.0085 | 0.0357 | 0.0007 | 0.0269 | 0.1674 | 117 | 28 | +| microsam-iso12.5-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0127 | 0.0402 | 87 | 37 | +| microsam-iso12.5-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0190 | 0.0671 | 65 | 28 | +| microsam-iso15-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0067 | 0.0098 | 67 | 37 | +| microsam-iso15-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0095 | 0.0190 | 27 | 28 | +| microsam-iso20-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0049 | 0.0046 | 36 | 37 | +| microsam-iso20-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 17 | 28 | +| microsam-iso25-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0046 | 0.0026 | 23 | 37 | +| microsam-iso25-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0051 | 0.0103 | 12 | 28 | +| microsam-iso30-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0016 | 0.0033 | 24 | 37 | +| microsam-iso30-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 4 | 28 | +| microsam-z25-xy10-xy | 4007 | 0.0157 | 0.0111 | 0.0270 | 0.0008 | 0.0252 | 0.0640 | 90 | 37 | +| microsam-z25-xy10-xy | 4009 | 0.0476 | 0.0357 | 0.0714 | 0.0036 | 0.0509 | 0.1729 | 56 | 28 | +| microsam-z25-xy12.5-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0114 | 0.0245 | 58 | 37 | +| microsam-z25-xy12.5-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0322 | 0.0692 | 34 | 28 | +| microsam-z25-xy15-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0050 | 0.0055 | 47 | 37 | +| microsam-z25-xy15-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0059 | 0.0109 | 17 | 28 | +| microsam-z25-xy20-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0044 | 0.0050 | 36 | 37 | +| microsam-z25-xy20-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 11 | 28 | +| microsam-z25-xy3.75-xy | 4007 | 0.0070 | 0.0040 | 0.0270 | 0.0025 | 0.0356 | 0.1448 | 249 | 37 | +| microsam-z25-xy3.75-xy | 4009 | 0.1073 | 0.0621 | 0.3929 | 0.0223 | 0.0909 | 0.4628 | 177 | 28 | +| microsam-z25-xy30-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0084 | 0.0055 | 29 | 37 | +| microsam-z25-xy30-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 4 | 28 | +| microsam-z25-xy5-xy | 4007 | 0.0717 | 0.0430 | 0.2162 | 0.0175 | 0.0799 | 0.3949 | 186 | 37 | +| microsam-z25-xy5-xy | 4009 | 0.1064 | 0.0625 | 0.3571 | 0.0307 | 0.1093 | 0.5167 | 160 | 28 | +| microsam-z25-xy7.5-xy | 4007 | 0.0306 | 0.0189 | 0.0811 | 0.0067 | 0.0573 | 0.2431 | 159 | 37 | +| microsam-z25-xy7.5-xy | 4009 | 0.0544 | 0.0336 | 0.1429 | 0.0083 | 0.0957 | 0.3842 | 119 | 28 | +| microsam-z25-xy8-xy | 4007 | 0.0455 | 0.0288 | 0.1081 | 0.0087 | 0.0569 | 0.1961 | 139 | 37 | +| microsam-z25-xy8-xy | 4009 | 0.0333 | 0.0217 | 0.0714 | 0.0076 | 0.0732 | 0.2756 | 92 | 28 | +| microsam-z30-xy8-xy | 4007 | 0.0261 | 0.0172 | 0.0541 | 0.0059 | 0.0677 | 0.2138 | 116 | 37 | +| microsam-z30-xy8-xy | 4009 | 0.0308 | 0.0196 | 0.0714 | 0.0070 | 0.0686 | 0.2798 | 102 | 28 | diff --git a/scripts/volume_em/RESULTS_mitonet_iso30.md b/scripts/volume_em/RESULTS_mitonet_iso30.md new file mode 100644 index 00000000..8e39ccfe --- /dev/null +++ b/scripts/volume_em/RESULTS_mitonet_iso30.md @@ -0,0 +1,72 @@ +# Volume EM mitochondria: MitoNet at 30 nm isotropic vs. the published model + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +`dice` is the semantic foreground dice. `val_metric` is the best validation DiceLoss of the run, +read from its checkpoint, where lower is better. + +All models were segmented with identical preprocessing and identical post-processing +(seed_distance 1, boundary_threshold 0.12, area_threshold 200, min_size 1000). + +Segmentations passed with `-s` were produced elsewhere and are only scored here, so the +sentence above does not apply to them. They were made with: + +- `mitonet-iso30-ortho`: allow_one_view False, cluster_iou_thr 0.75, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 29.91/29.96/29.96, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 23, mode ortho-plane consensus of xy, xz, yz, model MitoNet_v1, nms_kernel 3, nms_threshold 0.1, pixel_vote_thr 2 +- `mitonet-iso30-xy`: confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 29.91/29.96/29.96, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 23, mode xy stack, model MitoNet_v1, nms_kernel 3, nms_threshold 0.1 + +## Averaged over the test blocks + +| model | f1 | precision | recall | msa | sbd | dice | val_metric | iteration | +|---------------------|--------|-----------|--------|--------|--------|--------|------------|-----------| +| mitonet-iso30-ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0014 | 0.0054 | - | - | +| mitonet-iso30-xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0063 | 0.0036 | - | - | +| paper | 0.6905 | 0.5850 | 0.8469 | 0.3197 | 0.5722 | 0.8490 | 0.2927 | 14500 | + +## Per test block + +| model | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|---------------------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| mitonet-iso30-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0028 | 0.0107 | 16 | 30 | +| mitonet-iso30-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0126 | 0.0072 | 20 | 30 | +| paper | 4007 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 0.5188 | 0.8742 | 48 | 30 | +| mitonet-iso30-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| paper | 4009 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 0.6256 | 0.8237 | 57 | 43 | + +## Size filter sweep + +Applied to the finished segmentations, which is not the same as running the watershed with +that `min_size`, so read it as a diagnostic for how much of the error is sub-threshold +fragments. + +**Do not read a `min_size` off this table.** It would be tuning on the test data, and it +would generalize badly: these two test blocks happen to hold only large mitochondria, while +18% of the 717 mitochondria annotated in the training blocks are smaller than 20,000 voxels. +The default of 1,000 is the value consistent with the annotation, which has 99.7% of its +objects above it. + +| model | dataset | min_size | f1 | precision | recall | msa | n_pred | n_true | +|---------------------|---------|----------|--------|-----------|--------|--------|--------|--------| +| paper | 4007 | 1000 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 48 | 30 | +| paper | 4007 | 10000 | 0.8065 | 0.7812 | 0.8333 | 0.4515 | 32 | 30 | +| paper | 4007 | 20000 | 0.8333 | 0.8333 | 0.8333 | 0.4761 | 30 | 30 | +| paper | 4007 | 30000 | 0.8136 | 0.8276 | 0.8000 | 0.4573 | 29 | 30 | +| paper | 4009 | 1000 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 57 | 43 | +| paper | 4009 | 10000 | 0.8132 | 0.7708 | 0.8605 | 0.3697 | 48 | 43 | +| paper | 4009 | 20000 | 0.8605 | 0.8605 | 0.8605 | 0.4048 | 43 | 43 | +| paper | 4009 | 30000 | 0.8571 | 0.8780 | 0.8372 | 0.4134 | 41 | 43 | +| mitonet-iso30-ortho | 4007 | 1000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 16 | 30 | +| mitonet-iso30-ortho | 4007 | 10000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 10 | 30 | +| mitonet-iso30-ortho | 4007 | 20000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 4 | 30 | +| mitonet-iso30-ortho | 4007 | 30000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 1 | 30 | +| mitonet-iso30-ortho | 4009 | 1000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-ortho | 4009 | 10000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-ortho | 4009 | 20000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-ortho | 4009 | 30000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-xy | 4007 | 1000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 20 | 30 | +| mitonet-iso30-xy | 4007 | 10000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 15 | 30 | +| mitonet-iso30-xy | 4007 | 20000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 7 | 30 | +| mitonet-iso30-xy | 4007 | 30000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 6 | 30 | +| mitonet-iso30-xy | 4009 | 1000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-xy | 4009 | 10000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-xy | 4009 | 20000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| mitonet-iso30-xy | 4009 | 30000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | diff --git a/scripts/volume_em/RESULTS_mitonet_sweep_val.md b/scripts/volume_em/RESULTS_mitonet_sweep_val.md new file mode 100644 index 00000000..07b4ba91 --- /dev/null +++ b/scripts/volume_em/RESULTS_mitonet_sweep_val.md @@ -0,0 +1,166 @@ +# Volume EM mitochondria: MitoNet_v1 voxel-size sweep on the validation blocks + +MitoNet_v1 on the `val` blocks of `split-mito_vol_em_aniso2lvl_final.json` (4007, 4009), which it never saw. +Each block is resampled to the voxel size, segmented, and resized back to the native +25/5/5 nm grid, where it is scored against the untouched ground truth. `min_size` is scaled +to the same physical volume as synapse-net's 1,000 native voxels at every voxel size. + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +Metrics are averaged over the blocks, `n_pred` and `n_true` are summed. + +The best setting of each approach is selected here, by F1 with ties broken by msa, and then run +once on the test split. Nothing is selected on the test blocks. + +## A: z at 25 nm, xy resampled (xy stack) + +**Selected: `mitonet-z25-xy10-xy`** — raw contrast, xy mode, F1 0.1821, msa 0.0381. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 5 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0069 | 0.0226 | 412 | 65 | +| 7.5 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0125 | 0.0383 | 381 | 65 | +| 8 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0144 | 0.0399 | 357 | 65 | +| 10 | inverted | xy | 0.0073 | 0.0050 | 0.0135 | 0.0007 | 0.0175 | 0.0342 | 275 | 65 | +| 12.5 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0167 | 0.0304 | 228 | 65 | +| 15 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0065 | 0.0176 | 185 | 65 | +| 20 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0125 | 0.0283 | 185 | 65 | +| 25 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0131 | 0.0360 | 231 | 65 | +| 30 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0160 | 0.0363 | 243 | 65 | +| 5 | raw | xy | 0.0784 | 0.0517 | 0.1622 | 0.0162 | 0.1359 | 0.3454 | 150 | 65 | +| 7.5 | raw | xy | 0.1341 | 0.0889 | 0.2790 | 0.0332 | 0.1791 | 0.4810 | 184 | 65 | +| 8 | raw | xy | 0.1678 | 0.1160 | 0.3147 | 0.0399 | 0.2056 | 0.5018 | 168 | 65 | +| 10 | raw | xy | 0.1821 | 0.1270 | 0.3369 | 0.0381 | 0.2073 | 0.5219 | 173 | 65 | +| 12.5 | raw | xy | 0.1745 | 0.1250 | 0.3098 | 0.0372 | 0.1984 | 0.5040 | 168 | 65 | +| 15 | raw | xy | 0.1333 | 0.0992 | 0.2201 | 0.0376 | 0.1716 | 0.4603 | 150 | 65 | +| 20 | raw | xy | 0.0975 | 0.1063 | 0.1076 | 0.0235 | 0.1164 | 0.2787 | 105 | 65 | +| 25 | raw | xy | 0.0100 | 0.0079 | 0.0135 | 0.0010 | 0.0478 | 0.1367 | 78 | 65 | +| 30 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0192 | 0.0524 | 31 | 65 | + +## B: isotropic + +**Selected: `mitonet-iso12.5-ortho`** — raw contrast, ortho mode, F1 0.2346, msa 0.0561. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|-------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 10 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0090 | 0.0258 | 182 | 65 | +| 12.5 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0100 | 0.0246 | 158 | 65 | +| 15 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0081 | 0.0179 | 116 | 65 | +| 20 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0087 | 0.0131 | 117 | 65 | +| 25 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0086 | 0.0171 | 124 | 65 | +| 30 | inverted | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0046 | 0.0102 | 125 | 65 | +| 10 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0092 | 0.0332 | 596 | 65 | +| 12.5 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0093 | 0.0248 | 428 | 65 | +| 15 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0056 | 0.0143 | 283 | 65 | +| 20 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0091 | 0.0240 | 236 | 65 | +| 25 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0131 | 0.0360 | 231 | 65 | +| 30 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0336 | 223 | 65 | +| 10 | raw | ortho | 0.2170 | 0.1683 | 0.3325 | 0.0580 | 0.2564 | 0.5147 | 129 | 65 | +| 12.5 | raw | ortho | 0.2346 | 0.1857 | 0.3369 | 0.0561 | 0.2579 | 0.5133 | 120 | 65 | +| 15 | raw | ortho | 0.2014 | 0.1711 | 0.2601 | 0.0463 | 0.2370 | 0.4864 | 113 | 65 | +| 20 | raw | ortho | 0.0443 | 0.0463 | 0.0449 | 0.0074 | 0.1319 | 0.2894 | 69 | 65 | +| 25 | raw | ortho | 0.0286 | 0.0303 | 0.0270 | 0.0029 | 0.0728 | 0.1350 | 39 | 65 | +| 30 | raw | ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0251 | 12 | 65 | +| 10 | raw | xy | 0.1104 | 0.0650 | 0.3861 | 0.0231 | 0.1114 | 0.5549 | 396 | 65 | +| 12.5 | raw | xy | 0.1012 | 0.0611 | 0.3098 | 0.0251 | 0.1194 | 0.5469 | 335 | 65 | +| 15 | raw | xy | 0.0836 | 0.0541 | 0.1931 | 0.0257 | 0.1218 | 0.4899 | 239 | 65 | +| 20 | raw | xy | 0.0847 | 0.0805 | 0.1076 | 0.0217 | 0.1232 | 0.3025 | 135 | 65 | +| 25 | raw | xy | 0.0100 | 0.0079 | 0.0135 | 0.0010 | 0.0478 | 0.1367 | 78 | 65 | +| 30 | raw | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0253 | 0.0529 | 30 | 65 | + +## C: single settings (voxel size z/y/x in nm) + +**Selected: `mitonet-z30-xy8-xy`** — raw contrast, xy mode, F1 0.2228, msa 0.0602. + +| voxel_nm | contrast | mode | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|----------|----------|------|--------|-----------|--------|--------|--------|--------|--------|--------| +| 30/8/8 | inverted | xy | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0138 | 0.0367 | 308 | 65 | +| 30/8/8 | raw | xy | 0.2228 | 0.1579 | 0.3909 | 0.0602 | 0.2442 | 0.5533 | 157 | 65 | + +## Per block + +| setting | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|---------------------------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| mitonet-iso10-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0039 | 0.0201 | 71 | 37 | +| mitonet-iso10-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0315 | 111 | 28 | +| mitonet-iso10-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0052 | 0.0231 | 223 | 37 | +| mitonet-iso10-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0132 | 0.0433 | 373 | 28 | +| mitonet-iso10-ortho | 4007 | 0.2727 | 0.1895 | 0.4865 | 0.0730 | 0.2157 | 0.6378 | 95 | 37 | +| mitonet-iso10-ortho | 4009 | 0.1613 | 0.1471 | 0.1786 | 0.0429 | 0.2971 | 0.3917 | 34 | 28 | +| mitonet-iso10-xy | 4007 | 0.1176 | 0.0669 | 0.4865 | 0.0274 | 0.0932 | 0.5976 | 269 | 37 | +| mitonet-iso10-xy | 4009 | 0.1032 | 0.0630 | 0.2857 | 0.0188 | 0.1295 | 0.5122 | 127 | 28 | +| mitonet-iso12.5-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0052 | 0.0214 | 43 | 37 | +| mitonet-iso12.5-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0149 | 0.0278 | 115 | 28 | +| mitonet-iso12.5-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0044 | 0.0067 | 95 | 37 | +| mitonet-iso12.5-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0142 | 0.0430 | 333 | 28 | +| mitonet-iso12.5-ortho | 4007 | 0.2787 | 0.2000 | 0.4595 | 0.0660 | 0.2292 | 0.6336 | 85 | 37 | +| mitonet-iso12.5-ortho | 4009 | 0.1905 | 0.1714 | 0.2143 | 0.0462 | 0.2865 | 0.3931 | 35 | 28 | +| mitonet-iso12.5-xy | 4007 | 0.1128 | 0.0655 | 0.4054 | 0.0278 | 0.1044 | 0.5903 | 229 | 37 | +| mitonet-iso12.5-xy | 4009 | 0.0896 | 0.0566 | 0.2143 | 0.0224 | 0.1345 | 0.5035 | 106 | 28 | +| mitonet-iso15-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0050 | 0.0181 | 33 | 37 | +| mitonet-iso15-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0112 | 0.0177 | 83 | 28 | +| mitonet-iso15-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 50 | 37 | +| mitonet-iso15-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0113 | 0.0286 | 233 | 28 | +| mitonet-iso15-ortho | 4007 | 0.1695 | 0.1235 | 0.2703 | 0.0454 | 0.2023 | 0.5852 | 81 | 37 | +| mitonet-iso15-ortho | 4009 | 0.2333 | 0.2188 | 0.2500 | 0.0471 | 0.2716 | 0.3876 | 32 | 28 | +| mitonet-iso15-xy | 4007 | 0.0896 | 0.0549 | 0.2432 | 0.0232 | 0.1144 | 0.5336 | 164 | 37 | +| mitonet-iso15-xy | 4009 | 0.0777 | 0.0533 | 0.1429 | 0.0282 | 0.1292 | 0.4463 | 75 | 28 | +| mitonet-iso20-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0019 | 0.0080 | 38 | 37 | +| mitonet-iso20-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0154 | 0.0182 | 79 | 28 | +| mitonet-iso20-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0007 | 0.0012 | 72 | 37 | +| mitonet-iso20-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0176 | 0.0469 | 164 | 28 | +| mitonet-iso20-ortho | 4007 | 0.0460 | 0.0400 | 0.0541 | 0.0105 | 0.1446 | 0.3818 | 50 | 37 | +| mitonet-iso20-ortho | 4009 | 0.0426 | 0.0526 | 0.0357 | 0.0043 | 0.1192 | 0.1970 | 19 | 28 | +| mitonet-iso20-xy | 4007 | 0.0541 | 0.0360 | 0.1081 | 0.0111 | 0.0908 | 0.3665 | 111 | 37 | +| mitonet-iso20-xy | 4009 | 0.1154 | 0.1250 | 0.1071 | 0.0323 | 0.1557 | 0.2386 | 24 | 28 | +| mitonet-iso25-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0032 | 0.0100 | 50 | 37 | +| mitonet-iso25-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0243 | 74 | 28 | +| mitonet-iso25-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0007 | 0.0021 | 109 | 37 | +| mitonet-iso25-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0254 | 0.0699 | 122 | 28 | +| mitonet-iso25-ortho | 4007 | 0.0571 | 0.0606 | 0.0541 | 0.0058 | 0.0975 | 0.2067 | 33 | 37 | +| mitonet-iso25-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0482 | 0.0632 | 6 | 28 | +| mitonet-iso25-xy | 4007 | 0.0200 | 0.0159 | 0.0270 | 0.0020 | 0.0673 | 0.2037 | 63 | 37 | +| mitonet-iso25-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0282 | 0.0697 | 15 | 28 | +| mitonet-iso30-inv-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0006 | 0.0008 | 43 | 37 | +| mitonet-iso30-inv-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0086 | 0.0197 | 82 | 28 | +| mitonet-iso30-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0015 | 0.0038 | 101 | 37 | +| mitonet-iso30-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0264 | 0.0635 | 122 | 28 | +| mitonet-iso30-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0392 | 11 | 37 | +| mitonet-iso30-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0140 | 0.0109 | 1 | 28 | +| mitonet-iso30-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0435 | 0.0866 | 27 | 37 | +| mitonet-iso30-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0071 | 0.0193 | 3 | 28 | +| mitonet-z25-xy10-inv-xy | 4007 | 0.0146 | 0.0100 | 0.0270 | 0.0015 | 0.0143 | 0.0259 | 100 | 37 | +| mitonet-z25-xy10-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0206 | 0.0425 | 175 | 28 | +| mitonet-z25-xy10-xy | 4007 | 0.2179 | 0.1429 | 0.4595 | 0.0545 | 0.1877 | 0.5956 | 119 | 37 | +| mitonet-z25-xy10-xy | 4009 | 0.1463 | 0.1111 | 0.2143 | 0.0217 | 0.2269 | 0.4482 | 54 | 28 | +| mitonet-z25-xy12.5-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0155 | 0.0163 | 42 | 37 | +| mitonet-z25-xy12.5-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0180 | 0.0445 | 186 | 28 | +| mitonet-z25-xy12.5-xy | 4007 | 0.1911 | 0.1250 | 0.4054 | 0.0409 | 0.1691 | 0.5847 | 120 | 37 | +| mitonet-z25-xy12.5-xy | 4009 | 0.1579 | 0.1250 | 0.2143 | 0.0335 | 0.2276 | 0.4233 | 48 | 28 | +| mitonet-z25-xy15-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 37 | 37 | +| mitonet-z25-xy15-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0131 | 0.0352 | 148 | 28 | +| mitonet-z25-xy15-xy | 4007 | 0.1507 | 0.1009 | 0.2973 | 0.0370 | 0.1513 | 0.5270 | 109 | 37 | +| mitonet-z25-xy15-xy | 4009 | 0.1159 | 0.0976 | 0.1429 | 0.0382 | 0.1920 | 0.3937 | 41 | 28 | +| mitonet-z25-xy20-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0010 | 0.0013 | 50 | 37 | +| mitonet-z25-xy20-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0239 | 0.0553 | 135 | 28 | +| mitonet-z25-xy20-xy | 4007 | 0.0645 | 0.0460 | 0.1081 | 0.0149 | 0.1058 | 0.3529 | 87 | 37 | +| mitonet-z25-xy20-xy | 4009 | 0.1304 | 0.1667 | 0.1071 | 0.0321 | 0.1269 | 0.2044 | 18 | 28 | +| mitonet-z25-xy30-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0012 | 0.0037 | 113 | 37 | +| mitonet-z25-xy30-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0308 | 0.0689 | 130 | 28 | +| mitonet-z25-xy30-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0298 | 0.0674 | 28 | 37 | +| mitonet-z25-xy30-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0086 | 0.0374 | 3 | 28 | +| mitonet-z25-xy5-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0051 | 0.0168 | 197 | 37 | +| mitonet-z25-xy5-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0087 | 0.0284 | 215 | 28 | +| mitonet-z25-xy5-xy | 4007 | 0.1569 | 0.1034 | 0.3243 | 0.0324 | 0.1379 | 0.4890 | 116 | 37 | +| mitonet-z25-xy5-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.1339 | 0.2017 | 34 | 28 | +| mitonet-z25-xy7.5-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0096 | 0.0323 | 169 | 37 | +| mitonet-z25-xy7.5-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0153 | 0.0444 | 212 | 28 | +| mitonet-z25-xy7.5-xy | 4007 | 0.2222 | 0.1440 | 0.4865 | 0.0594 | 0.1833 | 0.5970 | 125 | 37 | +| mitonet-z25-xy7.5-xy | 4009 | 0.0460 | 0.0339 | 0.0714 | 0.0070 | 0.1749 | 0.3650 | 59 | 28 | +| mitonet-z25-xy8-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0118 | 0.0361 | 152 | 37 | +| mitonet-z25-xy8-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0170 | 0.0438 | 205 | 28 | +| mitonet-z25-xy8-xy | 4007 | 0.2368 | 0.1565 | 0.4865 | 0.0683 | 0.2002 | 0.6017 | 115 | 37 | +| mitonet-z25-xy8-xy | 4009 | 0.0988 | 0.0755 | 0.1429 | 0.0114 | 0.2110 | 0.4019 | 53 | 28 | +| mitonet-z30-xy8-inv-xy | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0088 | 0.0282 | 136 | 37 | +| mitonet-z30-xy8-inv-xy | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0188 | 0.0452 | 172 | 28 | +| mitonet-z30-xy8-xy | 4007 | 0.2937 | 0.1981 | 0.5676 | 0.0884 | 0.2217 | 0.6220 | 106 | 37 | +| mitonet-z30-xy8-xy | 4009 | 0.1519 | 0.1176 | 0.2143 | 0.0321 | 0.2667 | 0.4845 | 51 | 28 | diff --git a/scripts/volume_em/RESULTS_mitonet_voxel_size.md b/scripts/volume_em/RESULTS_mitonet_voxel_size.md new file mode 100644 index 00000000..f146efce --- /dev/null +++ b/scripts/volume_em/RESULTS_mitonet_voxel_size.md @@ -0,0 +1,41 @@ +# Volume EM mitochondria: MitoNet at the voxel sizes selected on the validation blocks + +Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95. +`dice` is the semantic foreground dice. `val_metric` is the best validation DiceLoss of the run, +read from its checkpoint, where lower is better. + +All models were segmented with identical preprocessing and identical post-processing +(seed_distance 1, boundary_threshold 0.12, area_threshold 200, min_size 1000). + +Segmentations passed with `-s` were produced elsewhere and are only scored here, so the +sentence above does not apply to them. They were made with: + +- `A-mitonet-z25-xy10-xy`: blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 25.0/10.0/10.0, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 250, mode xy stack, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, target_voxel_size_nm 25.0/10.0/10.0 +- `B-mitonet-iso12.5-ortho`: allow_one_view False, blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, cluster_iou_thr 0.75, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 12.5/12.5/12.5, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 320, mode ortho-plane consensus of xy, xz, yz, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, pixel_vote_thr 2, target_voxel_size_nm 12.5/12.5/12.5 +- `C-mitonet-z30-xy8-xy`: blocks /mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 29.91/8.0/8.0, invert False, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 326, mode xy stack, model MitoNet_v1, n_lost_to_native 0, nms_kernel 3, nms_threshold 0.1, target_voxel_size_nm 30.0/8.0/8.0 +- `mitonet-iso30-ortho`: allow_one_view False, cluster_iou_thr 0.75, confidence_thr 0.3, empanada_napari 1.2.1, inference_voxel_size_nm 29.91/29.96/29.96, label_divisor 1000, median_kernel_size 5, min_extent 4, min_size 23, mode ortho-plane consensus of xy, xz, yz, model MitoNet_v1, nms_kernel 3, nms_threshold 0.1, pixel_vote_thr 2 + +## Averaged over the test blocks + +| model | f1 | precision | recall | msa | sbd | dice | val_metric | iteration | +|-------------------------|--------|-----------|--------|--------|--------|--------|------------|-----------| +| A-mitonet-z25-xy10-xy | 0.0938 | 0.0698 | 0.1430 | 0.0263 | 0.1665 | 0.3675 | - | - | +| B-mitonet-iso12.5-ortho | 0.1013 | 0.0828 | 0.1314 | 0.0186 | 0.1831 | 0.3281 | - | - | +| C-mitonet-z30-xy8-xy | 0.1209 | 0.0867 | 0.1996 | 0.0267 | 0.1718 | 0.4073 | - | - | +| mitonet-iso30-ortho | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0014 | 0.0054 | - | - | +| paper | 0.6905 | 0.5850 | 0.8469 | 0.3197 | 0.5722 | 0.8490 | 0.2927 | 14500 | + +## Per test block + +| model | dataset | f1 | precision | recall | msa | sbd | dice | n_pred | n_true | +|-------------------------|---------|--------|-----------|--------|--------|--------|--------|--------|--------| +| A-mitonet-z25-xy10-xy | 4007 | 0.0625 | 0.0455 | 0.1000 | 0.0172 | 0.1413 | 0.3691 | 66 | 30 | +| B-mitonet-iso12.5-ortho | 4007 | 0.0706 | 0.0545 | 0.1000 | 0.0134 | 0.1519 | 0.3225 | 55 | 30 | +| C-mitonet-z30-xy8-xy | 4007 | 0.1010 | 0.0725 | 0.1667 | 0.0231 | 0.1587 | 0.4163 | 69 | 30 | +| mitonet-iso30-ortho | 4007 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0028 | 0.0107 | 16 | 30 | +| paper | 4007 | 0.6410 | 0.5208 | 0.8333 | 0.3196 | 0.5188 | 0.8742 | 48 | 30 | +| A-mitonet-z25-xy10-xy | 4009 | 0.1250 | 0.0941 | 0.1860 | 0.0355 | 0.1918 | 0.3658 | 85 | 43 | +| B-mitonet-iso12.5-ortho | 4009 | 0.1321 | 0.1111 | 0.1628 | 0.0239 | 0.2143 | 0.3338 | 63 | 43 | +| C-mitonet-z30-xy8-xy | 4009 | 0.1408 | 0.1010 | 0.2326 | 0.0304 | 0.1848 | 0.3982 | 99 | 43 | +| mitonet-iso30-ortho | 4009 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0.0000 | 0 | 43 | +| paper | 4009 | 0.7400 | 0.6491 | 0.8605 | 0.3198 | 0.6256 | 0.8237 | 57 | 43 | diff --git a/scripts/volume_em/evaluation/compare_models.py b/scripts/volume_em/evaluation/compare_models.py new file mode 100644 index 00000000..65bdac60 --- /dev/null +++ b/scripts/volume_em/evaluation/compare_models.py @@ -0,0 +1,403 @@ +"""Compare volume EM mitochondria models on the held-out test blocks. + +Segments every test block with every model, using identical preprocessing and identical +post-processing, and reports instance and semantic metrics per block. Holding the post-processing +fixed matters: tuning it per model would fold a post-processing difference into what is meant to be +a comparison of the trained networks. + +The segmentations are cached, so the script can be re-run to add a model without redoing the others. + +Segmentations produced elsewhere, such as the MitoNet baseline of segment_mitonet.py, can be scored +alongside with '-s'. They are only scored here, not segmented, so the note on identical +post-processing does not apply to them; their settings are read from the attributes of their 'seg'. + +Metrics: + f1 / precision / recall instance matching at an IoU of 0.5, the metric the volume EM pipeline + has always reported (`elf.evaluation.matching`) + msa mean segmentation accuracy, averaged over IoU thresholds 0.5 to 0.95 + sbd symmetric best dice + dice semantic dice of the foreground, independent of the instance splitting + n_pred / n_true instance counts, which show over- or under-segmentation directly + val_metric the best validation DiceLoss of the run, read from its checkpoint + +Usage: + python compare_models.py -m name=/path/to/checkpoint [-m ...] [-s name=/path/to/segmentations] -o RESULTS.md +""" + +import argparse +import importlib.util +import json +import os +import sys + +import h5py +import numpy as np +import pandas as pd +import torch +import torch_em +from elf.evaluation import dice_score, matching, mean_segmentation_accuracy, symmetric_best_dice_score + +from synapse_net.inference.util import parse_tiling +from synapse_net.training.split import _resolve_split + +TEST_SPLIT = "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split" +SEGMENTATION_ROOT = "/mnt/lustre-grete/usr/u15205/volume-em/repro-comparison" +TITLE = "Volume EM mitochondria: reproduction vs. the published model" + +_INFERENCE_SCRIPT = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "inference", "run_mitochondria_vol_em_segmentation.py", +) + + +def _load_inference_module(): + """Load the segmentation script, so the comparison uses exactly the shipped inference path.""" + spec = importlib.util.spec_from_file_location("run_mitochondria_vol_em_segmentation", _INFERENCE_SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def find_test_blocks(test_split: str): + """Find the held-out test blocks, as a mapping of dataset name to filepath.""" + blocks = {} + for dataset in sorted(os.listdir(test_split)): + folder = os.path.join(test_split, dataset) + if not os.path.isdir(folder): + continue + files = sorted(name for name in os.listdir(folder) if name.endswith(".h5")) + if len(files) != 1: + raise ValueError(f"Expect exactly one test block in {folder}, found {len(files)}.") + blocks[dataset] = os.path.join(folder, files[0]) + if not blocks: + raise ValueError(f"Did not find any test blocks in {test_split}.") + return blocks + + +def find_split_blocks(split_file: str, key: str): + """Find the blocks of one key of a split file, as a mapping of data root name to filepath. + + Naming a block by its data root keeps the names of the validation blocks parallel to the test + blocks ('4007', '4009'); it requires at most one block per root. + """ + with open(split_file) as f: + entries = json.load(f).get(key, []) + paths = _resolve_split(split_file, {}, [key])[0] + blocks = {} + for entry, path in zip(entries, paths): + name = entry.partition("/")[0] + if name in blocks: + raise ValueError(f"The '{key}' split has more than one block from '{name}', so they cannot " + "be named by their data root.") + blocks[name] = path + if not blocks: + raise ValueError(f"The split file {split_file} has no '{key}' entries.") + return blocks + + +def segment_block(inference, model_path, block_path, output_path, tiling, force=False): + """Segment one test block, reusing a cached result unless `force` is set.""" + if os.path.exists(output_path) and not force: + with h5py.File(output_path, "r") as f: + if "seg" in f: + return + with h5py.File(block_path, "r") as f: + raw = f["raw"][:] + + segmentation = inference.segment_mitochondria_vol_em( + raw, + white_patch_min_size=20, + model_path=model_path, + tiling=tiling, + preprocess=torch_em.transform.raw.normalize_percentile, + min_size=inference.MIN_SIZE, + seed_distance=inference.SEED_DISTANCE, + boundary_threshold=inference.BOUNDARY_THRESHOLD, + area_threshold=inference.AREA_THRESHOLD, + verbose=False, + ) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + with h5py.File(output_path, "a") as f: + if "seg" in f: + del f["seg"] + f.create_dataset("seg", data=segmentation, compression="gzip") + + +def evaluate_block(segmentation_path, block_path): + """Compute the instance and semantic metrics for one segmentation.""" + with h5py.File(segmentation_path, "r") as f: + seg = f["seg"][:] + with h5py.File(block_path, "r") as f: + gt = f["labels/mitochondria"][:] + + stats = matching(segmentation=seg, groundtruth=gt, threshold=0.5, criterion="iou", ignore_label=0) + return { + "f1": stats["f1"], + "precision": stats["precision"], + "recall": stats["recall"], + "msa": mean_segmentation_accuracy(seg, gt), + "sbd": symmetric_best_dice_score(seg, gt), + "dice": dice_score(seg, gt), + "n_pred": int(len(np.unique(seg)) - (1 if 0 in seg else 0)), + "n_true": int(len(np.unique(gt)) - (1 if 0 in gt else 0)), + } + + +def apply_size_filter(seg, min_size): + """Drop instances below a size, on a finished segmentation.""" + ids, counts = np.unique(seg, return_counts=True) + too_small = ids[(counts < min_size) & (ids != 0)] + if too_small.size == 0: + return seg + seg = seg.copy() + seg[np.isin(seg, too_small)] = 0 + return seg + + +def sweep_size_filter(segmentation_path, block_path, thresholds): + """Re-score a segmentation under several size filters. + + This filters the finished segmentation, which is not the same as running the watershed with that + 'min_size' -- the size filter inside 'segment_mitochondria' runs before the hole filling. + + Read it as a diagnostic for how much of the error is sub-threshold fragments, not as a tuning run. + Picking a threshold off these numbers would be tuning on the test data, and it would generalize + badly: the two test blocks happen to hold only large mitochondria, while 18% of the mitochondria + annotated in the training blocks are smaller than 20,000 voxels. + """ + with h5py.File(segmentation_path, "r") as f: + seg0 = f["seg"][:] + with h5py.File(block_path, "r") as f: + gt = f["labels/mitochondria"][:] + + rows = [] + for min_size in thresholds: + seg = apply_size_filter(seg0, min_size) + stats = matching(segmentation=seg, groundtruth=gt, threshold=0.5, criterion="iou", ignore_label=0) + rows.append({ + "min_size": min_size, + "f1": stats["f1"], "precision": stats["precision"], "recall": stats["recall"], + "msa": mean_segmentation_accuracy(seg, gt), + "n_pred": int(len(np.unique(seg)) - (1 if 0 in seg else 0)), + "n_true": int(len(np.unique(gt)) - (1 if 0 in gt else 0)), + }) + return rows + + +def read_val_metric(model_path): + """Read the best validation metric of a run out of its checkpoint.""" + checkpoint = model_path if os.path.isdir(model_path) else os.path.dirname(model_path) + best = os.path.join(checkpoint, "best.pt") + if not os.path.exists(best): + return None, None + try: + state = torch.load(best, map_location="cpu", weights_only=False) + except AttributeError: + # An unrepaired checkpoint of the published model, see repair_checkpoint.py. + return None, None + return state.get("best_metric"), state.get("iteration") + + +def find_precomputed_segmentations(folder, blocks): + """Map each test block to its segmentation in a folder of precomputed ones.""" + paths = {dataset: os.path.join(folder, f"{dataset}.h5") for dataset in blocks} + missing = [path for path in paths.values() if not os.path.exists(path)] + if missing: + raise ValueError(f"Missing segmentations in {folder}: {', '.join(missing)}") + return paths + + +def read_segmentation_settings(segmentation_path): + """Read the settings a precomputed segmentation was made with, from the attributes of its 'seg'.""" + with h5py.File(segmentation_path, "r") as f: + attrs = dict(f["seg"].attrs) + + def render(value): + if isinstance(value, np.ndarray): + return "/".join(str(v) for v in value.tolist()) + return str(value) + + return ", ".join(f"{key} {render(value)}" for key, value in attrs.items()) + + +def _markdown_table(frame: pd.DataFrame, float_format: str = "{:.4f}") -> str: + """Render a dataframe as a markdown table, without depending on tabulate.""" + def render(value): + if value is None or value is pd.NA or (isinstance(value, float) and np.isnan(value)): + return "-" + return float_format.format(value) if isinstance(value, float) else str(value) + + header = list(frame.columns) + rows = [[render(value) for value in row] for row in frame.itertuples(index=False)] + widths = [max(len(header[i]), *(len(row[i]) for row in rows)) if rows else len(header[i]) + for i in range(len(header))] + lines = ["| " + " | ".join(name.ljust(width) for name, width in zip(header, widths)) + " |", + "|" + "|".join("-" * (width + 2) for width in widths) + "|"] + lines += ["| " + " | ".join(value.ljust(width) for value, width in zip(row, widths)) + " |" for row in rows] + return "\n".join(lines) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("-m", "--model", action="append", default=[], metavar="NAME=PATH", + help="A model to compare, as 'name=checkpoint folder'. Repeat for each model.") + parser.add_argument("-s", "--segmentation", action="append", default=[], metavar="NAME=FOLDER", + help="Precomputed segmentations to score, as 'name=folder' with one '.h5' " + "per test block, holding 'seg'. Repeat for each. These are not segmented here.") + parser.add_argument("--title", default=TITLE, help="The title of the report.") + parser.add_argument("--test_split", default=TEST_SPLIT, help="The folder with the held-out test blocks.") + parser.add_argument("--split_file", help="Score the blocks of a split file instead of the test split.") + parser.add_argument("--split_key", default="val", help="The key of the split file to score.") + parser.add_argument("--segmentation_root", default=SEGMENTATION_ROOT, + help="Where the segmentations are cached.") + parser.add_argument("-o", "--output", help="Write the report to this markdown file.") + parser.add_argument("--tile_shape", type=int, nargs=3, help="The tile shape for prediction, in ZYX.") + parser.add_argument("--halo", type=int, nargs=3, help="The halo for prediction, in ZYX.") + parser.add_argument("--force", action="store_true", help="Recompute segmentations that are already cached.") + parser.add_argument("--size_filter_sweep", type=int, nargs="*", metavar="MIN_SIZE", + help="Also re-score the segmentations under these size filters, as a diagnostic " + "for how much of the error is sub-threshold fragments.") + args = parser.parse_args() + + def parse_entries(entries, flag): + parsed = {} + for entry in entries: + if "=" not in entry: + parser.error(f"Expect '{flag} name=path', got {entry!r}.") + name, path = entry.split("=", 1) + if not os.path.exists(path): + parser.error(f"The path of {name!r} does not exist: {path}") + parsed[name] = path + return parsed + + models = parse_entries(args.model, "--model") + segmentations = parse_entries(args.segmentation, "--segmentation") + if not models and not segmentations: + parser.error("Pass at least one '--model' or '--segmentation'.") + if set(models) & set(segmentations): + parser.error(f"Names are used twice: {', '.join(sorted(set(models) & set(segmentations)))}") + + inference = _load_inference_module() + tiling = parse_tiling(args.tile_shape, args.halo) + if args.split_file: + blocks = find_split_blocks(args.split_file, args.split_key) + else: + blocks = find_test_blocks(args.test_split) + print(f"Comparing {len(models)} model(s) and {len(segmentations)} precomputed segmentation(s) " + f"on {len(blocks)} test block(s): {', '.join(blocks)}\n") + + # The blocks of a split file are named like the test blocks, so their cached segmentations go + # into a folder of their own; otherwise a cached test segmentation would be scored against them. + cache_root = os.path.join(args.segmentation_root, args.split_key) if args.split_file else args.segmentation_root + segmentation_paths = { + name: {dataset: os.path.join(cache_root, name, f"{dataset}.h5") for dataset in blocks} + for name in models + } + for name, folder in segmentations.items(): + segmentation_paths[name] = find_precomputed_segmentations(folder, blocks) + + rows = [] + for name, model_path in models.items(): + val_metric, iteration = read_val_metric(model_path) + for dataset, block_path in blocks.items(): + output_path = segmentation_paths[name][dataset] + print(f" {name} / {dataset} ...", flush=True) + segment_block(inference, model_path, block_path, output_path, tiling, force=args.force) + metrics = evaluate_block(output_path, block_path) + metrics.update(model=name, dataset=dataset, val_metric=val_metric, iteration=iteration) + rows.append(metrics) + for name in segmentations: + for dataset, block_path in blocks.items(): + print(f" {name} / {dataset} (precomputed) ...", flush=True) + metrics = evaluate_block(segmentation_paths[name][dataset], block_path) + metrics.update(model=name, dataset=dataset, val_metric=None, iteration=None) + rows.append(metrics) + + results = pd.DataFrame(rows) + columns = ["model", "dataset", "f1", "precision", "recall", "msa", "sbd", "dice", "n_pred", "n_true"] + per_block = results[columns].sort_values(["dataset", "model"]) + + mean_columns = ["f1", "precision", "recall", "msa", "sbd", "dice"] + averaged = results.groupby("model")[mean_columns].mean().reset_index() + val = results.groupby("model")[["val_metric", "iteration"]].first().reset_index() + # A precomputed segmentation has no iteration, which would otherwise turn the column into floats. + val["iteration"] = val["iteration"].astype("Int64") + averaged = averaged.merge(val, on="model") + + baseline_only = [ + "", + "**Status: baseline only.** Only one model is in this report, so there is nothing to compare", + "against yet. Re-run with the reproduction runs added (`-m seed42=... -m seed43=...`) once they", + "finish; the segmentations already here are cached and are not recomputed.", + ] if len(models) + len(segmentations) < 2 else [] + + precomputed = [ + "", + "Segmentations passed with `-s` were produced elsewhere and are only scored here, so the", + "sentence above does not apply to them. They were made with:", + "", + *(f"- `{name}`: {read_segmentation_settings(next(iter(segmentation_paths[name].values())))}" + for name in segmentations), + ] if segmentations else [] + + report = [ + f"# {args.title}", + *baseline_only, + "", + "Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95.", + "`dice` is the semantic foreground dice. `val_metric` is the best validation DiceLoss of the run,", + "read from its checkpoint, where lower is better.", + "", + "All models were segmented with identical preprocessing and identical post-processing", + f"(seed_distance {inference.SEED_DISTANCE}, boundary_threshold {inference.BOUNDARY_THRESHOLD}, " + f"area_threshold {inference.AREA_THRESHOLD}, min_size {inference.MIN_SIZE}).", + *precomputed, + "", + "## Averaged over the test blocks", + "", + _markdown_table(averaged), + "", + "## Per test block", + "", + _markdown_table(per_block), + "", + ] + if args.size_filter_sweep: + sweep_rows = [] + for name, paths in segmentation_paths.items(): + for dataset, block_path in blocks.items(): + for row in sweep_size_filter(paths[dataset], block_path, args.size_filter_sweep): + row.update(model=name, dataset=dataset) + sweep_rows.append(row) + sweep = pd.DataFrame(sweep_rows)[ + ["model", "dataset", "min_size", "f1", "precision", "recall", "msa", "n_pred", "n_true"] + ] + report += [ + "## Size filter sweep", + "", + "Applied to the finished segmentations, which is not the same as running the watershed with", + "that `min_size`, so read it as a diagnostic for how much of the error is sub-threshold", + "fragments.", + "", + "**Do not read a `min_size` off this table.** It would be tuning on the test data, and it", + "would generalize badly: these two test blocks happen to hold only large mitochondria, while", + "18% of the 717 mitochondria annotated in the training blocks are smaller than 20,000 voxels.", + "The default of 1,000 is the value consistent with the annotation, which has 99.7% of its", + "objects above it.", + "", + _markdown_table(sweep), + "", + ] + + report = "\n".join(report) + print("\n" + report) + + if args.output: + with open(args.output, "w") as f: + f.write(report) + print(f"Wrote {args.output}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/evaluation/microsam_sweep_score.yaml b/scripts/volume_em/evaluation/microsam_sweep_score.yaml new file mode 100644 index 00000000..4437ed0c --- /dev/null +++ b/scripts/volume_em/evaluation/microsam_sweep_score.yaml @@ -0,0 +1,23 @@ +# Scores the micro-sam voxel-size sweep of microsam_sweep_val.yaml on the validation blocks and +# selects the best setting per approach, see score_voxel_size_sweep.py. Writes the report into the +# repo and the selection to best.json in the sweep root. Scoring is at the native resolution, +# ~6-8 GB per worker, so 8 workers fit the 128 GB of the profile. +# +# Submit after the sweep, with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/microsam_sweep_score.yaml --dependency afterok: + +slurm_profile: grete_std_cpu + +slurm: + job_name: "volem-microsam-sweep-score" + +env: + type: micromamba + path: /mnt/vast-nhr/home/freckmann15/u15205/micromamba/envs/synapse + +commands: + - "python /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/score_voxel_size_sweep.py + --sweep_root /mnt/lustre-grete/usr/u15205/volume-em/microsam-sweep/val + --workers 8 + -o /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/RESULTS_microsam_sweep_val.md" diff --git a/scripts/volume_em/evaluation/microsam_sweep_val.yaml b/scripts/volume_em/evaluation/microsam_sweep_val.yaml new file mode 100644 index 00000000..443882fa --- /dev/null +++ b/scripts/volume_em/evaluation/microsam_sweep_val.yaml @@ -0,0 +1,41 @@ +# micro-sam voxel-size sweep on the validation blocks of the pinned split, see segment_microsam.py. +# The same grid as the MitoNet sweep (mitonet_sweep_val.yaml), on the raw contrast only, since the +# inverted contrast scored F1 ~0 at every voxel size for MitoNet: +# B isotropic 30 25 20 15 12.5 10 nm +# A z at the native 25 nm, xy 30 20 15 12.5 10 8 7.5 5 3.75 nm; xy 25 is B's iso25. 3.75 extends +# the grid by a step, because the first run selected its edge, 5 nm +# C 30 nm in z and 8 nm in xy +# micro-sam has only the xy mode (2D slices merged in z). B runs first, so A at xy 25 finds B's +# segmentation and skips it; settings that already exist are skipped, so a grid extension only +# computes the new ones. Score with microsam_sweep_score.yaml. +# +# Every image or tile reaches the SAM encoder at exactly 1024 px, so the tiled settings (xy 7.5, 5 +# and 3.75) segment a padded canvas of 16, 25 and 25 tiles per slice; a full A100 keeps that affordable. +# vit_b_em_organelles is read from ~/.cache/micro_sam/models, so the job needs no internet access. +# +# Submit with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/microsam_sweep_val.yaml [--dry-run] + +slurm_profile: grete_shared + +slurm: + job_name: "volem-microsam-sweep-val" + time: "12:00:00" + +env: + type: micromamba + path: /mnt/lustre-grete/usr/u15205/envs/micro-sam + +commands: + - "export CUDA_VISIBLE_DEVICES=0 + SEG=/mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/segment_microsam.py + SPLIT=/mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json + ROOT=/mnt/lustre-grete/usr/u15205/volume-em/microsam-sweep/val" + - "for v in 30 25 20 15 12.5 10; do + python $SEG --voxel_size $v $v $v --split_file $SPLIT --split_key val --segmentation_root $ROOT; + done" + - "for xy in 30 25 20 15 12.5 10 8 7.5 5 3.75; do + python $SEG --voxel_size 25 $xy $xy --split_file $SPLIT --split_key val --segmentation_root $ROOT; + done" + - "python $SEG --voxel_size 30 8 8 --split_file $SPLIT --split_key val --segmentation_root $ROOT" diff --git a/scripts/volume_em/evaluation/microsam_test_selected.yaml b/scripts/volume_em/evaluation/microsam_test_selected.yaml new file mode 100644 index 00000000..1324b5d0 --- /dev/null +++ b/scripts/volume_em/evaluation/microsam_test_selected.yaml @@ -0,0 +1,23 @@ +# The one test run of the micro-sam setting selected on the validation blocks, see microsam_sweep_val.yaml +# and RESULTS_microsam_sweep_val.md. Only approach A has a setting above the F1 floor of 0.05: z at the +# native 25 nm and xy at the native 5 nm, i.e. no resampling, only the tiled 1024 px canvas. The grid was +# extended to 3.75 nm to check that 5 nm is not just its edge. Writes microsam-z25-xy5-xy/ under +# /mnt/lustre-grete/usr/u15205/volume-em/repro-comparison, which compare_models.py then scores with '-s'. +# +# Submit with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/microsam_test_selected.yaml [--dry-run] + +slurm_profile: grete_shared + +slurm: + job_name: "volem-microsam-test" + time: "02:00:00" + +env: + type: micromamba + path: /mnt/lustre-grete/usr/u15205/envs/micro-sam + +commands: + - "CUDA_VISIBLE_DEVICES=0 python /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/segment_microsam.py + --voxel_size 25 5 5" diff --git a/scripts/volume_em/evaluation/mitonet_iso30.yaml b/scripts/volume_em/evaluation/mitonet_iso30.yaml new file mode 100644 index 00000000..1bde0cab --- /dev/null +++ b/scripts/volume_em/evaluation/mitonet_iso30.yaml @@ -0,0 +1,25 @@ +# MitoNet baseline at 30 nm isotropic on the held-out test blocks, see segment_mitonet.py. +# Writes mitonet-iso30-ortho/ and mitonet-iso30-xy/ under /mnt/lustre-grete/usr/u15205/volume-em/repro-comparison, +# which compare_models.py then scores with '-s'. The volumes are small at 30 nm (107 x 267 x 267), so a +# MIG slice is ample. (The grete_interactive profile asks for 1g.20gb slices, which the partition no longer has.) +# +# The environment is built from the synapse repo's env_empanada.yaml. The MitoNet_v1 weights are +# downloaded to ~/.empanada on first use; fetch them on the login node if the compute nodes have no +# internet access. +# +# Submit with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/mitonet_iso30.yaml [--dry-run] + +slurm_profile: grete_interactive_10gb + +slurm: + job_name: "volem-mitonet-iso30" + time: "01:00:00" + +env: + type: micromamba + path: /mnt/lustre-grete/usr/u15205/envs/empanada + +commands: + - "CUDA_VISIBLE_DEVICES=0 python /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/segment_mitonet.py" diff --git a/scripts/volume_em/evaluation/mitonet_sweep_score.yaml b/scripts/volume_em/evaluation/mitonet_sweep_score.yaml new file mode 100644 index 00000000..efb42eb5 --- /dev/null +++ b/scripts/volume_em/evaluation/mitonet_sweep_score.yaml @@ -0,0 +1,23 @@ +# Scores the MitoNet voxel-size sweep of mitonet_sweep_val.yaml on the validation blocks and selects +# the best setting per approach, see score_voxel_size_sweep.py. Writes the report into the repo and the +# selection to best.json in the sweep root. Scoring is at the native resolution, ~6-8 GB per worker, +# so 8 workers fit the 128 GB of the profile. +# +# Submit after the sweep, with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/mitonet_sweep_score.yaml --dependency afterok: + +slurm_profile: grete_std_cpu + +slurm: + job_name: "volem-mitonet-sweep-score" + +env: + type: micromamba + path: /mnt/vast-nhr/home/freckmann15/u15205/micromamba/envs/synapse + +commands: + - "python /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/score_voxel_size_sweep.py + --sweep_root /mnt/lustre-grete/usr/u15205/volume-em/mitonet-sweep/val + --workers 8 + -o /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/RESULTS_mitonet_sweep_val.md" diff --git a/scripts/volume_em/evaluation/mitonet_sweep_val.yaml b/scripts/volume_em/evaluation/mitonet_sweep_val.yaml new file mode 100644 index 00000000..b60ff995 --- /dev/null +++ b/scripts/volume_em/evaluation/mitonet_sweep_val.yaml @@ -0,0 +1,36 @@ +# MitoNet_v1 voxel-size sweep on the validation blocks of the pinned split, see segment_mitonet.py. +# Every setting runs on the raw and on the inverted contrast: +# B isotropic 30 25 20 15 12.5 10 nm, ortho-plane consensus and xy stack +# A z at the native 25 nm, xy 30 20 15 12.5 10 8 7.5 5 nm, xy stack; xy 25 is B's iso25-xy +# C 30 nm in z and 8 nm in xy, the voxel size MitoNet_v1 was probably trained at, xy stack +# B runs first, so A at xy 25 finds B's segmentation and skips it. Settings that already exist are +# skipped, so a grid extension only computes the new ones. Score with mitonet_sweep_score.yaml. +# +# Submit with the runner of the synapse repo: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/evaluation/mitonet_sweep_val.yaml [--dry-run] + +slurm_profile: grete_interactive_10gb + +slurm: + job_name: "volem-mitonet-sweep-val" + time: "06:00:00" + +env: + type: micromamba + path: /mnt/lustre-grete/usr/u15205/envs/empanada + +commands: + - "export CUDA_VISIBLE_DEVICES=0 + SEG=/mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/evaluation/segment_mitonet.py + SPLIT=/mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json + ROOT=/mnt/lustre-grete/usr/u15205/volume-em/mitonet-sweep/val" + - "for inv in '' --invert; do for v in 30 25 20 15 12.5 10; do + python $SEG --voxel_size $v $v $v --modes ortho xy $inv --split_file $SPLIT --split_key val --segmentation_root $ROOT; + done; done" + - "for inv in '' --invert; do for xy in 30 25 20 15 12.5 10 8 7.5 5; do + python $SEG --voxel_size 25 $xy $xy --modes xy $inv --split_file $SPLIT --split_key val --segmentation_root $ROOT; + done; done" + - "for inv in '' --invert; do + python $SEG --voxel_size 30 8 8 --modes xy $inv --split_file $SPLIT --split_key val --segmentation_root $ROOT; + done" diff --git a/scripts/volume_em/evaluation/score_voxel_size_sweep.py b/scripts/volume_em/evaluation/score_voxel_size_sweep.py new file mode 100644 index 00000000..c23132e7 --- /dev/null +++ b/scripts/volume_em/evaluation/score_voxel_size_sweep.py @@ -0,0 +1,240 @@ +"""Score a voxel-size sweep of a baseline method and select the best setting per approach. + +The sweep is made by segment_mitonet.py or segment_microsam.py on the validation blocks of the pinned +split, one folder per setting and mode; a sweep root holds the folders of one method. The settings +fall into three groups: + A z stays at the native 25 nm and only xy is resampled; the xy stack only, since MitoNet's + ortho-plane consensus needs isotropic voxels and micro-sam has no ortho-plane mode + B isotropic resampling of all axes; for MitoNet the ortho-plane consensus and the xy stack + C any other single setting, such as 30 nm in z and 8 nm in xy, the voxel size MitoNet_v1 was + probably trained at +25 nm isotropic belongs to both A and B. A setting may be run on the raw and on the inverted contrast. + +Every folder is scored with compare_models.evaluate_block against the native ground truth, so the +numbers are the ones compare_models.py reports. The metrics are cached in a 'metrics.json' per +folder, so extending the grid only scores the new settings. + +The best setting of an approach is the one with the highest F1 averaged over the validation blocks, +ties broken by msa. It is selected here, on the validation blocks, and then run once on the test +split; nothing is selected on the test blocks. + +Usage: + python score_voxel_size_sweep.py --sweep_root ROOT [-o RESULTS.md] [--workers 8] +""" + +import argparse +import json +import os +import sys +from multiprocessing import Pool + +import h5py +import numpy as np +import pandas as pd + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import compare_models # noqa: E402 + +SPLIT_FILE = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "training", "split-mito_vol_em_aniso2lvl_final.json", +) +NATIVE_Z = 25.0 + +# Below this, no setting is presented as working: the best of a set of failures is not a result. +F1_FLOOR = 0.05 + +METRICS = ["f1", "precision", "recall", "msa", "sbd", "dice"] + + +def read_setting(folder, datasets): + """Read the setting of a sweep folder from the attributes of its segmentations.""" + paths = {dataset: os.path.join(folder, f"{dataset}.h5") for dataset in datasets} + missing = [path for path in paths.values() if not os.path.exists(path)] + if missing: + raise ValueError(f"Missing segmentations in {folder}: {', '.join(missing)}") + with h5py.File(next(iter(paths.values())), "r") as f: + attrs = f["seg"].attrs + z, y, x = (float(v) for v in attrs["target_voxel_size_nm"]) + setting = { + "folder": os.path.basename(folder), + "model": str(attrs["model"]), + "voxel_size": (z, y, x), + "contrast": "inverted" if bool(attrs["invert"]) else "raw", + "mode": "ortho" if str(attrs["mode"]).startswith("ortho") else "xy", + } + approaches = [] + if z == NATIVE_Z and y == x and setting["mode"] == "xy": + approaches.append("A") + if z == y == x: + approaches.append("B") + setting["approaches"] = approaches or ["C"] + return setting, paths + + +def voxel_label(approach, voxel_size): + """The voxel size a table row is keyed by: xy for A, the isotropic size for B, all axes for C.""" + z, y, x = voxel_size + if approach == "A": + return y + if approach == "B": + return z + return "/".join(f"{v:g}" for v in voxel_size) + + +def _score(task): + folder, dataset, segmentation_path, block_path = task + metrics = compare_models.evaluate_block(segmentation_path, block_path) + # Plain Python numbers, so the metrics can be cached as json. + metrics = {key: int(value) if key.startswith("n_") else float(value) for key, value in metrics.items()} + return folder, dataset, metrics + + +def score_sweep(folders, blocks, workers): + """Score every folder on every block, reusing the cached metrics of unchanged segmentations.""" + cached, tasks = {}, [] + for folder, paths in folders.items(): + cache_path = os.path.join(folder, "metrics.json") + cache = {} + if os.path.exists(cache_path): + with open(cache_path) as f: + cache = json.load(f) + cached[folder] = {} + for dataset, path in paths.items(): + entry = cache.get(dataset) + if entry is not None and entry.get("seg_mtime") == os.path.getmtime(path): + cached[folder][dataset] = entry + else: + tasks.append((folder, dataset, path, blocks[dataset])) + + print(f"Scoring {len(tasks)} segmentation(s), {sum(map(len, cached.values()))} cached", flush=True) + with Pool(workers) as pool: + for folder, dataset, metrics in pool.imap_unordered(_score, tasks): + metrics["seg_mtime"] = os.path.getmtime(folders[folder][dataset]) + cached[folder][dataset] = metrics + print(f" {os.path.basename(folder)} / {dataset}: f1 {metrics['f1']:.4f}", flush=True) + with open(os.path.join(folder, "metrics.json"), "w") as f: + json.dump(cached[folder], f, indent=2) + return cached + + +def select_best(frame): + """The best row of an approach: highest mean F1, ties broken by msa.""" + return frame.sort_values(["f1", "msa"], ascending=False).iloc[0] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("--sweep_root", required=True, + help="The folder with one folder per setting and mode, of one method.") + parser.add_argument("--split_file", default=SPLIT_FILE, help="The split file the sweep was run on.") + parser.add_argument("--split_key", default="val", help="The key of the split file the sweep was run on.") + parser.add_argument("-o", "--output", help="Write the report to this markdown file.") + parser.add_argument("--best", help="Write the selected settings to this json file. " + "Default: best.json in the sweep root.") + parser.add_argument("--workers", type=int, default=8, + help="Parallel scorings; each takes ~6-8 GB at the native resolution.") + args = parser.parse_args() + + blocks = compare_models.find_split_blocks(args.split_file, args.split_key) + settings, folders = [], {} + for name in sorted(os.listdir(args.sweep_root)): + folder = os.path.join(args.sweep_root, name) + if not os.path.isdir(folder): + continue + setting, paths = read_setting(folder, blocks) + settings.append(setting) + folders[folder] = paths + if not settings: + raise ValueError(f"Did not find any sweep folders in {args.sweep_root}.") + models = sorted({setting["model"] for setting in settings}) + if len(models) != 1: + raise ValueError(f"A sweep root must hold one model, found {', '.join(models)} in {args.sweep_root}.") + model = models[0] + + metrics = score_sweep(folders, blocks, args.workers) + + per_block, averaged = [], [] + for setting in settings: + folder_metrics = metrics[os.path.join(args.sweep_root, setting["folder"])] + rows = [{"setting": setting["folder"], "dataset": dataset, + **{key: folder_metrics[dataset][key] for key in METRICS + ["n_pred", "n_true"]}} + for dataset in blocks] + per_block += rows + mean = {key: float(np.mean([row[key] for row in rows])) for key in METRICS} + counts = {key: int(sum(row[key] for row in rows)) for key in ("n_pred", "n_true")} + for approach in setting["approaches"]: + averaged.append({"approach": approach, "voxel_nm": voxel_label(approach, setting["voxel_size"]), + "contrast": setting["contrast"], "mode": setting["mode"], + "setting": setting["folder"], "voxel_size": setting["voxel_size"], + **mean, **counts}) + averaged = pd.DataFrame(averaged) + per_block = pd.DataFrame(per_block).sort_values(["setting", "dataset"]) + + table_columns = ["voxel_nm", "contrast", "mode", *METRICS, "n_pred", "n_true"] + report = [ + f"# Volume EM mitochondria: {model} voxel-size sweep on the validation blocks", + "", + f"{model} on the `{args.split_key}` blocks of `{os.path.basename(args.split_file)}` " + f"({', '.join(blocks)}), which it never saw.", + "Each block is resampled to the voxel size, segmented, and resized back to the native", + "25/5/5 nm grid, where it is scored against the untouched ground truth. `min_size` is scaled", + "to the same physical volume as synapse-net's 1,000 native voxels at every voxel size.", + "", + "Instance metrics are matching at an IoU of 0.5; `msa` averages over the thresholds 0.5 to 0.95.", + "Metrics are averaged over the blocks, `n_pred` and `n_true` are summed.", + "", + "The best setting of each approach is selected here, by F1 with ties broken by msa, and then run", + "once on the test split. Nothing is selected on the test blocks.", + "", + ] + best = {} + for approach, title in (("A", "A: z at 25 nm, xy resampled (xy stack)"), + ("B", "B: isotropic"), + ("C", "C: single settings (voxel size z/y/x in nm)")): + frame = averaged[averaged["approach"] == approach].sort_values(["contrast", "mode", "voxel_nm"]) + if frame.empty: + continue + row = select_best(frame) + works = row["f1"] >= F1_FLOOR + # C is a set of single settings, not a grid, so it has no edge to extend. + grid = sorted(frame["voxel_nm"].unique()) if approach != "C" else [] + on_edge = works and len(grid) > 1 and row["voxel_nm"] in (grid[0], grid[-1]) + best[approach] = { + "setting": row["setting"], "voxel_size": list(row["voxel_size"]), "contrast": row["contrast"], + "mode": row["mode"], "f1": row["f1"], "msa": row["msa"], "works": bool(works), + "on_grid_edge": bool(on_edge), + "segment_args": " ".join( + ["--voxel_size", *(f"{v:g}" for v in row["voxel_size"]), "--modes", row["mode"]] + + (["--invert"] if row["contrast"] == "inverted" else []) + ), + } + if works: + verdict = (f"**Selected: `{row['setting']}`** — {row['contrast']} contrast, {row['mode']} mode, " + f"F1 {row['f1']:.4f}, msa {row['msa']:.4f}.") + if on_edge: + verdict += (f" It is on the edge of the grid ({grid[0]:g}–{grid[-1]:g} nm): extend the grid " + "by a step before relying on it.") + else: + verdict = (f"**No setting works:** the best is `{row['setting']}` at F1 {row['f1']:.4f}, " + f"below {F1_FLOOR}, so none is selected.") + table = frame[table_columns].assign(voxel_nm=frame["voxel_nm"].map( + lambda v: v if isinstance(v, str) else f"{v:g}")) + report += [f"## {title}", "", verdict, "", compare_models._markdown_table(table), ""] + report += ["## Per block", "", compare_models._markdown_table(per_block), ""] + + report = "\n".join(report) + print("\n" + report) + if args.output: + with open(args.output, "w") as f: + f.write(report) + print(f"Wrote {args.output}") + best_path = args.best or os.path.join(args.sweep_root, "best.json") + with open(best_path, "w") as f: + json.dump(best, f, indent=2) + print(f"Wrote {best_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/evaluation/segment_microsam.py b/scripts/volume_em/evaluation/segment_microsam.py new file mode 100644 index 00000000..b23e430d --- /dev/null +++ b/scripts/volume_em/evaluation/segment_microsam.py @@ -0,0 +1,221 @@ +"""Segment volume EM blocks with micro-sam at a chosen voxel size. + +This is the micro-sam baseline for compare_models.py, the counterpart of segment_mitonet.py. Each +block is resampled to the target voxel size, segmented with micro-sam's automatic instance +segmentation (AIS: the instance segmentation decoder of an EM organelle model), and the result is +resized back to the native grid with nearest-neighbor interpolation. The ground truth is never +resampled; compare_models.py scores the native 'seg' against it with '-s', and +score_voxel_size_sweep.py scores a whole voxel-size sweep. + +micro-sam segments the xy slices in 2D and merges them along z; it has no ortho-plane mode, so there +is one mode, 'xy', written to '-xy'. The prefix names the setting like segment_mitonet.py +does, e.g. 'microsam-iso12.5' or 'microsam-z25-xy10'. + +The encoder scale is controlled, not only the resampling. SAM resizes every image, and every tile, +so that its longest side is 1024 px, up or down. Passed as-is, a slice of these 8 um blocks would +reach the encoder at ~7.8 nm/px at every target voxel size, only more or less blurred. So every +image or tile that holds data is made exactly 1024 px, and SAM runs at scale 1: + a slice of at most 1024 px is reflect-padded to 1024 x 1024 and segmented untiled + a larger slice is reflect-padded onto a canvas on which micro-sam's tiling, with + tiles of 768 and a halo of 128, gives every tile holding data its full + 1024 px outer block; micro-sam clips the halo at the canvas border, + so the data starts one tile in +The padding is cropped away after the segmentation. + +The size filter is matched to synapse-net in physical volume, as for MitoNet +(voxel_sweep.min_size_for). Everything else is micro-sam's: the AIS defaults, and the 3D merge with +gap_closing 2 and min_z_extent 2, the defaults of its napari annotator for automatic 3D segmentation. + +This runs in the micro-sam environment, which does not have synapse_net, so it does not import +compare_models.py; the block lookup and resampling come from voxel_sweep.py. + +Usage: + python segment_microsam.py [--voxel_size Z Y X] [--model_type vit_b_em_organelles] [--force] + python segment_microsam.py --split_file SPLIT.json --split_key val --segmentation_root ROOT ... +""" + +import argparse +import importlib.metadata +import math +import os +import sys + +import h5py +import numpy as np +from skimage.segmentation import relabel_sequential + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from voxel_sweep import ( # noqa: E402 + NATIVE_VOXEL_SIZE, TEST_SPLIT, find_split_blocks, find_test_blocks, min_size_for, n_objects, + setting_prefix, to_native_resolution, to_target_resolution, +) + +SEGMENTATION_ROOT = "/mnt/lustre-grete/usr/u15205/volume-em/repro-comparison" +MODEL_TYPE = "vit_b_em_organelles" + +# The input size of the SAM image encoder, and a tiling whose tiles with halo have exactly that size. +SAM_INPUT = 1024 +TILE, HALO = 768, 128 +assert TILE + 2 * HALO == SAM_INPUT + +# The defaults of InstanceSegmentationWithDecoder.generate, set explicitly so the settings recorded +# with the segmentation stay true if a newer micro-sam changes them. +AIS_PARAMS = dict( + center_distance_threshold=0.5, boundary_distance_threshold=0.5, foreground_threshold=0.5, + foreground_smoothing=1.0, distance_smoothing=1.6, min_size=0, +) +# The defaults of micro-sam's napari annotator for automatic 3D segmentation; the merge along z +# itself is a multicut with beta 0.5, fixed inside automatic_3d_segmentation. +GAP_CLOSING = 2 +MIN_Z_EXTENT = 2 + + +def _axis_layout(n, tiled): + """The padding before and after one axis of n px, so that SAM sees the data at scale 1.""" + if not tiled: + before = (SAM_INPUT - n) // 2 + return before, SAM_INPUT - n - before + # micro-sam's blocking starts at 0 and clips the halo at the border, so the first tile never has + # its full outer block. The data starts at TILE, which leaves that tile to the padding, and the + # canvas ends one halo after the last tile holding data. + n_data_tiles = math.ceil(n / TILE) + total = (n_data_tiles + 1) * TILE + HALO + return TILE, total - TILE - n + + +def to_sam_canvas(volume): + """Reflect-pad a volume in y and x so that every image or tile holding data is SAM_INPUT px. + + Returns the canvas, the slice that crops the data back out of it, and the tile shape and halo + for micro-sam, which are None for an untiled canvas. + """ + tiled = max(volume.shape[1:]) > SAM_INPUT + pads = [(0, 0)] + [_axis_layout(n, tiled) for n in volume.shape[1:]] + canvas = np.pad(volume, pads, mode="reflect") + crop = tuple(slice(before, before + n) for (before, _), n in zip(pads, volume.shape)) + if tiled: + return canvas, crop, (TILE, TILE), (HALO, HALO) + return canvas, crop, None, None + + +def filter_small_objects(seg, min_size): + """Drop the objects below a size and relabel the rest consecutively.""" + ids, counts = np.unique(seg, return_counts=True) + too_small = ids[(counts < min_size) & (ids != 0)] + if too_small.size: + seg = seg.copy() + seg[np.isin(seg, too_small)] = 0 + seg, _, _ = relabel_sequential(seg) + return seg.astype(np.uint32) + + +def segment_microsam(predictor, segmenter, volume, min_size): + """Segment a volume with micro-sam's AIS at scale 1 and return the segmentation of the data.""" + from micro_sam.multi_dimensional_segmentation import automatic_3d_segmentation + + canvas, crop, tile_shape, halo = to_sam_canvas(volume) + seg = automatic_3d_segmentation( + canvas, predictor, segmenter, gap_closing=GAP_CLOSING, min_z_extent=MIN_Z_EXTENT, + tile_shape=tile_shape, halo=halo, verbose=False, **AIS_PARAMS, + ) + return filter_small_objects(seg[crop], min_size), canvas.shape, tile_shape, halo + + +def settings(model_type, voxel_size, invert, min_size, inference_shape, native_shape, canvas_shape, + tile_shape, halo, n_lost, source): + """The settings a segmentation was made with, stored as the attributes of its 'seg'.""" + effective = [n * size / m for n, size, m in zip(native_shape, NATIVE_VOXEL_SIZE, inference_shape)] + return { + "model": model_type, + "mode": "xy slices merged in z", + "segmentation_mode": "ais", + "target_voxel_size_nm": np.array(voxel_size, dtype=float), + "inference_voxel_size_nm": np.round(effective, 2), + "invert": bool(invert), + # AIS has a 'min_size' of its own, per slice; the 3D size filter keeps the plain name. + **{("ais_min_size" if key == "min_size" else key): value for key, value in AIS_PARAMS.items()}, + "min_size": min_size, + "gap_closing": GAP_CLOSING, + "min_z_extent": MIN_Z_EXTENT, + "sam_input": SAM_INPUT, + "canvas_shape": np.array(canvas_shape), + "tile_shape": np.array(tile_shape) if tile_shape else "none", + "halo": np.array(halo) if halo else "none", + "n_lost_to_native": n_lost, + "blocks": source, + "micro_sam": importlib.metadata.version("micro_sam"), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("--voxel_size", type=float, nargs=3, default=[30.0, 30.0, 30.0], metavar=("Z", "Y", "X"), + help="The voxel size to segment at, in nm.") + parser.add_argument("--modes", nargs="+", choices=["xy"], default=["xy"], + help="Only 'xy'; accepted so the arguments of a sweep's best.json can be passed on.") + parser.add_argument("--invert", action="store_true", help="Invert the contrast, i.e. segment 255 - raw.") + parser.add_argument("--model_type", default=MODEL_TYPE, help="The micro-sam model, with an AIS decoder.") + parser.add_argument("--test_split", default=TEST_SPLIT, help="The folder with the held-out test blocks.") + parser.add_argument("--split_file", help="Segment the blocks of a split file instead of the test split.") + parser.add_argument("--split_key", default="val", help="The key of the split file to segment.") + parser.add_argument("--segmentation_root", default=SEGMENTATION_ROOT, + help="Where the segmentations are written, one folder per setting.") + parser.add_argument("--force", action="store_true", help="Recompute segmentations that already exist.") + args = parser.parse_args() + + import torch + from micro_sam.automatic_segmentation import get_predictor_and_segmenter + + if args.split_file: + blocks = find_split_blocks(args.split_file, args.split_key) + source = f"{args.split_file}:{args.split_key}" + else: + blocks = find_test_blocks(args.test_split) + source = args.test_split + voxel_size = tuple(args.voxel_size) + min_size = min_size_for(voxel_size) + prefix = setting_prefix("microsam", voxel_size, args.invert) + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"{prefix}: micro-sam {args.model_type} at {voxel_size} nm on {len(blocks)} block(s), " + f"device {device}, min_size {min_size}\n") + + segmenters = {} + for dataset, block_path in blocks.items(): + output_path = os.path.join(args.segmentation_root, f"{prefix}-xy", f"{dataset}.h5") + if not args.force and os.path.exists(output_path): + print(f" {dataset}: already segmented, skipping") + continue + + with h5py.File(block_path, "r") as f: + raw = f["raw"][:] + if args.invert: + raw = np.iinfo(raw.dtype).max - raw + volume = to_target_resolution(raw, voxel_size) + + tiled = max(volume.shape[1:]) > SAM_INPUT + if tiled not in segmenters: + segmenters[tiled] = get_predictor_and_segmenter( + args.model_type, segmentation_mode="ais", is_tiled=tiled, + ) + predictor, segmenter = segmenters[tiled] + seg, canvas_shape, tile_shape, halo = segment_microsam(predictor, segmenter, volume, min_size) + print(f" {dataset}: {raw.shape} at {NATIVE_VOXEL_SIZE} nm -> {volume.shape}, " + f"SAM canvas {canvas_shape}" + (f", tiles {tile_shape} + halo {halo}" if tile_shape else ""), + flush=True) + + native = to_native_resolution(seg, raw.shape) + n_lost = n_objects(seg) - n_objects(native) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + with h5py.File(output_path, "w") as f: + ds = f.create_dataset("seg", data=native, compression="gzip") + ds.attrs.update(settings(args.model_type, voxel_size, args.invert, min_size, seg.shape, raw.shape, + canvas_shape, tile_shape, halo, n_lost, source)) + f.create_dataset("seg_inference", data=seg, compression="gzip") + lost = f", {n_lost} lost on the native grid" if n_lost else "" + print(f" xy: {n_objects(seg)} objects{lost} -> {output_path}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/evaluation/segment_mitonet.py b/scripts/volume_em/evaluation/segment_mitonet.py new file mode 100644 index 00000000..9ccf6ef9 --- /dev/null +++ b/scripts/volume_em/evaluation/segment_mitonet.py @@ -0,0 +1,184 @@ +"""Segment volume EM blocks with MitoNet at a chosen voxel size. + +This is the MitoNet baseline for compare_models.py. MitoNet (empanada) is a 2D generalist, and its 3D +mode, the ortho-plane consensus, infers on the xy, xz and yz planes, so it assumes isotropic voxels. +The blocks are at 25 nm in z and 5 nm in xy. Each block is resampled to the target voxel size, +segmented there, and the result is resized back to the native grid with nearest-neighbor +interpolation. The ground truth is never resampled; compare_models.py scores the native 'seg' +against it with '-s', and score_mitonet_sweep.py scores a whole voxel-size sweep. + +Two modes can come out of the same run, and each is written to its own folder: + -ortho the ortho-plane consensus of the xy, xz and yz stacks; only meaningful isotropic + -xy the xy stack alone, the mode of the earlier MitoNet runs +The prefix names the setting: 'mitonet-iso30' for 30 nm isotropic, 'mitonet-z25-xy10' for 25 nm in z +and 10 nm in xy, with '-inv' appended when the contrast is inverted. + +The size filter is matched to synapse-net in physical volume rather than left at empanada's default, +see voxel_sweep.min_size_for. All other parameters are the Engine3d defaults. + +This runs in the empanada environment, which has neither synapse_net nor torch_em nor elf, so it does +not import compare_models.py or synapse_net; the block lookup and resampling come from voxel_sweep.py. + +Usage: + python segment_mitonet.py [--voxel_size Z Y X] [--modes ortho xy] [--invert] [--force] + python segment_mitonet.py --split_file SPLIT.json --split_key val --segmentation_root ROOT ... +""" + +import argparse +import importlib.metadata +import os +import sys + +import h5py +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from voxel_sweep import ( # noqa: E402 + NATIVE_VOXEL_SIZE, TEST_SPLIT, find_split_blocks, find_test_blocks, min_size_for, n_objects, + setting_prefix, to_native_resolution, to_target_resolution, +) + +SEGMENTATION_ROOT = "/mnt/lustre-grete/usr/u15205/volume-em/repro-comparison" +MIN_EXTENT = 4 + +# The Engine3d defaults, which the earlier MitoNet runs used too. They are set explicitly so the +# settings recorded with the segmentation stay true if a newer empanada changes its defaults. The +# napari widget has different defaults (confidence_thr 0.5, median_kernel_size 3, min_extent 5). +ENGINE_PARAMS = dict( + inference_scale=1, confidence_thr=0.3, median_kernel_size=5, + nms_threshold=0.1, nms_kernel=3, label_divisor=1000, +) +# The consensus defaults, which the napari widget uses as well. +CONSENSUS_PARAMS = dict(pixel_vote_thr=2, cluster_iou_thr=0.75, allow_one_view=False) + +# MitoNet_v1.yaml, without the FINETUNE block that inference does not read. +MITONET_CONFIG = { + "class_names": {1: "mito"}, + "labels": [1], + "thing_list": [1], + "model": "https://zenodo.org/record/6861565/files/MitoNet_v1.pth?download=1", + "model_quantized": "https://zenodo.org/record/6861565/files/MitoNet_v1_quantized.pth?download=1", + "padding_factor": 16, + "norms": {"mean": 0.57571, "std": 0.12765}, +} + +MODES = ("ortho", "xy") + + +def segment_mitonet(engine, volume, modes, min_size): + """Segment a volume with MitoNet and return a segmentation per mode.""" + from empanada_napari.inference import stack_postprocessing, tracker_consensus + + axes = ("xy", "xz", "yz") if "ortho" in modes else ("xy",) + trackers = {} + for axis_name in axes: + _, trackers[axis_name] = engine.infer_on_axis(volume, axis_name) + + # Both functions are napari thread workers; call the generators they wrap, which yield one + # (volume, class name, instances) per class. MitoNet has only the one class. + segmentations = {} + if "xy" in modes: + segmentations["xy"], _, _ = next(stack_postprocessing.__wrapped__( + {"xy": trackers["xy"]}, None, MITONET_CONFIG, label_divisor=engine.label_divisor, + min_size=min_size, min_extent=MIN_EXTENT, dtype=np.uint32, + )) + if "ortho" in modes: + segmentations["ortho"], _, _ = next(tracker_consensus.__wrapped__( + trackers, None, MITONET_CONFIG, label_divisor=engine.label_divisor, **CONSENSUS_PARAMS, + min_size=min_size, min_extent=MIN_EXTENT, dtype=np.uint32, + )) + return segmentations + + +def settings(mode, voxel_size, invert, min_size, inference_shape, native_shape, n_lost, source): + """The settings a segmentation was made with, stored as the attributes of its 'seg'.""" + effective = [n * size / m for n, size, m in zip(native_shape, NATIVE_VOXEL_SIZE, inference_shape)] + attrs = { + "model": "MitoNet_v1", + "mode": "ortho-plane consensus of xy, xz, yz" if mode == "ortho" else "xy stack", + "target_voxel_size_nm": np.array(voxel_size, dtype=float), + "inference_voxel_size_nm": np.round(effective, 2), + "invert": bool(invert), + "min_size": min_size, + "min_extent": MIN_EXTENT, + **{key: value for key, value in ENGINE_PARAMS.items() if key != "inference_scale"}, + } + if mode == "ortho": + attrs.update(CONSENSUS_PARAMS) + attrs["n_lost_to_native"] = n_lost + attrs["blocks"] = source + attrs["empanada_napari"] = importlib.metadata.version("empanada-napari") + return attrs + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("--voxel_size", type=float, nargs=3, default=[30.0, 30.0, 30.0], metavar=("Z", "Y", "X"), + help="The voxel size to segment at, in nm.") + parser.add_argument("--modes", nargs="+", choices=MODES, default=list(MODES), + help="The modes to write. 'xy' alone skips the xz and yz inference.") + parser.add_argument("--invert", action="store_true", help="Invert the contrast, i.e. segment 255 - raw.") + parser.add_argument("--test_split", default=TEST_SPLIT, help="The folder with the held-out test blocks.") + parser.add_argument("--split_file", help="Segment the blocks of a split file instead of the test split.") + parser.add_argument("--split_key", default="val", help="The key of the split file to segment.") + parser.add_argument("--segmentation_root", default=SEGMENTATION_ROOT, + help="Where the segmentations are written, one folder per setting and mode.") + parser.add_argument("--force", action="store_true", help="Recompute segmentations that already exist.") + args = parser.parse_args() + + from empanada_napari.inference import Engine3d + + if args.split_file: + blocks = find_split_blocks(args.split_file, args.split_key) + source = f"{args.split_file}:{args.split_key}" + else: + blocks = find_test_blocks(args.test_split) + source = args.test_split + voxel_size = tuple(args.voxel_size) + modes = [mode for mode in MODES if mode in args.modes] + if "ortho" in modes and len(set(voxel_size)) > 1: + print(f"Warning: ortho-plane inference at the anisotropic voxel size {voxel_size}.") + min_size = min_size_for(voxel_size) + prefix = setting_prefix("mitonet", voxel_size, args.invert) + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"{prefix}: MitoNet at {voxel_size} nm on {len(blocks)} block(s), modes {', '.join(modes)}, " + f"device {device}, min_size {min_size}, min_extent {MIN_EXTENT}\n") + + engine = None + for dataset, block_path in blocks.items(): + output_paths = {mode: os.path.join(args.segmentation_root, f"{prefix}-{mode}", f"{dataset}.h5") + for mode in modes} + if not args.force and all(os.path.exists(path) for path in output_paths.values()): + print(f" {dataset}: already segmented, skipping") + continue + + with h5py.File(block_path, "r") as f: + raw = f["raw"][:] + if args.invert: + raw = np.iinfo(raw.dtype).max - raw + volume = to_target_resolution(raw, voxel_size) + print(f" {dataset}: {raw.shape} at {NATIVE_VOXEL_SIZE} nm -> {volume.shape}", flush=True) + + if engine is None: + engine = Engine3d(MITONET_CONFIG, min_size=min_size, min_extent=MIN_EXTENT, **ENGINE_PARAMS) + segmentations = segment_mitonet(engine, volume, modes, min_size) + + for mode, seg in segmentations.items(): + native = to_native_resolution(seg, raw.shape) + n_lost = n_objects(seg) - n_objects(native) + output_path = output_paths[mode] + os.makedirs(os.path.dirname(output_path), exist_ok=True) + with h5py.File(output_path, "w") as f: + ds = f.create_dataset("seg", data=native, compression="gzip") + ds.attrs.update(settings(mode, voxel_size, args.invert, min_size, seg.shape, raw.shape, + n_lost, source)) + f.create_dataset("seg_inference", data=seg.astype(np.uint32), compression="gzip") + lost = f", {n_lost} lost on the native grid" if n_lost else "" + print(f" {mode}: {n_objects(seg)} objects{lost} -> {output_path}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/evaluation/voxel_sweep.py b/scripts/volume_em/evaluation/voxel_sweep.py new file mode 100644 index 00000000..7bd9fa12 --- /dev/null +++ b/scripts/volume_em/evaluation/voxel_sweep.py @@ -0,0 +1,129 @@ +"""Shared parts of the voxel-size sweeps of the baseline methods, segment_mitonet.py and segment_microsam.py. + +Block lookup, resampling to a target voxel size and back, the physically matched size filter, and the +names of the setting folders. It only needs numpy, h5py and skimage, so that the environments of the +baseline methods, which have neither synapse_net nor torch_em, can import it. +""" + +import json +import os + +import numpy as np +from skimage.transform import resize + +TEST_SPLIT = "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/test_split" + +# The voxel size of the blocks, in nm and ZYX. It is hardcoded because the 'voxel_size' attributes +# of the files cannot be trusted: 4007 stores 17.388 for every axis, 4009 stores +# [0.025, 0.005, 0.005], and the downsampled test_split_* copies disagree with both. +NATIVE_VOXEL_SIZE = (25.0, 5.0, 5.0) + +# The synapse-net post-processing drops objects below 1000 native voxels, see min_size_for. +MIN_SIZE_NATIVE = 1000 + + +def find_test_blocks(test_split: str): + """Find the held-out test blocks, as a mapping of dataset name to filepath. + + The same lookup as compare_models.find_test_blocks, which cannot be imported in these environments. + """ + blocks = {} + for dataset in sorted(os.listdir(test_split)): + folder = os.path.join(test_split, dataset) + if not os.path.isdir(folder): + continue + files = sorted(name for name in os.listdir(folder) if name.endswith(".h5")) + if len(files) != 1: + raise ValueError(f"Expect exactly one test block in {folder}, found {len(files)}.") + blocks[dataset] = os.path.join(folder, files[0]) + if not blocks: + raise ValueError(f"Did not find any test blocks in {test_split}.") + return blocks + + +def find_split_blocks(split_file: str, key: str): + """Find the blocks of one key of a split file, as a mapping of data root name to filepath. + + Resolves the '/' entries like synapse_net.training.split._resolve_split, + which cannot be imported in these environments, with the roots the split file carries. + """ + with open(split_file) as f: + split = json.load(f) + roots = split.get("roots", {}) + if not isinstance(roots, dict): + roots = {os.path.basename(root.rstrip("/")): root for root in roots} + roots = {name: root.rstrip("/") for name, root in roots.items()} + + blocks = {} + for entry in split.get(key, []): + name, _, relative_path = entry.partition("/") + if name not in roots: + raise ValueError(f"The split file {split_file} refers to the unknown data root '{name}'.") + if name in blocks: + raise ValueError(f"The '{key}' split has more than one block from '{name}', so they cannot " + "be named by their data root.") + blocks[name] = os.path.join(roots[name], *relative_path.split("/")) + if not blocks: + raise ValueError(f"The split file {split_file} has no '{key}' entries.") + return blocks + + +def min_size_for(voxel_size): + """The size filter at a voxel size, as the same physical volume as synapse-net's. + + synapse-net drops objects below 1000 native voxels. At 30 nm isotropic that is 23 voxels, while + empanada's default of 500 voxels would be ~21,600 native voxels, which would drop 18% of the + mitochondria annotated in the training blocks (see the README) and would flatter MitoNet on + blocks with large mitochondria. Scaling it also keeps a voxel-size sweep from changing the + filter strength along with the voxel size. + """ + return round(MIN_SIZE_NATIVE * float(np.prod(NATIVE_VOXEL_SIZE)) / float(np.prod(voxel_size))) + + +def setting_prefix(method, voxel_size, invert): + """The folder prefix that names a setting, e.g. 'mitonet-iso30' or 'microsam-z25-xy10'.""" + z, y, x = (f"{v:g}" for v in voxel_size) + if z == y == x: + prefix = f"{method}-iso{z}" + elif y == x: + prefix = f"{method}-z{z}-xy{y}" + else: + prefix = f"{method}-z{z}-y{y}-x{x}" + return prefix + ("-inv" if invert else "") + + +def to_target_resolution(raw, voxel_size): + """Resample the raw data from the native voxel size to the target voxel size. + + An axis whose shape does not change is left untouched, and a block whose shape does not change + at all is returned as it is. + """ + if raw.dtype != np.uint8: + # empanada normalizes by the maximum of the integer dtype and refuses floats. + raise ValueError(f"Expect uint8 raw data, got {raw.dtype}.") + shape = tuple(int(round(n * size / target)) for n, size, target in zip(raw.shape, NATIVE_VOXEL_SIZE, voxel_size)) + if shape == raw.shape: + return raw + # skimage only smooths axes that are downsampled, and interpolating an axis at its own grid + # points is exact, so the axes that keep their shape come out unchanged. + resampled = resize(raw, shape, order=1, anti_aliasing=True, preserve_range=True) + return np.clip(np.round(resampled), 0, 255).astype(np.uint8) + + +def to_native_resolution(seg, shape): + """Resize a segmentation to the native grid with nearest-neighbor interpolation. + + Every output voxel takes the input voxel its center falls into. Along an upsampled axis every + input voxel is hit, so nothing is lost; along a downsampled one, such as z for isotropic + inference below 25 nm, an object thinner than a native slice can be. + """ + index = [ + np.minimum(((np.arange(n_out) + 0.5) * n_in / n_out).astype(int), n_in - 1) + for n_in, n_out in zip(seg.shape, shape) + ] + return seg[np.ix_(*index)].astype(np.uint32) + + +def n_objects(seg): + ids = np.unique(seg) + return int(len(ids) - (1 if 0 in ids else 0)) diff --git a/scripts/volume_em/inference/run_mitochondria_vol_em_segmentation.py b/scripts/volume_em/inference/run_mitochondria_vol_em_segmentation.py new file mode 100644 index 00000000..13d6e75d --- /dev/null +++ b/scripts/volume_em/inference/run_mitochondria_vol_em_segmentation.py @@ -0,0 +1,167 @@ +"""Segment mitochondria in volume EM data with SynapseNet's volume EM mitochondria model. + +The model is not part of the model registry, so its checkpoint has to be passed explicitly via '-m'. +See `synapse_net.training.mitochondria_vol_em` for how it was trained. + +The preprocessing is part of the model and is reproduced here in two steps: + +- The white filler borders are removed from the **whole volume** before prediction. This cannot be + done in the per-block 'preprocess' of 'segment_mitochondria', because 'get_prediction' standardizes + a numpy input volume before it runs that, and the filler is identified by its literal value of 255, + which no longer exists after standardizing. +- The blocks are then percentile-normalized via 'preprocess', exactly as during training. The + percentile normalization is invariant under the standardization that 'get_prediction' applies, so + the network sees the same values either way. + +Pass '--no_white_patch_fix' to skip the filler removal, which is only correct for data that was not +cut out of a larger volume. On a block that is 14% filler, leaving it in cost 0.04 Dice (0.79 instead +of 0.84), because the filler skews the percentile normalization of every tile it overlaps. + +The data is segmented at its native resolution by default, because the model is meant to be applied to +data at its own training resolution of 25 nm in z and 5 nm in xy. Pass '--scale' for data at a +different voxel size; it cannot be derived automatically, because the voxel size is only read for mrc +files and the volume EM data is stored as hdf5. +""" + +import argparse +from functools import partial + +import torch_em + +from synapse_net.inference.mitochondria import segment_mitochondria +from synapse_net.inference.util import inference_helper, parse_tiling +from synapse_net.training.transform import remove_white_patches + +# The post-processing defaults were found by a grid search on the 4007 dataset. Note that the grid +# search used a different (out-of-core) watershed implementation, so these are a starting point for +# the watershed of 'segment_mitochondria' rather than a tuned optimum for it. +SEED_DISTANCE = 1 +BOUNDARY_THRESHOLD = 0.12 +AREA_THRESHOLD = 200 +MIN_SIZE = 1000 + + +def segment_mitochondria_vol_em(input_volume, white_patch_min_size=20, **kwargs): + """Remove the white filler borders of a volume EM block and then segment its mitochondria. + + Args: + input_volume: The volume to segment, before any normalization. + white_patch_min_size: The minimal size (in voxels) of a filler component that is removed. + Pass None to segment without removing the filler. + kwargs: Additional keyword arguments for `synapse_net.inference.mitochondria.segment_mitochondria`. + + Returns: + The mitochondria segmentation. + """ + if white_patch_min_size is not None: + # This has to happen before the rescaling inside of 'segment_mitochondria', which would + # interpolate the filler away from its literal value. + input_volume = remove_white_patches(input_volume, min_size=white_patch_min_size) + return segment_mitochondria(input_volume, **kwargs) + + +def run_mitochondria_segmentation(args): + tiling = parse_tiling(args.tile_shape, args.halo) + segmentation_function = partial( + segment_mitochondria_vol_em, + white_patch_min_size=None if args.no_white_patch_fix else args.white_patch_min_size, + model_path=args.model, + tiling=tiling, + preprocess=torch_em.transform.raw.normalize_percentile, + min_size=args.min_size, + seed_distance=args.seed_distance, + boundary_threshold=args.boundary_threshold, + area_threshold=args.area_threshold, + ws_block_shape=tuple(args.ws_block_shape), + ws_halo=tuple(args.ws_halo), + verbose=args.verbose, + ) + inference_helper( + args.input_path, args.output_path, segmentation_function, + mask_input_path=args.mask_path, force=args.force, data_ext=args.data_ext, + output_key=args.segmentation_key, scale=args.scale, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument( + "--input_path", "-i", required=True, + help="The filepath to the hdf5 file or the directory containing the volume EM data." + ) + parser.add_argument( + "--output_path", "-o", required=True, + help="The filepath to the directory where the segmentation will be saved." + ) + parser.add_argument( + "--model", "-m", required=True, + help="The filepath to the volume EM mitochondria model. This model is not part of the model registry." + ) + parser.add_argument( + "--segmentation_key", "-s", default="seg", + help="The key the segmentation is saved under. Pass an empty string to save tif files instead." + ) + parser.add_argument( + "--data_ext", default=".h5", help="The extension of the volume EM data. By default .h5." + ) + parser.add_argument( + "--mask_path", help="The filepath to a tif file with a mask that will be used to restrict the segmentation." + "Can also be a directory with tifs if the filestructure matches input_path." + ) + parser.add_argument( + "--force", action="store_true", help="Whether to over-write already present segmentation results." + ) + parser.add_argument( + "--scale", type=float, nargs=3, + help="The factor for rescaling the data before prediction, in ZYX. By default the data is segmented at its " + "native resolution, which is correct for data at the training voxel size of 25 x 5 x 5 nm." + ) + parser.add_argument( + "--no_white_patch_fix", action="store_true", + help="Do not remove the white filler borders before normalizing. Only correct for data that was not cut out " + "of a larger volume." + ) + parser.add_argument( + "--white_patch_min_size", type=int, default=20, + help="The minimal size (in voxels) of a white filler component that is removed." + ) + parser.add_argument( + "--tile_shape", type=int, nargs=3, help="The tile shape for prediction, in ZYX. Lower it if GPU memory is " + "insufficient." + ) + parser.add_argument( + "--halo", type=int, nargs=3, help="The halo for prediction, in ZYX. Increase it to minimize boundary artifacts." + ) + parser.add_argument( + "--min_size", type=int, default=MIN_SIZE, help="The minimum size of a mitochondrion, in voxels." + ) + parser.add_argument( + "--seed_distance", type=int, default=SEED_DISTANCE, help="The distance threshold for the seeded watershed." + ) + parser.add_argument( + "--boundary_threshold", type=float, default=BOUNDARY_THRESHOLD, + help="The boundary threshold for the distance calculation." + ) + parser.add_argument( + "--area_threshold", type=int, default=AREA_THRESHOLD, + help="The maximum area (in pixels) of holes that are filled in the segmentation." + ) + parser.add_argument( + "--ws_block_shape", type=int, nargs=3, default=[128, 256, 256], + help="The block shape for the seeded watershed, in ZYX." + ) + parser.add_argument( + "--ws_halo", type=int, nargs=3, default=[48, 48, 48], help="The halo for the seeded watershed, in ZYX." + ) + parser.add_argument( + "--verbose", "-v", action="store_true", help="Whether to print information about the segmentation progress." + ) + + args = parser.parse_args() + if not args.segmentation_key: + args.segmentation_key = None + run_mitochondria_segmentation(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/volume_em/training/compare_training_setup.py b/scripts/volume_em/training/compare_training_setup.py new file mode 100644 index 00000000..fcc57dae --- /dev/null +++ b/scripts/volume_em/training/compare_training_setup.py @@ -0,0 +1,315 @@ +"""Check that the ported volume EM mitochondria training builds the same trainer as the published model. + +This constructs the trainer through `synapse_net.training.mitochondria_vol_em.vol_em_mitochondria_training`, +stops it just before it would start training, and diffs its serialized setup against the setup stored +in the published checkpoint. It is the cheap part of reproducing that model: if this passes, the only +thing a training run can still tell us is whether the optimization lands in the same place. + +Three differences are expected and are checked rather than ignored: + +- The raw transform. The published run pickled a function that was local to its training script, which + is why its checkpoint cannot be loaded without a stand-in. The port uses + `synapse_net.training.transform.RemoveWhitePatchesAndNormalize`. They are compared numerically, on + real patches drawn from the training data. +- The joint transform. The port wraps the augmentations in `PadIfNecessary`, which the published run + did not. This is checked to be a no-op for this data, i.e. that no block is smaller than the patch. +- 'name', 'id_', 'save_root', 'device' and 'rank', which differ by construction. + +Usage: + python compare_training_setup.py [-c ] [--split_file ] +Exits non-zero if anything else differs. +""" + +import argparse +import os +import sys +import warnings + +import numpy as np +import torch +import torch_em +from skimage.measure import label as connected_components + +from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths, vol_em_mitochondria_training +from synapse_net.training.transform import RemoveWhitePatchesAndNormalize + +REFERENCE_CHECKPOINT = ( + "/mnt/lustre-grete/usr/u15205/volume-em/models/checkpoints/" + "volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final" +) +SPLIT_FILE = os.path.join(os.path.dirname(__file__), "split-mito_vol_em_aniso2lvl_final.json") + +# These differ by construction and carry no information about the recipe. +IGNORED_FIELDS = ("name", "id_", "save_root", "device", "rank", "logger_kwargs") +# Fields that a newer torch-em adds to the trainer, together with why they cannot change this recipe. +# A field that a newer torch-em adds and that is not listed here is reported as a difference, because +# it may well change what is trained. +TOLERATED_ADDITIONS = { + "mixed_precision_dtype": + "only read by the mixed-precision train and validate loops, and this recipe sets " + "mixed_precision to False", +} +# These hold live objects and are compared structurally further down instead. +STRUCTURAL_FIELDS = ("train_dataset", "val_dataset") + + +def reference_raw_transform(x): + """The raw transform of the published run, transcribed from 5aa482e:training/train_mito_generic.py. + + It is reimplemented here rather than imported, because importing it from the other repository pulls + in napari, mrcfile and z5py. Keeping it standalone is also what makes this an independent check. + """ + img = x + if img.dtype != np.uint8: + if not (np.issubdtype(img.dtype, np.floating) and img.min() >= 0 and img.max() <= 255): + warnings.warn("img must be uint8, converting to uint8 from " + str(img.dtype)) + img = img.astype(np.uint8) + if img.ndim not in (2, 3): + raise ValueError(f"img must be a 2D or 3D array, but got {img.ndim}D") + + white_mask = img == 255 + labeled = connected_components(white_mask) + if labeled.max() == 0: + out = img + else: + sizes = np.bincount(labeled.ravel()) + large_labels = np.where(sizes >= 20)[0] + large_labels = large_labels[large_labels != 0] + if large_labels.size == 0: + out = img + else: + out = img.copy() + out[np.isin(labeled, large_labels)] = 0 + return torch_em.transform.raw.normalize_percentile(out) + + +def load_reference_init(checkpoint): + """Read the setup out of the published checkpoint.""" + # The checkpoint pickled a raw transform that lived in '__main__' of the training script, so + # unpickling it needs a stand-in under that name. + sys.modules["__main__"]._raw_transform_with_fix_white_patches = reference_raw_transform + state = torch.load(os.path.join(checkpoint, "best.pt"), map_location="cpu", weights_only=False) + return state["init"] + + +def build_our_init(train_paths, val_paths): + """Build the trainer through the ported recipe and serialize it, without training.""" + captured = {} + build_trainer = torch_em.default_segmentation_trainer + + def capture(**kwargs): + trainer = build_trainer(**kwargs) + captured["trainer"] = trainer + trainer.fit = lambda *args, **kwargs: None + return trainer + + torch_em.default_segmentation_trainer = capture + try: + vol_em_mitochondria_training( + name="compare-training-setup", train_paths=train_paths, val_paths=val_paths, + save_root="/tmp/compare-training-setup", + ) + finally: + torch_em.default_segmentation_trainer = build_trainer + + trainer = captured["trainer"] + # These two lines are what 'DefaultTrainer._initialize' does before it serializes the setup. + trainer._model_class = f"{type(trainer.model).__module__}.{type(trainer.model).__name__}" + trainer._model_kwargs = torch_em.util.get_constructor_arguments(trainer.model) + return trainer._build_init() + + +def _describe_dataset(dataset): + """Reduce a ConcatDataset to the properties that define the training data.""" + subsets = dataset.datasets + first = subsets[0] + return { + "n_datasets": len(subsets), + "n_samples": len(dataset), + "lengths": [len(ds) for ds in subsets], + "files": sorted(os.path.basename(ds.raw_path) for ds in subsets), + "patch_shape": list(first.patch_shape), + "raw_key": first.raw_key, + "label_key": first.label_key, + "ndim": first.ndim, + "dtype": str(first.dtype), + "label_dtype": str(first.label_dtype), + "sampler": (type(first.sampler).__name__, vars(first.sampler)), + "label_transform": type(first.label_transform).__name__, + "label_transform2": type(first.label_transform2).__name__, + "augmentations": [type(aug).__name__ for aug in _augmentations_of(first.transform)], + } + + +def _augmentations_of(transform): + """Get the kornia augmentations of a joint transform, looking through a Compose.""" + if hasattr(transform, "augmentations"): + return transform.augmentations + # The port wraps the augmentations in a Compose together with PadIfNecessary. + for step in getattr(transform, "transforms", []): + if hasattr(step, "augmentations"): + return step.augmentations + return [] + + +def _pads_of(transform): + """Get the PadIfNecessary steps of a joint transform, which the published run did not have.""" + return [step for step in getattr(transform, "transforms", []) if type(step).__name__ == "PadIfNecessary"] + + +def compare_fields(reference, ours): + """Compare the flat fields of the two setups. + + Returns the unexpected differences, and separately the fields that only the version of torch-em + running here produces, which are tolerated only if they are known to be inert for this recipe. + """ + differences, additions = [], [] + fields = sorted(set(reference) | set(ours)) + print(f"{'field':<24} {'match':<7} value") + print("-" * 100) + for field in fields: + if field in IGNORED_FIELDS or field in STRUCTURAL_FIELDS: + continue + + # A field that the published checkpoint does not have at all was added by a newer torch-em, + # rather than being a difference in how the recipe is built. + if field not in reference: + if field in TOLERATED_ADDITIONS: + additions.append((field, ours[field], TOLERATED_ADDITIONS[field])) + print(f"{field:<24} {'added':<7} {ours[field]} (not in the published setup)") + else: + differences.append((field, "", ours[field])) + print(f"{field:<24} {'NEW':<7} {ours[field]} (not in the published setup)") + continue + + ref_value, our_value = reference[field], ours.get(field, "") + match = ref_value == our_value + if not match: + differences.append((field, ref_value, our_value)) + print(f"{field:<24} {'ok' if match else 'DIFFERS':<7} {ref_value}") + if not match: + print(f"{'':<24} {'':<7} ours: {our_value}") + return differences, additions + + +def compare_datasets(reference, ours): + """Compare the training and validation data of the two setups.""" + differences = [] + for field in STRUCTURAL_FIELDS: + ref_desc = _describe_dataset(reference[field]) + our_desc = _describe_dataset(ours[field]) + print(f"\n{field}") + print("-" * 100) + for key in ref_desc: + match = ref_desc[key] == our_desc[key] + if not match: + differences.append((f"{field}.{key}", ref_desc[key], our_desc[key])) + shown = ref_desc[key] if key != "files" else f"{len(ref_desc[key])} files" + print(f" {key:<20} {'ok' if match else 'DIFFERS':<7} {shown}") + if not match: + print(f" {'':<20} {'':<7} ours: {our_desc[key]}") + return differences + + +def check_raw_transform(train_paths, n_patches=24, patch_shape=(32, 512, 512), seed=0): + """Check that the ported raw transform agrees with the published one, on real data.""" + import h5py + + rng = np.random.default_rng(seed) + ours = RemoveWhitePatchesAndNormalize() + max_difference, n_with_filler = 0.0, 0 + + for i in range(n_patches): + path = train_paths[i % len(train_paths)] + with h5py.File(path, "r") as f: + raw = f["raw"] + # Draw patches from the volume border half of the time, where the filler lives. + offsets = [ + 0 if (i % 2 == 0 and axis > 0) else int(rng.integers(0, max(size - extent, 1))) + for axis, (size, extent) in enumerate(zip(raw.shape, patch_shape)) + ] + bb = tuple(slice(off, off + extent) for off, extent in zip(offsets, patch_shape)) + patch = raw[bb] + + if (patch == 255).sum() >= 20: + n_with_filler += 1 + difference = np.abs(ours(patch) - reference_raw_transform(patch)).max() + max_difference = max(max_difference, float(difference)) + + print(f"\nraw transform: {n_patches} real patches, {n_with_filler} of them containing filler") + print(f" maximal absolute difference to the published transform: {max_difference:g}") + return max_difference, n_with_filler + + +def check_padding_is_a_noop(ours, patch_shape=(32, 512, 512)): + """Check that the PadIfNecessary the port adds never fires, so the joint transforms agree.""" + import h5py + + too_small = [] + for dataset in list(ours["train_dataset"].datasets) + list(ours["val_dataset"].datasets): + with h5py.File(dataset.raw_path, "r") as f: + shape = f[dataset.raw_key].shape + if any(size < extent for size, extent in zip(shape, patch_shape)): + too_small.append((os.path.basename(dataset.raw_path), shape)) + n_pads = len(_pads_of(ours["train_dataset"].datasets[0].transform)) + print(f"\njoint transform: the port adds {n_pads} PadIfNecessary step(s)") + print(f" blocks smaller than {tuple(patch_shape)} in any axis: {len(too_small)}") + for name, shape in too_small: + print(f" {name} {shape}") + return too_small + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("-c", "--checkpoint", default=REFERENCE_CHECKPOINT, + help="The checkpoint folder of the published model.") + parser.add_argument("--split_file", default=SPLIT_FILE, help="The split file of the published run.") + args = parser.parse_args() + + train_paths, val_paths = get_vol_em_mitochondria_paths(split_file=args.split_file) + print(f"Comparing against {args.checkpoint}") + print(f"{len(train_paths)} training and {len(val_paths)} validation blocks from {args.split_file}\n") + + reference = load_reference_init(args.checkpoint) + ours = build_our_init(train_paths, val_paths) + + differences, additions = compare_fields(reference, ours) + differences += compare_datasets(reference, ours) + + max_difference, n_with_filler = check_raw_transform(train_paths) + too_small = check_padding_is_a_noop(ours) + + print("\n" + "=" * 100) + failed = False + if differences: + failed = True + print(f"{len(differences)} unexpected difference(s) in the training setup:") + for field, ref_value, our_value in differences: + print(f" {field}: published {ref_value!r} vs ours {our_value!r}") + else: + print("The training setup matches the published model in every compared field.") + + for field, value, reason in additions: + print(f"This torch-em adds '{field}' = {value!r}, which the published run did not have. " + f"It is inert here: {reason}.") + + if max_difference > 1e-6: + failed = True + print(f"The raw transform does not agree with the published one ({max_difference:g} > 1e-6).") + elif n_with_filler == 0: + failed = True + print("No patch contained filler, so the raw transform check did not exercise the filler removal.") + else: + print("The raw transform agrees with the published one on real data, including on filler.") + + if too_small: + failed = True + print(f"{len(too_small)} block(s) are smaller than the patch, so the added padding is not a no-op.") + else: + print("The padding the port adds to the joint transform never fires for this data.") + + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/training/repair_checkpoint.py b/scripts/volume_em/training/repair_checkpoint.py new file mode 100644 index 00000000..180bae02 --- /dev/null +++ b/scripts/volume_em/training/repair_checkpoint.py @@ -0,0 +1,97 @@ +"""Make a volume EM mitochondria checkpoint loadable again. + +torch-em pickles the whole training dataset into a checkpoint, including its raw transform. The +published volume EM mitochondria model was trained by a script that defined that transform as a +local function, so the checkpoint references `__main__._raw_transform_with_fix_white_patches` and +cannot be loaded anywhere else: + + AttributeError: Can't get attribute '_raw_transform_with_fix_white_patches' on + +This rewrites the reference to `synapse_net.training.transform.RemoveWhitePatchesAndNormalize`, which +is the same transform as an importable class, and writes the result to a new folder. The weights, the +optimizer state and every hyperparameter are copied through unchanged; only the dangling reference is +replaced. The original is never modified. + +Models trained with `synapse_net.training.mitochondria_vol_em` do not need this: their raw transform +is a module-level class, so their checkpoints load as they are. + +Usage: + python repair_checkpoint.py -i -o +""" + +import argparse +import os +import shutil +import sys + +import torch + +from synapse_net.training.transform import RemoveWhitePatchesAndNormalize + + +def repair_checkpoint(input_dir: str, output_dir: str, min_size: int = 20) -> None: + """Rewrite the dangling raw transform reference of a checkpoint. + + Args: + input_dir: The checkpoint folder to repair. It is not modified. + output_dir: The folder to write the repaired checkpoint to. + min_size: The minimal size of a filler component, for the transform that replaces the + dangling reference. The published model used the default of 20. + + Raises: + FileNotFoundError: If `input_dir` holds no checkpoint. + """ + checkpoints = [name for name in ("best.pt", "latest.pt") if os.path.exists(os.path.join(input_dir, name))] + if not checkpoints: + raise FileNotFoundError(f"Did not find best.pt or latest.pt in {input_dir}.") + + # Give the unpickler something to bind the dangling reference to. The published transform removed + # the white filler and then normalized with the 1st and 99th percentile, which is exactly what + # 'RemoveWhitePatchesAndNormalize' does; this was checked numerically in compare_training_setup.py. + replacement = RemoveWhitePatchesAndNormalize(min_size=min_size) + sys.modules["__main__"]._raw_transform_with_fix_white_patches = replacement + + os.makedirs(output_dir, exist_ok=True) + for name in checkpoints: + state = torch.load(os.path.join(input_dir, name), map_location="cpu", weights_only=False) + n_patched = 0 + for key in ("train_dataset", "val_dataset"): + dataset = state["init"].get(key) + if dataset is None: + continue + for sub_dataset in getattr(dataset, "datasets", [dataset]): + sub_dataset.raw_transform = RemoveWhitePatchesAndNormalize(min_size=min_size) + n_patched += 1 + torch.save(state, os.path.join(output_dir, name)) + print(f"{name}: iteration {state['iteration']}, best metric {state['best_metric']:.6f}, " + f"{n_patched} dataset(s) patched") + + for name in os.listdir(input_dir): + if not name.endswith(".pt"): + shutil.copy2(os.path.join(input_dir, name), output_dir) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("-i", "--input", required=True, help="The checkpoint folder to repair.") + parser.add_argument("-o", "--output", required=True, help="Where to write the repaired checkpoint.") + parser.add_argument("--white_patch_min_size", type=int, default=20, + help="The minimal size of a filler component for the replacement transform.") + args = parser.parse_args() + + repair_checkpoint(args.input, args.output, min_size=args.white_patch_min_size) + + # Prove that the result is loadable without any stand-in. + import subprocess + check = subprocess.run( + [sys.executable, "-c", + f"import torch_em; m = torch_em.util.load_model(checkpoint={args.output!r}, device='cpu'); " + "print('the repaired checkpoint loads:', type(m).__name__)"], + capture_output=True, text=True, + ) + print(check.stdout.strip() or check.stderr.strip()) + return check.returncode + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/volume_em/training/repro_aniso2lvl_seed42.yaml b/scripts/volume_em/training/repro_aniso2lvl_seed42.yaml new file mode 100644 index 00000000..2ca9eacc --- /dev/null +++ b/scripts/volume_em/training/repro_aniso2lvl_seed42.yaml @@ -0,0 +1,29 @@ +# Reproduction of the volume EM mitochondria model with the synapse-net training module. +# Baseline: volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final (val DiceLoss 0.292714). +# One clean run per seed; the recipe stops on early stopping, not on n_iterations. +# The save_root is deliberately separate from the original models/ directory: the published +# checkpoint was corrupted by two concurrent jobs sharing one, and synapse-net now refuses that. +# +# Submit with the runner of the synapse repo, which resolves 'slurm_profile' against its own +# slurm_profiles/ and writes the generated job script to its logs/sbatch_jobs/: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/training/repro_aniso2lvl_seed42.yaml [--dry-run] + +slurm_profile: grete_shared_80gb + +slurm: + job_name: "volem-mito-repro-seed42" + # The profile caps at 36 h. Historical runs early-stopped after 26-29 h, so raise this to the + # partition maximum to keep a margin at the measured 5.13 s/it. + time: "2-00:00:00" + +env: + type: micromamba + path: /mnt/vast-nhr/home/freckmann15/u15205/micromamba/envs/synapse + +commands: + - "CUDA_VISIBLE_DEVICES=0 synapse_net.run_vol_em_mitochondria_training + -n volume-em-mito-aniso2lvl-repro-seed42 + --split_file /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json + --seed 42 + --save_root /mnt/lustre-grete/usr/u15205/volume-em/models-repro" diff --git a/scripts/volume_em/training/repro_aniso2lvl_seed43.yaml b/scripts/volume_em/training/repro_aniso2lvl_seed43.yaml new file mode 100644 index 00000000..df678d2f --- /dev/null +++ b/scripts/volume_em/training/repro_aniso2lvl_seed43.yaml @@ -0,0 +1,29 @@ +# Reproduction of the volume EM mitochondria model with the synapse-net training module. +# Baseline: volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final (val DiceLoss 0.292714). +# One clean run per seed; the recipe stops on early stopping, not on n_iterations. +# The save_root is deliberately separate from the original models/ directory: the published +# checkpoint was corrupted by two concurrent jobs sharing one, and synapse-net now refuses that. +# +# Submit with the runner of the synapse repo, which resolves 'slurm_profile' against its own +# slurm_profiles/ and writes the generated job script to its logs/sbatch_jobs/: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/training/repro_aniso2lvl_seed43.yaml [--dry-run] + +slurm_profile: grete_shared_80gb + +slurm: + job_name: "volem-mito-repro-seed43" + # The profile caps at 36 h. Historical runs early-stopped after 26-29 h, so raise this to the + # partition maximum to keep a margin at the measured 5.13 s/it. + time: "2-00:00:00" + +env: + type: micromamba + path: /mnt/vast-nhr/home/freckmann15/u15205/micromamba/envs/synapse + +commands: + - "CUDA_VISIBLE_DEVICES=0 synapse_net.run_vol_em_mitochondria_training + -n volume-em-mito-aniso2lvl-repro-seed43 + --split_file /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json + --seed 43 + --save_root /mnt/lustre-grete/usr/u15205/volume-em/models-repro" diff --git a/scripts/volume_em/training/repro_aniso2lvl_seed44.yaml b/scripts/volume_em/training/repro_aniso2lvl_seed44.yaml new file mode 100644 index 00000000..a7de4335 --- /dev/null +++ b/scripts/volume_em/training/repro_aniso2lvl_seed44.yaml @@ -0,0 +1,29 @@ +# Reproduction of the volume EM mitochondria model with the synapse-net training module. +# Baseline: volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final (val DiceLoss 0.292714). +# One clean run per seed; the recipe stops on early stopping, not on n_iterations. +# The save_root is deliberately separate from the original models/ directory: the published +# checkpoint was corrupted by two concurrent jobs sharing one, and synapse-net now refuses that. +# +# Submit with the runner of the synapse repo, which resolves 'slurm_profile' against its own +# slurm_profiles/ and writes the generated job script to its logs/sbatch_jobs/: +# python /mnt/vast-nhr/home/freckmann15/u15205/synapse/sbatch_runner.py \ +# scripts/volume_em/training/repro_aniso2lvl_seed44.yaml [--dry-run] + +slurm_profile: grete_shared_80gb + +slurm: + job_name: "volem-mito-repro-seed44" + # The profile caps at 36 h. Historical runs early-stopped after 26-29 h, so raise this to the + # partition maximum to keep a margin at the measured 5.13 s/it. + time: "2-00:00:00" + +env: + type: micromamba + path: /mnt/vast-nhr/home/freckmann15/u15205/micromamba/envs/synapse + +commands: + - "CUDA_VISIBLE_DEVICES=0 synapse_net.run_vol_em_mitochondria_training + -n volume-em-mito-aniso2lvl-repro-seed44 + --split_file /mnt/vast-nhr/home/freckmann15/u15205/synapse-net/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json + --seed 44 + --save_root /mnt/lustre-grete/usr/u15205/volume-em/models-repro" diff --git a/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json b/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json new file mode 100644 index 00000000..78e272fd --- /dev/null +++ b/scripts/volume_em/training/split-mito_vol_em_aniso2lvl_final.json @@ -0,0 +1,26 @@ +{ + "roots": { + "4007": "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/4007_hdf5/all_cutouts_s0", + "4009": "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/4009_hdf5/cutouts_segmented" + }, + "train": [ + "4007/cutout_1_hdf5/cutout_1.h5", + "4007/cutout_2_hdf5/cutout_2.h5", + "4007/cutouts_h5_final/4007_block_z000256_000384_y001936_003872_x003312_004968.h5", + "4007/cutouts_h5_final/4007_block_z000384_000512_y001936_003872_x003312_004968.h5", + "4007/cutouts_h5_final/4007_block_z000512_000640_y001936_003872_x003312_004968.h5", + "4007/cutouts_h5_final/4007_block_z000640_000768_y001936_003872_x003312_004968.h5", + "4007/cutouts_h5_final/4007_block_z000768_000896_y001936_003872_x003312_004968.h5", + "4007/cutouts_h5_final/4007_block_z000896_001024_y001936_003872_x003312_004968.h5", + "4009/4009_raw_z0-128_y0-1600_x1600-3200_1.h5", + "4009/4009_raw_z0-128_y1600-3200_x0-1600_2.h5", + "4009/4009_raw_z0-128_y1600-3200_x1600-3200_3.h5", + "4009/4009_raw_z128-256_y0-1600_x0-1600_4.h5", + "4009/4009_raw_z128-256_y0-1600_x1600-3200_5.h5", + "4009/4009_z427_y1084_x1020.h5" + ], + "val": [ + "4007/cutouts_h5_final/4007_block_z000128_000256_y001936_003872_x003312_004968.h5", + "4009/4009_raw_z0-128_y0-1600_x0-1600_0.h5" + ] +} diff --git a/scripts/volume_em/training/train_mitochondria_vol_em.py b/scripts/volume_em/training/train_mitochondria_vol_em.py new file mode 100644 index 00000000..6455169f --- /dev/null +++ b/scripts/volume_em/training/train_mitochondria_vol_em.py @@ -0,0 +1,79 @@ +"""Train SynapseNet's volume EM mitochondria model. + +This reproduces the published model, which was trained on the FIB-SEM datasets 4007 and 4009 at a +voxel size of 25 nm in z and 5 nm in xy. The split is pinned in 'split-mito_vol_em_aniso2lvl_final.json', +which also carries the data roots, so the script runs unchanged after the data has been moved. + +How the published checkpoint came about is worth knowing, because its counters do not mean what they +look like. Two jobs for it were submitted 24 minutes apart and ran concurrently for about 23 hours, +both writing the same checkpoint folder. The second one picked up the first one's 'best.pt' as it +stood after 15 epochs and warm-started from it, weights only, so its optimizer, scheduler and +counters all restarted. Both then stopped on early stopping, neither on the 50,000 iterations or on +the wall clock. The surviving 'best.pt' is the second job's: 14,500 of its own iterations on top of +the 2,000 it inherited, about 16,500 gradient steps in total, with an optimizer reset in the middle. +The first job's remaining 14,125 steps, including its own best of 0.296378, were overwritten. + +That is an accident, not a recipe, and it gained little: 0.292714 against the 0.296378 the first job +had reached on its own, and 0.298842 for the sibling 'finalv2' run. This script therefore trains a +single clean run, which is what the recipe describes. Note that running it twice with the same name +and save_root is now refused outright rather than silently racing. + +To deviate from the recipe, use the CLI instead, which exposes all hyperparameters: +synapse_net.run_vol_em_mitochondria_training -h +""" + +import argparse +import os + +from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths, vol_em_mitochondria_training + +TRAIN_ROOTS = { + "4007": "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/4007_hdf5/all_cutouts_s0", + "4009": "/mnt/lustre-grete/projects/nim00020/data/volume-em/moebius/4009_hdf5/cutouts_segmented", +} +OUTPUT_ROOT = "/mnt/lustre-grete/usr/u15205/volume-em/models" +SPLIT_FILE = os.path.join(os.path.dirname(__file__), "split-mito_vol_em_aniso2lvl_final.json") + +PATCH_SHAPE = (32, 512, 512) +BATCH_SIZE = 4 +LEARNING_RATE = 1e-4 +N_ITERATIONS = 50000 +EARLY_STOPPING = 20 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("-n", "--name", default="volume-em-mito-aniso2lvl-lr1e-4-bs4-ps32x512x512-blacked-final", + help="The name of the model to be trained.") + parser.add_argument("-i", "--train_root", nargs="+", help="Override the data roots of the split file.") + parser.add_argument("-o", "--output_root", default=OUTPUT_ROOT, help="Where to save the checkpoint and log dir.") + parser.add_argument("--random_split", action="store_true", + help="Split the data randomly instead of using the pinned split.") + parser.add_argument("--resume", action="store_true", help="Continue a previous run with the same name.") + parser.add_argument("--check", action="store_true", + help="Visualize samples from the data loaders instead of running training.") + args = parser.parse_args() + + train_roots = TRAIN_ROOTS if args.train_root is None else args.train_root + train_paths, val_paths = get_vol_em_mitochondria_paths( + train_roots, split_file=None if args.random_split else SPLIT_FILE, + ) + print("Training on", len(train_paths), "blocks and validating on", len(val_paths), "blocks.") + + vol_em_mitochondria_training( + name=args.name, + train_paths=train_paths, + val_paths=val_paths, + save_root=args.output_root, + patch_shape=PATCH_SHAPE, + batch_size=BATCH_SIZE, + lr=LEARNING_RATE, + n_iterations=N_ITERATIONS, + early_stopping=EARLY_STOPPING, + resume=args.resume, + check=args.check, + ) + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index a981a0ef..71218ef9 100644 --- a/setup.py +++ b/setup.py @@ -45,6 +45,9 @@ "synapse_net.export_to_imod_objects = synapse_net.tools.cli:imod_object_cli", "synapse_net.run_supervised_training = synapse_net.training.supervised_training:main", "synapse_net.run_domain_adaptation = synapse_net.training.domain_adaptation:main", + "synapse_net.run_mitochondria_training = synapse_net.training.mitochondria:main", + "synapse_net.run_cristae_training = synapse_net.training.cristae:main", + "synapse_net.run_vol_em_mitochondria_training = synapse_net.training.mitochondria_vol_em:main", "synapse_net.visualize_vesicle_pools = synapse_net.tools.cli:pool_visualization_cli", "synapse_net.run_cristae_analysis = synapse_net.tools.cli:cristae_analysis_cli", ], diff --git a/synapse_net/training/__init__.py b/synapse_net/training/__init__.py index 84204c95..94cb4aa0 100644 --- a/synapse_net/training/__init__.py +++ b/synapse_net/training/__init__.py @@ -3,4 +3,7 @@ from .supervised_training import supervised_training from .semisupervised_training import semisupervised_training from .domain_adaptation import mean_teacher_adaptation +from .mitochondria import mitochondria_training +from .cristae import cristae_training +from .mitochondria_vol_em import vol_em_mitochondria_training from .transform import AZDistanceLabelTransform diff --git a/synapse_net/training/cristae.py b/synapse_net/training/cristae.py new file mode 100644 index 00000000..51003a13 --- /dev/null +++ b/synapse_net/training/cristae.py @@ -0,0 +1,387 @@ +"""Training for cristae segmentation in electron tomography. + +This module reproduces the recipe that was used to train SynapseNet's `cristae5` model, which +segments cristae in electron tomograms at a resolution of 1.74 nm. + +Cristae training differs from the other training functions in this package in three ways: + +- The network sees **two input channels**: the tomogram and a semantic mitochondria state, where 0 is + background, 1 is a mitochondrion that carries cristae annotations and 2 is a mitochondrion that + does not. Both are stored in a single hdf5 dataset of shape (2, z, y, x). +- The voxels of unannotated mitochondria (state 2) are **excluded from the loss**, so that the network + is not penalized for predicting cristae where no annotation exists. The background stays in the + loss, so the network still learns that there are no cristae outside of mitochondria. The loss mask + is carried in the second half of the target channels, see + `synapse_net.training.loss.MaskedDiceLossPerSample`. +- The loss is optionally **weighted towards the mitochondria membrane**, which markedly improves the + detection of cristae junctions, see `synapse_net.training.transform.membrane_proximity_weight`. + +The training data is expected as hdf5 files, one per tomogram, that contain the tomogram and the +mitochondria state under the key 'raw_mitos_combined' and the cristae annotations under +'labels/cristae'. The files may be spread over several data roots and their sub-directories. +""" + +import os +import random +from glob import glob +from typing import List, Mapping, Optional, Sequence, Tuple, Union + +import torch_em +from torch_em.data import MinInstanceSampler + +from .loss import MaskedDiceLossPerSample +from .split import _normalize_roots, _resolve_split +from .supervised_training import supervised_training +from .transform import ( + AugmentedMitoStateMaskTransform, CRISTAE_VOXEL_SIZE, MitoStateMaskTransform, standardize_channel +) + + +def get_cristae_test_paths( + split_file: str, + data_roots: Optional[Union[Mapping[str, str], Sequence[str], str]] = None, +) -> List[str]: + """Find the tomograms that a split file holds out for testing. + + The training functions only use the 'train' and 'val' entries of a split file. Use this to get + the 'test' entries, either to evaluate on them or to keep them out of a random split, which + would otherwise train on them. + + Args: + split_file: Path to the json file with the split. + data_roots: The folders that contain the data. Can be left out if the split file carries + the roots it was created with. + + Returns: + The filepaths of the test data. Empty if the split file has no 'test' entry. + + Raises: + ValueError: If a root referenced by `split_file` is not given. + """ + roots = {} if data_roots is None else _normalize_roots(data_roots) + test_paths, = _resolve_split(split_file, roots, ("test",)) + return test_paths + + +def get_cristae_paths( + data_roots: Optional[Union[Mapping[str, str], Sequence[str], str]] = None, + file_pattern: str = "*_combined.h5", + val_fraction: float = 0.1, + split_file: Optional[str] = None, + exclude: Sequence[str] = (), + seed: int = 42, +) -> Tuple[List[str], List[str]]: + """Find the tomograms for cristae training and split them into training and validation data. + + The files are found recursively in each of the `data_roots` and are then shuffled with a fixed + seed, so that the split is deterministic for a given set of files. Pass `split_file` to use an + explicit split instead; this is the only way to reproduce the split of a previous training run, + because the seeded split depends on which files are present. + + Args: + data_roots: The folders that contain the training data. Either a mapping of a short name to + the folder, or a list of folders, in which case the names are derived from them. The names + are used to refer to the files in `split_file`. It can be left out if `split_file` contains + the roots it was created with, and takes precedence over them, so that data can be moved. + file_pattern: The pattern for selecting the files with training data. + val_fraction: The fraction of the data to use for validation. + split_file: Path to a json file with the keys 'train' and 'val', which each hold a list of + entries of the form '/'. If given, `val_fraction`, + `exclude` and `seed` are not used. + exclude: Substrings of filepaths that should be excluded from the training data, for example + to leave out volumes with unreliable annotations. + seed: The seed for shuffling the files before splitting them. + + Returns: + The filepaths for training. + The filepaths for validation. + + Raises: + ValueError: If no data roots are given, if no files are found, if a root referenced by + `split_file` is not given, or if a file listed in `split_file` does not exist. + """ + roots = {} if data_roots is None else _normalize_roots(data_roots) + + if split_file is not None: + train_paths, val_paths = _resolve_split(split_file, roots, ("train", "val")) + missing = [path for path in train_paths + val_paths if not os.path.exists(path)] + if missing: + raise ValueError(f"{len(missing)} files from the split file {split_file} do not exist, e.g. {missing[0]}.") + return train_paths, val_paths + + if not roots: + raise ValueError("No data roots were given, and the split file does not contain any.") + + paths = [] + for root in roots.values(): + paths.extend(glob(os.path.join(root, "**", file_pattern), recursive=True)) + paths = sorted(path for path in paths if not any(pattern in path for pattern in exclude)) + if len(paths) == 0: + raise ValueError(f"Did not find any files matching {file_pattern} in {list(roots.values())}.") + + # We use a separate random state so that the global one is not affected by this function. + random.Random(seed).shuffle(paths) + + n_val = int(len(paths) * val_fraction) + # Make sure that we validate on at least one tomogram if there is more than one. + if n_val == 0 and len(paths) > 1: + n_val = 1 + n_train = len(paths) - n_val + return paths[:n_train], paths[n_train:] + + +def cristae_training( + name: str, + train_paths: Sequence[str], + val_paths: Sequence[str], + save_root: Optional[str] = None, + raw_key: str = "raw_mitos_combined", + label_key: str = "labels/cristae", + patch_shape: Tuple[int, int, int] = (32, 256, 256), + batch_size: int = 24, + lr: float = 1.7e-4, + n_iterations: int = int(1e5), + n_samples_train: Optional[int] = None, + n_samples_val: Optional[int] = None, + initial_features: int = 32, + early_stopping: Optional[int] = 25, + log_image_interval: int = 50, + checkpoint_path: Optional[str] = None, + resume: bool = False, + overwrite: bool = False, + seed: Optional[int] = None, + deterministic: bool = False, + sampler: Optional[Union[callable, bool]] = None, + raw_transform: Optional[callable] = None, + state_channel: int = 1, + ignore_state_value: float = 2.0, + membrane_w_pos: float = 3.0, + membrane_w_neg: float = 2.0, + membrane_band_nm: float = 12.0, + membrane_offset_nm: float = 0.0, + voxel_size: Tuple[float, float, float] = CRISTAE_VOXEL_SIZE, + augmentations: bool = True, + num_workers: int = 8, + persistent_workers: bool = True, + prefetch_factor: int = 4, + shuffle: bool = False, + mixed_precision: bool = True, + check: bool = False, + **kwargs, +) -> None: + """Run supervised training for cristae segmentation in electron tomograms. + + The default arguments reproduce SynapseNet's `cristae5` model: an anisotropic U-Net **without + normalization layers** that takes the tomogram and the mitochondria state as input, predicts + foreground and boundaries, and is trained with a per-sample masked dice loss that is weighted + towards the mitochondria membrane. + + Two things about this recipe are easy to get wrong: + + - The loss mask is carried in the **target**, not in the model input. It is produced by the joint + transform and consumed by `MaskedDiceLossPerSample`, so a custom `loss_fn` has to understand + that layout. + - Setting `augmentations=False` means training **without any augmentation at all**, because the + joint transform occupies the slot of the dataloader that otherwise holds them. + + The normalization is part of the model: only the tomogram channel is standardized, and the + mitochondria state channel is left untouched so that its values stay comparable against + `ignore_state_value`. + + Args: + name: The name for the checkpoint to be trained. + train_paths: Filepaths to the hdf5 files for the training data. + val_paths: Filepaths to the hdf5 files for the validation data. + save_root: Folder where the checkpoint will be saved. + raw_key: The key that holds the tomogram and the mitochondria state inside of the hdf5. + label_key: The key that holds the cristae annotations inside of the hdf5. + patch_shape: The patch shape used for a training example. + batch_size: The batch size for training. + lr: The initial learning rate. + n_iterations: The maximal number of iterations to train for. + n_samples_train: The number of train samples per epoch. By default this will be estimated + based on the patch shape and the size of the volumes used for training. + n_samples_val: The number of val samples per epoch. By default this will be estimated + based on the patch shape and the size of the volumes used for validation. + initial_features: The number of features in the first level of the U-Net. + This has no effect if the model is initialized from `checkpoint_path`. + early_stopping: The number of epochs without improvement after which training is stopped. + log_image_interval: The interval (in iterations) at which images are written to the training log. + checkpoint_path: Path to a model checkpoint to initialize the weights from. Only the + weights are loaded; pass `resume` instead to continue a previous run. + resume: Whether to continue the previous run with this name in `save_root`, restoring + its weights, optimizer and iteration count. `n_iterations` then counts the total + iterations, including the ones the previous run already did. + overwrite: Whether to replace the checkpoints of a previous run with this name. + seed: The seed for the random number generators used in training. By default it is + not set, so that repeated runs differ. + deterministic: Whether to also disable cudnn benchmarking when a `seed` is given. + sampler: Sampler to accept or reject patches for training. + By default a minimum instance sampler with a rejection probability of 0.95 is used. + raw_transform: Transformation applied to the input before it is passed to the network. + By default the tomogram channel is standardized and the state channel is left unchanged. + state_channel: The channel of `raw_key` that holds the mitochondria state. + ignore_state_value: The state value that is excluded from the loss, i.e. the value marking + mitochondria without cristae annotations. + membrane_w_pos: The loss weight for cristae voxels close to the mitochondria membrane. + membrane_w_neg: The loss weight for non-cristae voxels close to the mitochondria membrane. + Pass 1.0 for both weights to train without membrane weighting. + membrane_band_nm: The thickness of the membrane band in nanometer. + membrane_offset_nm: Start the membrane band this far inside the membrane. + voxel_size: The voxel size of the training data in nanometer, in the order (z, y, x). + It is only used to convert the membrane band from nanometer to voxels. + augmentations: Whether to augment the training data. Note that this is not a no-op switch: + with `False` the training runs without any augmentation. + num_workers: The number of workers for the dataloaders. + persistent_workers: Whether to keep the dataloader workers alive between epochs. + prefetch_factor: The number of batches prefetched by each dataloader worker. + shuffle: Whether to shuffle the datasets in the dataloaders. This is disabled by default, + to match the published model; enabling it is recommended for new trainings, as it mixes + the tomograms within a batch. + mixed_precision: Whether to train with mixed precision. + check: Whether to check the training and validation loaders instead of running training. + kwargs: Additional keyword arguments for `synapse_net.training.supervised_training`. + """ + if sampler is None: + sampler = MinInstanceSampler(p_reject=0.95) + if raw_transform is None: + raw_transform = standardize_channel + + mask_kwargs = dict( + mito_channel=state_channel, exclude_state_value=ignore_state_value, + band_nm=membrane_band_nm, offset_nm=membrane_offset_nm, + w_pos=membrane_w_pos, w_neg=membrane_w_neg, voxel_size=voxel_size, + ) + if augmentations: + transform = AugmentedMitoStateMaskTransform(torch_em.transform.get_augmentations(3), **mask_kwargs) + else: + transform = MitoStateMaskTransform(**mask_kwargs) + + supervised_training( + name=name, + train_paths=tuple(train_paths), + val_paths=tuple(val_paths), + label_key=label_key, + patch_shape=patch_shape, + save_root=save_root, + raw_key=raw_key, + batch_size=batch_size, + lr=lr, + n_iterations=n_iterations, + sampler=sampler, + n_samples_train=n_samples_train, + n_samples_val=n_samples_val, + check=check, + loss_fn=MaskedDiceLossPerSample(), + in_channels=2, + out_channels=2, + initial_features=initial_features, + norm=None, + transform=transform, + checkpoint_path=checkpoint_path, + resume=resume, + overwrite=overwrite, + seed=seed, + deterministic=deterministic, + mixed_precision=mixed_precision, + early_stopping=early_stopping, + log_image_interval=log_image_interval, + raw_transform=raw_transform, + with_channels=True, + num_workers=num_workers, + persistent_workers=persistent_workers, + prefetch_factor=prefetch_factor, + shuffle=shuffle, + **kwargs, + ) + + +def main(): + """@private + """ + import argparse + + parser = argparse.ArgumentParser( + description="Train a model for cristae segmentation in electron tomograms.\n\n" + "The training data has to be stored in hdf5 files, one per tomogram, which contain the tomogram and the\n" + "mitochondria state in one dataset of shape (2, z, y, x), and the cristae annotations in another. The files\n" + "are found recursively in the folders passed via '-i'. For example:\n" + "synapse_net.run_cristae_training -n my_cristae_model -i /path/to/tomograms\n" + "The trained model will be saved in the folder 'checkpoints/my_cristae_model'.\n" + "The default arguments reproduce the training of SynapseNet's 'cristae5' model. Note that this model was\n" + "trained with a batch size of 24 and a patch shape of 32 x 256 x 256, which requires a large GPU.\n" + "Check out the information below for details on the arguments of this function.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("-n", "--name", required=True, help="The name of the model to be trained.") + parser.add_argument("-i", "--data_root", nargs="+", help="One or more folders with the training data. They are searched recursively. Can be left out if 'split_file' contains the data roots.") # noqa + parser.add_argument("--file_pattern", default="*_combined.h5", help="The pattern for selecting the files with training data.") # noqa + parser.add_argument("--raw_key", default="raw_mitos_combined", + help="The internal path for the tomogram and the mitochondria state.") + parser.add_argument("--label_key", default="labels/cristae", help="The internal path for the cristae annotations.") # noqa + + # How to split the data into training and validation data. + parser.add_argument("--split_file", help="A json file with the keys 'train' and 'val', which each hold a list of entries '/'. If not given, the data is split randomly.") # noqa + parser.add_argument("--val_fraction", type=float, default=0.1, help="The fraction of the data to use for validation. This has no effect if 'split_file' was passed.") # noqa + parser.add_argument("--exclude", nargs="*", default=[], help="Substrings of filepaths to exclude from the training data.") # noqa + parser.add_argument("--exclude_test_from", help="A split file whose 'test' entries are excluded from the training data. Use this when splitting randomly, so that the run does not train on the test volumes of a previous split.") # noqa + parser.add_argument("--seed", type=int, default=42, help="The seed for shuffling the files before splitting them, and for the random number generators used in training.") # noqa + + # The training hyperparameters. + parser.add_argument("-p", "--patch_shape", nargs=3, type=int, default=[32, 256, 256], + help="The patch shape for training in ZYX.") + parser.add_argument("--batch_size", type=int, default=24, help="The batch size for training.") + parser.add_argument("--learning_rate", type=float, default=1.7e-4, help="The initial learning rate.") + parser.add_argument("--n_iterations", type=int, default=int(1e5), help="The maximal number of iterations to train for.") # noqa + parser.add_argument("--n_samples_train", type=int, help="The number of samples per epoch for training. If not given will be derived from the data size.") # noqa + parser.add_argument("--n_samples_val", type=int, help="The number of samples per epoch for validation. If not given will be derived from the data size.") # noqa + parser.add_argument("--initial_features", type=int, default=32, help="The number of features in the first level of the U-Net.") # noqa + parser.add_argument("--early_stopping", type=int, default=25, help="The number of epochs without improvement after which training is stopped.") # noqa + parser.add_argument("--num_workers", type=int, default=8, help="The number of workers for the dataloaders.") + parser.add_argument("--no_augmentations", action="store_true", help="Train without any augmentation. This is not recommended, it is only here to reproduce earlier trainings.") # noqa + parser.add_argument("--shuffle", action="store_true", help="Shuffle the datasets in the dataloaders. This is recommended for new trainings, but was not used for the published model.") # noqa + + # The loss masking and the membrane weighting. + parser.add_argument("--state_channel", type=int, default=1, help="The channel holding the mitochondria state.") + parser.add_argument("--ignore_state_value", type=float, default=2.0, help="The mitochondria state that is excluded from the loss, i.e. the value marking mitochondria without cristae annotations.") # noqa + parser.add_argument("--membrane_w_pos", type=float, default=3.0, help="The loss weight for cristae voxels close to the mitochondria membrane.") # noqa + parser.add_argument("--membrane_w_neg", type=float, default=2.0, help="The loss weight for non-cristae voxels close to the mitochondria membrane. Pass 1.0 for both weights to train without membrane weighting.") # noqa + parser.add_argument("--membrane_band_nm", type=float, default=12.0, help="The thickness of the membrane band in nanometer.") # noqa + parser.add_argument("--membrane_offset_nm", type=float, default=0.0, help="Start the membrane band this far inside the membrane.") # noqa + parser.add_argument("--voxel_size", type=float, nargs=3, default=[1.74, 1.74, 1.74], help="The voxel size of the training data in nanometer, in ZYX. Only used for the membrane band.") # noqa + + # Where to save the model, and how to initialize it. + parser.add_argument("--save_root", help="Root path for saving the checkpoint and log dir.") + parser.add_argument("--checkpoint_path", help="A model checkpoint to initialize the weights from.") + parser.add_argument("--resume", action="store_true", help="Continue the previous run with the same name in 'save_root', restoring its weights, optimizer and iteration count. '--n_iterations' is the total number of iterations, including the ones already done.") # noqa + parser.add_argument("--overwrite", action="store_true", help="Replace the checkpoints of a previous run with the same name. By default training refuses to overwrite them.") # noqa + parser.add_argument("--deterministic", action="store_true", help="Disable cudnn benchmarking so that runs with the same seed match exactly. This costs throughput.") # noqa + parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.") # noqa + args = parser.parse_args() + + exclude = list(args.exclude) + if args.exclude_test_from is not None: + exclude += get_cristae_test_paths(args.exclude_test_from, args.data_root) + + train_paths, val_paths = get_cristae_paths( + args.data_root, file_pattern=args.file_pattern, val_fraction=args.val_fraction, + split_file=args.split_file, exclude=exclude, seed=args.seed, + ) + print("Training on", len(train_paths), "tomograms and validating on", len(val_paths), "tomograms.") + + checkpoint_path = args.checkpoint_path + + cristae_training( + name=args.name, train_paths=train_paths, val_paths=val_paths, save_root=args.save_root, + raw_key=args.raw_key, label_key=args.label_key, patch_shape=tuple(args.patch_shape), + batch_size=args.batch_size, lr=args.learning_rate, n_iterations=args.n_iterations, + n_samples_train=args.n_samples_train, n_samples_val=args.n_samples_val, + initial_features=args.initial_features, early_stopping=args.early_stopping, + checkpoint_path=checkpoint_path, resume=args.resume, overwrite=args.overwrite, + seed=args.seed, deterministic=args.deterministic, + num_workers=args.num_workers, shuffle=args.shuffle, + state_channel=args.state_channel, ignore_state_value=args.ignore_state_value, + membrane_w_pos=args.membrane_w_pos, membrane_w_neg=args.membrane_w_neg, + membrane_band_nm=args.membrane_band_nm, membrane_offset_nm=args.membrane_offset_nm, + voxel_size=tuple(args.voxel_size), augmentations=not args.no_augmentations, check=args.check, + ) diff --git a/synapse_net/training/dataloaders.py b/synapse_net/training/dataloaders.py index 29c5dd9f..c9089b27 100644 --- a/synapse_net/training/dataloaders.py +++ b/synapse_net/training/dataloaders.py @@ -39,6 +39,7 @@ def get_supervised_loader( label_transform: Optional[callable] = None, label_paths: Optional[Tuple[str]] = None, raw_transform: Optional[callable] = None, + transform: Optional[callable] = None, **loader_kwargs, ) -> torch.utils.data.DataLoader: """Get a dataloader for supervised segmentation training. @@ -65,6 +66,8 @@ def get_supervised_loader( label_paths: Optional paths containing the labels / annotations for training. If not given, the labels are expected to be contained in the `data_paths`. raw_transform: Normalization applied to the raw data. By default the torch_em standardization is used. + transform: Joint transformation applied to the raw data and the labels, after the + `label_transform`. By default padding and the standard augmentations are applied. loader_kwargs: Additional keyword arguments for the dataloader. Returns: @@ -86,7 +89,9 @@ def get_supervised_loader( raise NotImplementedError label_transform = torch_em.transform.label.connected_components - if ndim == 2: + if transform is not None: # A specific joint transform was passed, do nothing. + pass + elif ndim == 2: adjusted_patch_shape = _adjust_patch_shape(ndim, patch_shape) transform = torch_em.transform.Compose( torch_em.transform.PadIfNecessary(adjusted_patch_shape), torch_em.transform.get_augmentations(2) diff --git a/synapse_net/training/loss.py b/synapse_net/training/loss.py new file mode 100644 index 00000000..06a92f1f --- /dev/null +++ b/synapse_net/training/loss.py @@ -0,0 +1,100 @@ +"""Loss functions that read a per-voxel loss mask from the training target. + +These losses expect a target with twice the number of channels of the prediction: the first half +holds the actual targets, the second half holds a per-channel mask (or weight) that says how much +each voxel contributes to the loss. Targets of this shape are produced by +`synapse_net.training.transform.MitoStateMaskTransform` and are used for cristae training, where +mitochondria without cristae annotations have to be excluded from the loss. + +Carrying the mask in the target keeps the loss a plain two-argument `loss(prediction, target)`, +so it needs no special support from the trainer. +""" + +import torch +import torch.nn as nn + + +def _split_target(prediction: torch.Tensor, target: torch.Tensor, name: str): + n_channels = prediction.size(1) + if target.size(1) != 2 * n_channels: + raise ValueError(f"{name} expects a target with {2 * n_channels} channels, got {target.size(1)}.") + # The square root is applied so that a continuous weight w enters both the numerator and the + # denominator of the dice linearly. For a binary mask this is the identity, so an unweighted + # run is not affected by it. + mask = torch.sqrt(target[:, n_channels:]) + return prediction * mask, target[:, :n_channels] * mask + + +class MaskedDiceLoss(nn.Module): + """Dice loss that reads a per-channel loss mask from the second half of the target. + + The dice is computed per channel over the full batch (ratio of sums). See + `MaskedDiceLossPerSample` for the variant that averages the dice over the batch instead. + + Args: + eps: The epsilon value used to avoid division by zero. + """ + + def __init__(self, eps: float = 1e-7): + super().__init__() + self.eps = eps + # Enables torch-em to serialize this loss into the checkpoint, so that the checkpoint can be + # deserialized again without having to reconstruct the loss by hand. + self.init_kwargs = {"eps": eps} + + def forward(self, prediction: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor: + """Compute the masked dice loss. + + Args: + prediction: The predictions of the network. + target: The targets, with the loss mask in the second half of the channels. + + Returns: + The loss value. + """ + pred, tgt = _split_target(prediction, target, "MaskedDiceLoss") + # Flatten to [C, B * spatial] so that the dice is computed per channel across the batch, + # matching the reduction of torch_em's DiceLoss. + n_channels = pred.size(1) + pred = pred.permute(1, 0, *range(2, pred.dim())).reshape(n_channels, -1) + tgt = tgt.permute(1, 0, *range(2, tgt.dim())).reshape(n_channels, -1) + numerator = (pred * tgt).sum(-1) + denominator = (pred * pred).sum(-1) + (tgt * tgt).sum(-1) + return (1.0 - 2.0 * numerator / denominator.clamp(min=self.eps)).sum() + + +class MaskedDiceLossPerSample(nn.Module): + """Masked dice loss that averages the dice over the batch instead of pooling it. + + The dice is computed for each element of the batch, reducing over the spatial dimensions only, + and the results are then averaged (mean of ratios). Compared to `MaskedDiceLoss` (ratio of sums) + this weights every patch equally, regardless of how many foreground voxels it contains, which + counters the class imbalance of sparse structures such as cristae. + + Args: + eps: The epsilon value used to avoid division by zero. + """ + + def __init__(self, eps: float = 1e-7): + super().__init__() + self.eps = eps + # Enables torch-em to serialize this loss into the checkpoint, so that the checkpoint can be + # deserialized again without having to reconstruct the loss by hand. + self.init_kwargs = {"eps": eps} + + def forward(self, prediction: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor: + """Compute the per-sample masked dice loss. + + Args: + prediction: The predictions of the network. + target: The targets, with the loss mask in the second half of the channels. + + Returns: + The loss value. + """ + pred, tgt = _split_target(prediction, target, "MaskedDiceLossPerSample") + spatial_dims = tuple(range(2, pred.dim())) + numerator = (pred * tgt).sum(spatial_dims) + denominator = (pred * pred).sum(spatial_dims) + (tgt * tgt).sum(spatial_dims) + dice = (2.0 * numerator / denominator.clamp(min=self.eps)).mean(dim=0) + return (1.0 - dice).sum() diff --git a/synapse_net/training/mitochondria.py b/synapse_net/training/mitochondria.py new file mode 100644 index 00000000..a2f4663b --- /dev/null +++ b/synapse_net/training/mitochondria.py @@ -0,0 +1,263 @@ +"""Training for mitochondria segmentation in electron tomography. + +This module reproduces the recipe that was used to train SynapseNet's `mitochondria2` model, +which segments mitochondria in electron tomograms at a resolution of 2.87 nm. +It is a thin layer on top of `synapse_net.training.supervised_training` that fixes the +hyperparameters and the data layout of that model. + +The training data is expected as hdf5 files, one per tomogram, that contain the tomogram +under the key 'raw' and the corresponding instance annotations under 'labels/mitochondria'. +The files may be spread over sub-directories of the data root; they are found recursively. +""" + +import json +import os +import random +from glob import glob +from typing import List, Optional, Sequence, Tuple, Union + +import torch_em +from torch_em.data import MinInstanceSampler + +from .supervised_training import supervised_training + + +def get_mitochondria_paths( + data_root: str, + file_pattern: str = "*.h5", + val_fraction: float = 0.15, + split_file: Optional[str] = None, + seed: int = 42, +) -> Tuple[List[str], List[str]]: + """Find the tomograms for mitochondria training and split them into training and validation data. + + The files are found recursively in `data_root` and are then shuffled with a fixed seed, + so that the split is deterministic for a given set of files. Pass `split_file` to use an + explicit split instead, for example to reproduce the split of a previous training run. + + Args: + data_root: The folder that contains the training data. + file_pattern: The pattern for selecting the files with training data. + val_fraction: The fraction of the data to use for validation. + split_file: Path to a json file with the keys 'train' and 'val', which each hold a list of + filepaths relative to `data_root`. If given, `val_fraction` and `seed` are not used. + seed: The seed for shuffling the files before splitting them. + + Returns: + The filepaths for training. + The filepaths for validation. + + Raises: + ValueError: If no files are found, or if a file listed in `split_file` does not exist. + """ + if split_file is not None: + with open(split_file) as f: + split = json.load(f) + train_paths = [os.path.join(data_root, name) for name in split["train"]] + val_paths = [os.path.join(data_root, name) for name in split["val"]] + missing = [path for path in train_paths + val_paths if not os.path.exists(path)] + if missing: + raise ValueError(f"{len(missing)} files from the split file {split_file} do not exist, e.g. {missing[0]}.") + return train_paths, val_paths + + paths = sorted(glob(os.path.join(data_root, "**", file_pattern), recursive=True)) + if len(paths) == 0: + raise ValueError(f"Did not find any files matching {file_pattern} in {data_root}.") + + # We use a separate random state so that the global one is not affected by this function. + random.Random(seed).shuffle(paths) + + n_val = int(len(paths) * val_fraction) + # Make sure that we validate on at least one tomogram if there is more than one. + if n_val == 0 and len(paths) > 1: + n_val = 1 + n_train = len(paths) - n_val + return paths[:n_train], paths[n_train:] + + +def mitochondria_training( + name: str, + train_paths: Sequence[str], + val_paths: Sequence[str], + save_root: Optional[str] = None, + raw_key: str = "raw", + label_key: str = "labels/mitochondria", + patch_shape: Tuple[int, int, int] = (32, 256, 256), + batch_size: int = 8, + lr: float = 1e-4, + n_iterations: int = int(1.5e5), + n_samples_train: int = 500, + n_samples_val: int = 500, + initial_features: int = 32, + early_stopping: Optional[int] = 20, + log_image_interval: int = 50, + checkpoint_path: Optional[str] = None, + resume: bool = False, + overwrite: bool = False, + seed: Optional[int] = None, + deterministic: bool = False, + sampler: Optional[Union[callable, bool]] = None, + raw_transform: Optional[callable] = None, + num_workers: int = 8, + shuffle: bool = False, + mixed_precision: bool = False, + check: bool = False, + **kwargs, +) -> None: + """Run supervised training for mitochondria segmentation in electron tomograms. + + The default arguments reproduce SynapseNet's `mitochondria2` model: an anisotropic U-Net + with instance normalization that predicts foreground and boundaries, trained with a dice loss + on percentile-normalized tomograms. + + The raw transform is part of the model, not just of the training: the same normalization has to + be applied at inference time, by passing `preprocess=torch_em.transform.raw.normalize_percentile` + to `synapse_net.inference.mitochondria.segment_mitochondria`. A model trained with one + normalization and applied with another will not predict anything useful. + + Args: + name: The name for the checkpoint to be trained. + train_paths: Filepaths to the hdf5 files for the training data. + val_paths: Filepaths to the hdf5 files for the validation data. + save_root: Folder where the checkpoint will be saved. + raw_key: The key that holds the tomogram inside of the hdf5. + label_key: The key that holds the mitochondria annotations inside of the hdf5. + patch_shape: The patch shape used for a training example. + batch_size: The batch size for training. + lr: The initial learning rate. + n_iterations: The maximal number of iterations to train for. + n_samples_train: The number of train samples per epoch. + n_samples_val: The number of val samples per epoch. + initial_features: The number of features in the first level of the U-Net. + This has no effect if the model is initialized from `checkpoint_path`. + early_stopping: The number of epochs without improvement after which training is stopped. + log_image_interval: The interval (in iterations) at which images are written to the training log. + checkpoint_path: Path to a model checkpoint to initialize the weights from. Only the + weights are loaded; pass `resume` instead to continue a previous run. + resume: Whether to continue the previous run with this name in `save_root`, restoring + its weights, optimizer and iteration count. `n_iterations` then counts the total + iterations, including the ones the previous run already did. + overwrite: Whether to replace the checkpoints of a previous run with this name. + seed: The seed for the random number generators used in training. By default it is + not set, so that repeated runs differ. + deterministic: Whether to also disable cudnn benchmarking when a `seed` is given. + sampler: Sampler to accept or reject patches for training. + By default a minimum instance sampler with a rejection probability of 0.95 is used. + raw_transform: Transformation applied to the tomogram before it is passed to the network. + By default the tomogram is normalized with the 1st and 99th percentile. + num_workers: The number of workers for the dataloaders. + shuffle: Whether to shuffle the datasets in the dataloaders. This is disabled by default, + to match the published model; enabling it is recommended for new trainings, as it mixes + the tomograms within a batch. + mixed_precision: Whether to train with mixed precision. This is disabled by default, to match + the published model; enabling it reduces the memory consumption and speeds up training. + check: Whether to check the training and validation loaders instead of running training. + kwargs: Additional keyword arguments for `synapse_net.training.supervised_training`. + """ + if sampler is None: + sampler = MinInstanceSampler(p_reject=0.95) + if raw_transform is None: + raw_transform = torch_em.transform.raw.normalize_percentile + + supervised_training( + name=name, + train_paths=tuple(train_paths), + val_paths=tuple(val_paths), + label_key=label_key, + patch_shape=patch_shape, + save_root=save_root, + raw_key=raw_key, + batch_size=batch_size, + lr=lr, + n_iterations=n_iterations, + sampler=sampler, + n_samples_train=n_samples_train, + n_samples_val=n_samples_val, + check=check, + initial_features=initial_features, + checkpoint_path=checkpoint_path, + resume=resume, + overwrite=overwrite, + seed=seed, + deterministic=deterministic, + mixed_precision=mixed_precision, + early_stopping=early_stopping, + log_image_interval=log_image_interval, + raw_transform=raw_transform, + num_workers=num_workers, + shuffle=shuffle, + **kwargs, + ) + + +def main(): + """@private + """ + import argparse + + parser = argparse.ArgumentParser( + description="Train a model for mitochondria segmentation in electron tomograms.\n\n" + "The training data has to be stored in hdf5 files, one per tomogram, which contain the tomogram and the\n" + "mitochondria annotations. The files are found recursively in the folder passed via '-i'. For example:\n" + "synapse_net.run_mitochondria_training -n my_mito_model -i /path/to/tomograms\n" + "The trained model will be saved in the folder 'checkpoints/my_mito_model'.\n" + "The default arguments reproduce the training of SynapseNet's 'mitochondria2' model. Note that this model\n" + "was trained with a batch size of 8 and a patch shape of 32 x 256 x 256, which requires a large GPU.\n" + "Check out the information below for details on the arguments of this function.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("-n", "--name", required=True, help="The name of the model to be trained.") + parser.add_argument("-i", "--data_root", required=True, + help="The folder with the training data. It is searched recursively for matching files.") + parser.add_argument("--file_pattern", default="*.h5", help="The pattern for selecting the files with training data.") # noqa + parser.add_argument("--raw_key", default="raw", help="The internal path for the tomogram.") + parser.add_argument("--label_key", default="labels/mitochondria", + help="The internal path for the mitochondria annotations.") + + # How to split the data into training and validation data. + parser.add_argument("--split_file", help="A json file with the keys 'train' and 'val', which each hold a list of filepaths relative to 'data_root'. If not given, the data is split randomly.") # noqa + parser.add_argument("--val_fraction", type=float, default=0.15, help="The fraction of the data to use for validation. This has no effect if 'split_file' was passed.") # noqa + parser.add_argument("--seed", type=int, default=42, help="The seed for shuffling the files before splitting them, and for the random number generators used in training.") # noqa + + # The training hyperparameters. + parser.add_argument("-p", "--patch_shape", nargs=3, type=int, default=[32, 256, 256], + help="The patch shape for training in ZYX.") + parser.add_argument("--batch_size", type=int, default=8, help="The batch size for training.") + parser.add_argument("--learning_rate", type=float, default=1e-4, help="The initial learning rate.") + parser.add_argument("--n_iterations", type=int, default=int(1.5e5), help="The maximal number of iterations to train for.") # noqa + parser.add_argument("--n_samples_train", type=int, default=500, help="The number of samples per epoch for training.") # noqa + parser.add_argument("--n_samples_val", type=int, default=500, help="The number of samples per epoch for validation.") # noqa + parser.add_argument("--initial_features", type=int, default=32, help="The number of features in the first level of the U-Net.") # noqa + parser.add_argument("--early_stopping", type=int, default=20, help="The number of epochs without improvement after which training is stopped.") # noqa + parser.add_argument("--num_workers", type=int, default=8, help="The number of workers for the dataloaders.") + parser.add_argument("--shuffle", action="store_true", help="Shuffle the datasets in the dataloaders. This is recommended for new trainings, but was not used for the published model.") # noqa + parser.add_argument("--mixed_precision", action="store_true", help="Train with mixed precision. This reduces the memory consumption, but was not used for the published model.") # noqa + + # Where to save the model, and how to initialize it. + parser.add_argument("--save_root", help="Root path for saving the checkpoint and log dir.") + parser.add_argument("--checkpoint_path", help="A model checkpoint to initialize the weights from.") + parser.add_argument("--resume", action="store_true", help="Continue the previous run with the same name in 'save_root', restoring its weights, optimizer and iteration count. '--n_iterations' is the total number of iterations, including the ones already done.") # noqa + parser.add_argument("--overwrite", action="store_true", help="Replace the checkpoints of a previous run with the same name. By default training refuses to overwrite them.") # noqa + parser.add_argument("--deterministic", action="store_true", help="Disable cudnn benchmarking so that runs with the same seed match exactly. This costs throughput.") # noqa + parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.") # noqa + args = parser.parse_args() + + train_paths, val_paths = get_mitochondria_paths( + args.data_root, file_pattern=args.file_pattern, val_fraction=args.val_fraction, + split_file=args.split_file, seed=args.seed, + ) + print("Training on", len(train_paths), "tomograms and validating on", len(val_paths), "tomograms.") + + checkpoint_path = args.checkpoint_path + + mitochondria_training( + name=args.name, train_paths=train_paths, val_paths=val_paths, save_root=args.save_root, + raw_key=args.raw_key, label_key=args.label_key, patch_shape=tuple(args.patch_shape), + batch_size=args.batch_size, lr=args.learning_rate, n_iterations=args.n_iterations, + n_samples_train=args.n_samples_train, n_samples_val=args.n_samples_val, + initial_features=args.initial_features, early_stopping=args.early_stopping, + checkpoint_path=checkpoint_path, resume=args.resume, overwrite=args.overwrite, + seed=args.seed, deterministic=args.deterministic, + num_workers=args.num_workers, shuffle=args.shuffle, + mixed_precision=args.mixed_precision, check=args.check, + ) diff --git a/synapse_net/training/mitochondria_vol_em.py b/synapse_net/training/mitochondria_vol_em.py new file mode 100644 index 00000000..d7415f7f --- /dev/null +++ b/synapse_net/training/mitochondria_vol_em.py @@ -0,0 +1,332 @@ +"""Training for mitochondria segmentation in volume electron microscopy. + +This module reproduces the recipe that was used to train SynapseNet's volume EM mitochondria model on +FIB-SEM data with a voxel size of 25 nm in z and 5 nm in xy. It is a thin layer on top of +`synapse_net.training.supervised_training` that fixes the hyperparameters and the data layout of that +model. See `synapse_net.training.mitochondria` for the corresponding electron tomography model. + +Volume EM training differs from the tomography training in two ways: + +- The data is **strongly anisotropic** (5:1), so the U-Net downsamples only in xy for its first two + levels instead of one, see `VOL_EM_SCALE_FACTORS`. The network is also trained **without + normalization layers**. +- The blocks were cut out of a larger FIB-SEM volume and carry **white filler borders** where the + cutout extends past the imaged region. They are removed before the normalization, see + `synapse_net.training.transform.remove_white_patches`. + +The training data is expected as hdf5 files, one per block, that contain the image data under the key +'raw' and the corresponding instance annotations under 'labels/mitochondria'. The files may be spread +over several data roots and their sub-directories. +""" + +import os +import random +from glob import glob +from typing import List, Mapping, Optional, Sequence, Tuple, Union + +import torch_em +from torch_em.data import MinInstanceSampler + +from .split import _normalize_roots, _resolve_split +from .supervised_training import supervised_training +from .transform import RemoveWhitePatchesAndNormalize + +VOL_EM_SCALE_FACTORS = [[1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2]] +"""The downscaling factors of the U-Net encoder: two anisotropic levels, for the 5:1 anisotropy of the +volume EM data. After them the voxels are approximately isotropic, so the deeper levels downsample in +all three axes. +""" + + +def get_vol_em_mitochondria_paths( + data_roots: Optional[Union[Mapping[str, str], Sequence[str], str]] = None, + file_pattern: str = "*.h5", + val_fraction: float = 0.15, + split_file: Optional[str] = None, + exclude: Sequence[str] = (), + seed: int = 42, +) -> Tuple[List[str], List[str]]: + """Find the blocks for volume EM mitochondria training and split them into training and validation data. + + The files are found recursively in each of the `data_roots` and are then shuffled with a fixed + seed, so that the split is deterministic for a given set of files. Pass `split_file` to use an + explicit split instead; this is the only way to reproduce the split of a previous training run, + because the seeded split depends on which files are present. + + Args: + data_roots: The folders that contain the training data. Either a mapping of a short name to + the folder, or a list of folders, in which case the names are derived from them. The names + are used to refer to the files in `split_file`. It can be left out if `split_file` contains + the roots it was created with, and takes precedence over them, so that data can be moved. + file_pattern: The pattern for selecting the files with training data. + val_fraction: The fraction of the data to use for validation. + split_file: Path to a json file with the keys 'train' and 'val', which each hold a list of + entries of the form '/'. If given, `val_fraction`, + `exclude` and `seed` are not used. + exclude: Substrings of filepaths that should be excluded from the training data, for example + to leave out blocks with unreliable annotations. + seed: The seed for shuffling the files before splitting them. + + Returns: + The filepaths for training. + The filepaths for validation. + + Raises: + ValueError: If no data roots are given, if no files are found, if a root referenced by + `split_file` is not given, or if a file listed in `split_file` does not exist. + """ + roots = {} if data_roots is None else _normalize_roots(data_roots) + + if split_file is not None: + train_paths, val_paths = _resolve_split(split_file, roots, ("train", "val")) + missing = [path for path in train_paths + val_paths if not os.path.exists(path)] + if missing: + raise ValueError(f"{len(missing)} files from the split file {split_file} do not exist, e.g. {missing[0]}.") + return train_paths, val_paths + + if not roots: + raise ValueError("No data roots were given, and the split file does not contain any.") + + paths = [] + for root in roots.values(): + paths.extend(glob(os.path.join(root, "**", file_pattern), recursive=True)) + paths = sorted(path for path in paths if not any(pattern in path for pattern in exclude)) + if len(paths) == 0: + raise ValueError(f"Did not find any files matching {file_pattern} in {list(roots.values())}.") + + # We use a separate random state so that the global one is not affected by this function. + random.Random(seed).shuffle(paths) + + n_val = int(len(paths) * val_fraction) + # Make sure that we validate on at least one block if there is more than one. + if n_val == 0 and len(paths) > 1: + n_val = 1 + n_train = len(paths) - n_val + return paths[:n_train], paths[n_train:] + + +def vol_em_mitochondria_training( + name: str, + train_paths: Sequence[str], + val_paths: Sequence[str], + save_root: Optional[str] = None, + raw_key: str = "raw", + label_key: str = "labels/mitochondria", + patch_shape: Tuple[int, int, int] = (32, 512, 512), + batch_size: int = 4, + lr: float = 1e-4, + n_iterations: int = int(5e4), + n_samples_train: int = 500, + n_samples_val: int = 500, + initial_features: int = 32, + scale_factors: Sequence[Sequence[int]] = VOL_EM_SCALE_FACTORS, + norm: Optional[str] = None, + early_stopping: Optional[int] = 20, + log_image_interval: int = 50, + checkpoint_path: Optional[str] = None, + resume: bool = False, + overwrite: bool = False, + seed: Optional[int] = None, + deterministic: bool = False, + sampler: Optional[Union[callable, bool]] = None, + raw_transform: Optional[callable] = None, + fix_white_patches: bool = True, + white_patch_min_size: int = 20, + num_workers: int = 8, + shuffle: bool = False, + mixed_precision: bool = False, + check: bool = False, + **kwargs, +) -> None: + """Run supervised training for mitochondria segmentation in volume electron microscopy. + + The default arguments reproduce SynapseNet's volume EM mitochondria model: an anisotropic U-Net + **without normalization layers** and with **two anisotropic downsampling levels** that predicts + foreground and boundaries, trained with a dice loss on percentile-normalized FIB-SEM blocks whose + white filler borders have been removed. + + The raw transform is part of the model, not just of the training, and it has to be reproduced at + inference time in two steps: + + - Remove the filler from the **whole volume** with + `synapse_net.training.transform.remove_white_patches` before segmenting. It cannot be done in + the per-block `preprocess` of `synapse_net.inference.mitochondria.segment_mitochondria`, because + `get_prediction` standardizes a numpy input volume before it runs that, and the filler is + identified by its literal value, which no longer exists after standardizing. + - Pass `preprocess=torch_em.transform.raw.normalize_percentile` to `segment_mitochondria`. The + percentile normalization is invariant under that standardization, so the network sees the same + values it saw during training. + + Skipping the first step is not a small deviation: the filler skews the percentile normalization + of every tile it overlaps, which degrades the prediction on the surrounding tissue. On one block + that is 14% filler, leaving it in cost 0.04 Dice (0.79 instead of 0.84). See + `scripts/volume_em/inference/run_mitochondria_vol_em_segmentation.py` for a script that does both. + + Two training runs of this model exist and both are named 'final'. This function reproduces the run + with a patch shape of 32 x 512 x 512, a batch size of 4, without mixed precision and with the + filler removal. The other run used a patch shape of 64 x 256 x 256, a batch size of 8, mixed + precision and no filler removal; all other hyperparameters are the same. Both process the same + number of voxels per batch, so they trade field of view in-plane against field of view in z. + + Args: + name: The name for the checkpoint to be trained. + train_paths: Filepaths to the hdf5 files for the training data. + val_paths: Filepaths to the hdf5 files for the validation data. + save_root: Folder where the checkpoint will be saved. + raw_key: The key that holds the image data inside of the hdf5. + label_key: The key that holds the mitochondria annotations inside of the hdf5. + patch_shape: The patch shape used for a training example. + batch_size: The batch size for training. + lr: The initial learning rate. + n_iterations: The maximal number of iterations to train for. + n_samples_train: The number of train samples per epoch. + n_samples_val: The number of val samples per epoch. + initial_features: The number of features in the first level of the U-Net. + This has no effect if the model is initialized from `checkpoint_path`. + scale_factors: The downscaling factors for each level of the U-Net encoder, in ZYX. The + default matches the 5:1 anisotropy of the volume EM data. + This has no effect if the model is initialized from `checkpoint_path`. + norm: The normalization layer used in the convolutional blocks of the U-Net. The published + model was trained without normalization layers. + early_stopping: The number of epochs without improvement after which training is stopped. + log_image_interval: The interval (in iterations) at which images are written to the training log. + checkpoint_path: Path to a model checkpoint to initialize the weights from. Only the + weights are loaded; pass `resume` instead to continue a previous run. + resume: Whether to continue the previous run with this name in `save_root`, restoring + its weights, optimizer and iteration count. `n_iterations` then counts the total + iterations, including the ones the previous run already did. + overwrite: Whether to replace the checkpoints of a previous run with this name. + seed: The seed for the random number generators used in training. By default it is + not set, so that repeated runs differ. + deterministic: Whether to also disable cudnn benchmarking when a `seed` is given. + sampler: Sampler to accept or reject patches for training. + By default a minimum instance sampler with a rejection probability of 0.95 is used. + raw_transform: Transformation applied to the image data before it is passed to the network. + By default the white filler borders are removed and the block is then normalized with the + 1st and 99th percentile. Passing this overrides `fix_white_patches`. + fix_white_patches: Whether to remove the white filler borders of the blocks before normalizing + them. Disable this for data that was not cut out of a larger volume. + white_patch_min_size: The minimal size (in voxels) of a filler component that is removed. + num_workers: The number of workers for the dataloaders. + shuffle: Whether to shuffle the datasets in the dataloaders. This is disabled by default, + to match the published model; enabling it is recommended for new trainings, as it mixes + the blocks within a batch. + mixed_precision: Whether to train with mixed precision. This is disabled by default, to match + the published model; enabling it reduces the memory consumption and speeds up training. + check: Whether to check the training and validation loaders instead of running training. + kwargs: Additional keyword arguments for `synapse_net.training.supervised_training`. + """ + if sampler is None: + sampler = MinInstanceSampler(p_reject=0.95) + if raw_transform is None: + raw_transform = ( + RemoveWhitePatchesAndNormalize(min_size=white_patch_min_size) if fix_white_patches + else torch_em.transform.raw.normalize_percentile + ) + + supervised_training( + name=name, + train_paths=tuple(train_paths), + val_paths=tuple(val_paths), + label_key=label_key, + patch_shape=patch_shape, + save_root=save_root, + raw_key=raw_key, + batch_size=batch_size, + lr=lr, + n_iterations=n_iterations, + sampler=sampler, + n_samples_train=n_samples_train, + n_samples_val=n_samples_val, + check=check, + initial_features=initial_features, + scale_factors=scale_factors, + norm=norm, + checkpoint_path=checkpoint_path, + resume=resume, + overwrite=overwrite, + seed=seed, + deterministic=deterministic, + mixed_precision=mixed_precision, + early_stopping=early_stopping, + log_image_interval=log_image_interval, + raw_transform=raw_transform, + num_workers=num_workers, + shuffle=shuffle, + **kwargs, + ) + + +def main(): + """@private + """ + import argparse + + parser = argparse.ArgumentParser( + description="Train a model for mitochondria segmentation in volume electron microscopy.\n\n" + "The training data has to be stored in hdf5 files, one per block, which contain the image data and the\n" + "mitochondria annotations. The files are found recursively in the folders passed via '-i'. For example:\n" + "synapse_net.run_vol_em_mitochondria_training -n my_mito_model -i /path/to/blocks\n" + "The trained model will be saved in the folder 'checkpoints/my_mito_model'.\n" + "The default arguments reproduce the training of SynapseNet's volume EM mitochondria model, which was\n" + "trained on FIB-SEM data with a voxel size of 25 nm in z and 5 nm in xy. Note that this model was trained\n" + "with a batch size of 4 and a patch shape of 32 x 512 x 512, which requires a large GPU.\n" + "Check out the information below for details on the arguments of this function.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("-n", "--name", required=True, help="The name of the model to be trained.") + parser.add_argument("-i", "--data_root", nargs="+", help="One or more folders with the training data. They are searched recursively. Can be left out if 'split_file' contains the data roots.") # noqa + parser.add_argument("--file_pattern", default="*.h5", help="The pattern for selecting the files with training data.") # noqa + parser.add_argument("--raw_key", default="raw", help="The internal path for the image data.") + parser.add_argument("--label_key", default="labels/mitochondria", + help="The internal path for the mitochondria annotations.") + + # How to split the data into training and validation data. + parser.add_argument("--split_file", help="A json file with the keys 'train' and 'val', which each hold a list of entries '/'. If not given, the data is split randomly.") # noqa + parser.add_argument("--val_fraction", type=float, default=0.15, help="The fraction of the data to use for validation. This has no effect if 'split_file' was passed.") # noqa + parser.add_argument("--exclude", nargs="*", default=[], help="Substrings of filepaths to exclude from the training data.") # noqa + parser.add_argument("--seed", type=int, default=42, help="The seed for shuffling the files before splitting them, and for the random number generators used in training.") # noqa + + # The training hyperparameters. + parser.add_argument("-p", "--patch_shape", nargs=3, type=int, default=[32, 512, 512], + help="The patch shape for training in ZYX.") + parser.add_argument("--batch_size", type=int, default=4, help="The batch size for training.") + parser.add_argument("--learning_rate", type=float, default=1e-4, help="The initial learning rate.") + parser.add_argument("--n_iterations", type=int, default=int(5e4), help="The maximal number of iterations to train for.") # noqa + parser.add_argument("--n_samples_train", type=int, default=500, help="The number of samples per epoch for training.") # noqa + parser.add_argument("--n_samples_val", type=int, default=500, help="The number of samples per epoch for validation.") # noqa + parser.add_argument("--initial_features", type=int, default=32, help="The number of features in the first level of the U-Net.") # noqa + parser.add_argument("--early_stopping", type=int, default=20, help="The number of epochs without improvement after which training is stopped.") # noqa + parser.add_argument("--num_workers", type=int, default=8, help="The number of workers for the dataloaders.") + parser.add_argument("--no_white_patch_fix", action="store_true", help="Do not remove the white filler borders of the blocks before normalizing them. Use this for data that was not cut out of a larger volume.") # noqa + parser.add_argument("--white_patch_min_size", type=int, default=20, help="The minimal size (in voxels) of a white filler component that is removed.") # noqa + parser.add_argument("--shuffle", action="store_true", help="Shuffle the datasets in the dataloaders. This is recommended for new trainings, but was not used for the published model.") # noqa + parser.add_argument("--mixed_precision", action="store_true", help="Train with mixed precision. This reduces the memory consumption, but was not used for the published model.") # noqa + + # Where to save the model, and how to initialize it. + parser.add_argument("--save_root", help="Root path for saving the checkpoint and log dir.") + parser.add_argument("--checkpoint_path", help="A model checkpoint to initialize the weights from.") + parser.add_argument("--resume", action="store_true", help="Continue the previous run with the same name in 'save_root', restoring its weights, optimizer and iteration count. '--n_iterations' is the total number of iterations, including the ones already done.") # noqa + parser.add_argument("--overwrite", action="store_true", help="Replace the checkpoints of a previous run with the same name. By default training refuses to overwrite them.") # noqa + parser.add_argument("--deterministic", action="store_true", help="Disable cudnn benchmarking so that runs with the same seed match exactly. This costs throughput.") # noqa + parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.") # noqa + args = parser.parse_args() + + train_paths, val_paths = get_vol_em_mitochondria_paths( + args.data_root, file_pattern=args.file_pattern, val_fraction=args.val_fraction, + split_file=args.split_file, exclude=args.exclude, seed=args.seed, + ) + print("Training on", len(train_paths), "blocks and validating on", len(val_paths), "blocks.") + + vol_em_mitochondria_training( + name=args.name, train_paths=train_paths, val_paths=val_paths, save_root=args.save_root, + raw_key=args.raw_key, label_key=args.label_key, patch_shape=tuple(args.patch_shape), + batch_size=args.batch_size, lr=args.learning_rate, n_iterations=args.n_iterations, + n_samples_train=args.n_samples_train, n_samples_val=args.n_samples_val, + initial_features=args.initial_features, early_stopping=args.early_stopping, + checkpoint_path=args.checkpoint_path, resume=args.resume, overwrite=args.overwrite, + seed=args.seed, deterministic=args.deterministic, + fix_white_patches=not args.no_white_patch_fix, white_patch_min_size=args.white_patch_min_size, + num_workers=args.num_workers, shuffle=args.shuffle, + mixed_precision=args.mixed_precision, check=args.check, + ) diff --git a/synapse_net/training/models.py b/synapse_net/training/models.py index bc35201e..f580513d 100644 --- a/synapse_net/training/models.py +++ b/synapse_net/training/models.py @@ -1,6 +1,6 @@ """Model factories and pretrained backbone utilities for training.""" import warnings -from typing import Tuple +from typing import Optional, Tuple import numpy as np import torch @@ -34,6 +34,7 @@ def get_3d_model( scale_factors: Tuple[Tuple[int, int, int]] = [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]], initial_features: int = 32, final_activation: str = "Sigmoid", + norm: Optional[str] = "InstanceNorm", ) -> torch.nn.Module: """Get the U-Net model for 3D segmentation tasks. @@ -43,6 +44,8 @@ def get_3d_model( initial_features: The number of features in the first level of the U-Net. The number of features increases by a factor of two in each level. final_activation: The activation applied to the last output layer. + norm: The normalization layer used in the convolutional blocks. + Pass None to build the network without normalization layers. Returns: The U-Net. @@ -54,6 +57,7 @@ def get_3d_model( initial_features=initial_features, gain=2, final_activation=final_activation, + norm=norm, ) return model @@ -106,6 +110,7 @@ def _check_backbone(backbone, model_type, patch_shape): def _get_model( patch_shape, out_channels, in_channels=1, backbone=None, model_type=None, encoder_checkpoint=None, final_activation="Sigmoid", + initial_features=32, scale_factors=None, norm="InstanceNorm", ): is_2d, ndim = _determine_ndim(patch_shape) @@ -124,9 +129,16 @@ def _get_model( ) if is_2d: - return get_2d_model(out_channels=out_channels, in_channels=in_channels, final_activation=final_activation) + return get_2d_model( + out_channels=out_channels, in_channels=in_channels, initial_features=initial_features, + final_activation=final_activation, + ) - return get_3d_model(out_channels=out_channels, in_channels=in_channels, final_activation=final_activation) + scale_factor_kwargs = {} if scale_factors is None else {"scale_factors": scale_factors} + return get_3d_model( + out_channels=out_channels, in_channels=in_channels, initial_features=initial_features, + final_activation=final_activation, norm=norm, **scale_factor_kwargs, + ) def normalize_percentile_to_0_1_inverted(raw): diff --git a/synapse_net/training/split.py b/synapse_net/training/split.py new file mode 100644 index 00000000..04591d95 --- /dev/null +++ b/synapse_net/training/split.py @@ -0,0 +1,36 @@ +"""@private +""" + +import json +import os +from typing import Dict, Mapping, Sequence, Union + + +def _normalize_roots(data_roots: Union[Mapping[str, str], Sequence[str], str]) -> Dict[str, str]: + if isinstance(data_roots, str): + data_roots = [data_roots] + if not isinstance(data_roots, Mapping): + data_roots = {os.path.basename(root.rstrip("/")): root for root in data_roots} + return {name: root.rstrip("/") for name, root in data_roots.items()} + + +def _resolve_split(split_file, roots, keys): + """Resolve the '/' entries of a split file to filepaths.""" + with open(split_file) as f: + split = json.load(f) + # A split file may carry the data roots it was created with, so that it is self-contained. + # Roots passed by the caller take precedence, so that the data can be moved. + roots = {**_normalize_roots(split.get("roots", {})), **roots} + + resolved = [] + for key in keys: + paths = [] + for entry in split.get(key, []): + name, _, relative_path = entry.partition("/") + if name not in roots: + raise ValueError(f"The split file {split_file} refers to the unknown data root '{name}'.") + # The entries always use '/', so split on it to get native separators. Otherwise the + # paths are mixed on Windows and do not match the ones found by glob, e.g. in 'exclude'. + paths.append(os.path.join(roots[name], *relative_path.split("/"))) + resolved.append(paths) + return resolved diff --git a/synapse_net/training/supervised_training.py b/synapse_net/training/supervised_training.py index 626721ba..13574230 100644 --- a/synapse_net/training/supervised_training.py +++ b/synapse_net/training/supervised_training.py @@ -1,7 +1,10 @@ import os +import random from glob import glob +from shutil import copyfile from typing import Optional, Tuple, Union +import numpy as np import torch import torch_em from sklearn.model_selection import train_test_split @@ -12,6 +15,114 @@ from synapse_net.training.dataloaders import get_supervised_loader +def _checkpoint_folder(save_root, name): + """The folder that torch-em writes the checkpoints of this run to.""" + return os.path.join(save_root, "checkpoints", name) + + +def _set_seed(seed, deterministic=False): + """Seed the random number generators that affect training. + + Args: + seed: The seed for python's `random`, numpy and torch. Seeding alone does not make a + training run bit-exact on the GPU, because some cudnn kernels accumulate with atomics + and are non-deterministic regardless of the seed. Pass `deterministic` for that. + deterministic: Whether to also use deterministic kernels where torch has them, and pin the + cuBLAS workspace. Without this, two runs with the same seed still differ slightly, + because cudnn picks algorithms by timing. Note that torch has no deterministic backward + pass for the max pooling that the U-Net downsamples with, so this is requested with + 'warn_only' and bit-exact equality is not guaranteed on every GPU. It costs throughput. + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + if deterministic: + torch.backends.cudnn.benchmark = False + # cuBLAS needs this to be set before the CUDA context is created to be reproducible. + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True, warn_only=True) + + +def _check_overwrite(save_root, name, overwrite, resume): + """Refuse to start a run that would replace the checkpoints of a previous one. + + torch-em writes into '/checkpoints/' without any check, so a second run with + the same name silently replaces the checkpoints of the first. Its own `overwrite_training` does + not help: that only skips runs which already reached exactly the requested iteration count. + + Raises: + ValueError: If a previous run with this name exists and neither `overwrite` nor `resume` + was asked for. + """ + if save_root is None or overwrite or resume: + return + + previous_run = _checkpoint_folder(save_root, name) + if os.path.exists(os.path.join(previous_run, "latest.pt")): + raise ValueError( + f"The training run '{name}' already has checkpoints in {previous_run}. Training would " + "overwrite them. Pass a different name, or 'overwrite=True' ('--overwrite' on the " + "command line) to replace them, or 'resume=True' ('--resume') to continue that run." + ) + + +def _prepare_resume(save_root, name, n_iterations, mixed_precision, verbose=True): + """Work out how to continue a previous run, and how many iterations are left for it. + + Two details of torch-em are handled here. It only writes the scaler state if the previous run + used mixed precision, but reads it unconditionally when the new trainer has one, so a mismatch + is caught up front instead of failing with a `KeyError` mid-run. And it resets the best metric + when it initializes the trainer, so the first validation of the resumed run would overwrite + 'best.pt' even if it is worse; the previous best is copied aside before that can happen. + + Returns: + The argument for `trainer.fit(load_from_checkpoint=...)`. torch-em joins it with '.pt', so + it is the checkpoint path without the extension. A path that does not resolve is only + warned about and training then silently starts from scratch, so the file is checked to + exist before its stem is returned. + The number of iterations left to reach `n_iterations` in total. + + Raises: + ValueError: If there is no previous run, if it already reached `n_iterations`, or if it + used a different mixed precision setting. + """ + checkpoint_folder = _checkpoint_folder(save_root, name) + checkpoint_path = os.path.join(checkpoint_folder, "latest.pt") + if not os.path.exists(checkpoint_path): + raise ValueError( + f"Cannot resume the training run '{name}': {checkpoint_path} does not exist." + ) + + state = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + iteration = state["iteration"] + remaining = n_iterations - iteration + if remaining <= 0: + raise ValueError( + f"The training run '{name}' is already at iteration {iteration} of the requested " + f"{n_iterations}. Increase 'n_iterations' to train it further." + ) + + if mixed_precision and "scaler_state" not in state: + raise ValueError( + f"The training run '{name}' was trained without mixed precision, so it cannot be " + "resumed with mixed precision. Pass 'mixed_precision=False' to continue it." + ) + + best_path = os.path.join(checkpoint_folder, "best.pt") + if os.path.exists(best_path): + backup_path = os.path.join(checkpoint_folder, "best-before-resume.pt") + copyfile(best_path, backup_path) + if verbose: + print("The previous best checkpoint was copied to", backup_path) + + if verbose: + print(f"Resuming '{name}' from iteration {iteration}, training for {remaining} more.") + + return os.path.splitext(checkpoint_path)[0], remaining + + def supervised_training( name: str, train_paths: Tuple[str], @@ -37,12 +148,23 @@ def supervised_training( loss_fn: Optional[torch.nn.Module] = None, in_channels: int = 1, out_channels: int = 2, + initial_features: int = 32, + scale_factors: Optional[Tuple[Tuple[int, int, int]]] = None, + norm: Optional[str] = "InstanceNorm", + transform: Optional[callable] = None, mask_channel: bool = False, backbone: Optional[str] = None, model_type: Optional[str] = None, encoder_checkpoint: Optional[str] = None, checkpoint_path: Optional[str] = None, + resume: bool = False, + overwrite: bool = False, + seed: Optional[int] = None, + deterministic: bool = False, save_every_kth_epoch: Optional[int] = None, + mixed_precision: bool = True, + early_stopping: Optional[int] = None, + log_image_interval: int = 100, **loader_kwargs, ): """Run supervised segmentation training. @@ -85,7 +207,19 @@ def supervised_training( label_transform: Label transform that is applied to the segmentation to compute the targets. If no label transform is passed (the default) a boundary transform is used. loss_fn: Custom loss function. If None, will default to `torch_em.loss.DiceLoss`. + in_channels: The number of input channels of the UNet. Use a value larger than one to train + on data with multiple channels, and pass `with_channels=True` so that the loader reads them. out_channels: The number of output channels of the UNet. + initial_features: The number of features in the first level of the UNet. + The number of features increases by a factor of two in each level. + This has no effect if a pretrained `backbone` is used. + scale_factors: The downscaling factors for each level of the UNet encoder, in ZYX. By default + the levels of `get_3d_model` are used. This has no effect for 2d training, if a pretrained + `backbone` is used, or if the model is initialized from `checkpoint_path`. + norm: The normalization layer used in the convolutional blocks of the 3d UNet. + Pass None to train a network without normalization layers. + transform: Joint transformation applied to the raw data and the labels. + By default padding and the standard augmentations are applied. mask_channel: Whether the last channels in the labels should be used for masking the loss. This can be used to implement more complex masking operations and is not compatible with `ignore_label`. backbone: The pretrained encoder. Options: "sam", "sam2", "dinov2", "dinov3" (ViT encoders of a UNETR), @@ -95,12 +229,33 @@ def supervised_training( "vit_b_em_organelles" for "sam". encoder_checkpoint: Path to the pretrained encoder weights for `backbone`. By default they are fetched for `backbone` and `model_type`. - checkpoint_path: Path to the directory with 'best.pt'; continue training this model. + checkpoint_path: Path to a checkpoint to initialize the weights of a new run from. Only the + weights are loaded, the run starts at iteration zero with a fresh optimizer. Pass + `resume` instead to continue a previous run of the same name. + resume: Whether to continue the previous training run with this name in `save_root`, + restoring its weights, optimizer, learning rate schedule and iteration count. + `n_iterations` is the total number of iterations to train for, including the ones that + the previous run already did. + overwrite: Whether to replace the checkpoints of a previous training run with this name in + `save_root`. By default training refuses to overwrite them. + seed: The seed for the random number generators used in training, for example for the + weight initialization and the augmentations. By default the seed is not set, so that + repeated runs differ. + deterministic: Whether to also disable the cudnn benchmarking when a `seed` is given. This + makes repeated runs on the same hardware match exactly, at the cost of throughput. save_every_kth_epoch: Save checkpoints after every kth epoch in a separate file. The corresponding checkpoints will be saved with the naming scheme 'epoch-{epoch}.pt'. + mixed_precision: Whether to train with mixed precision. + early_stopping: The number of epochs without improvement after which training is stopped. + By default training runs for the full number of iterations. + log_image_interval: The interval (in iterations) at which images are written to the training log. loader_kwargs: Additional keyword arguments for the dataloader. """ _check_backbone(backbone, model_type, patch_shape) + # Before the loaders, so that this fails immediately instead of after the dataset construction. + _check_overwrite(save_root, name, overwrite, resume) + if seed is not None: + _set_seed(seed, deterministic=deterministic) if raw_transform is None: raw_transform = get_raw_transform(backbone)[0] @@ -109,13 +264,13 @@ def supervised_training( train_paths, raw_key, label_key, patch_shape, batch_size, n_samples=n_samples_train, rois=train_rois, sampler=sampler, ignore_label=ignore_label, label_transform=label_transform, label_paths=train_label_paths, raw_transform=raw_transform, - **loader_kwargs, + transform=transform, **loader_kwargs, ) val_loader = get_supervised_loader( val_paths, raw_key, label_key, patch_shape, batch_size, n_samples=n_samples_val, rois=val_rois, sampler=sampler, ignore_label=ignore_label, label_transform=label_transform, label_paths=val_label_paths, raw_transform=raw_transform, - **loader_kwargs, + transform=transform, **loader_kwargs, ) if check: @@ -124,12 +279,18 @@ def supervised_training( check_loader(val_loader, n_samples=4) return - if checkpoint_path is not None: + # When resuming, 'fit' loads the weights together with the rest of the trainer state below. + load_from_checkpoint, n_iterations = ( + _prepare_resume(save_root, name, n_iterations, mixed_precision) if resume else (None, n_iterations) + ) + + if checkpoint_path is not None and not resume: model = _restore_legacy_attributes(torch_em.util.load_model(checkpoint=checkpoint_path)) else: model = _get_model( patch_shape, out_channels, in_channels=in_channels, backbone=backbone, model_type=model_type, encoder_checkpoint=encoder_checkpoint, + initial_features=initial_features, scale_factors=scale_factors, norm=norm, ) base_loss = loss_fn if loss_fn is not None else torch_em.loss.DiceLoss() @@ -166,14 +327,16 @@ def supervised_training( train_loader=train_loader, val_loader=val_loader, learning_rate=lr, - mixed_precision=True, - log_image_interval=100, + mixed_precision=mixed_precision, + log_image_interval=log_image_interval, + early_stopping=early_stopping, compile_model=False, save_root=save_root, loss=loss, metric=metric, ) - trainer.fit(n_iterations, save_every_kth_epoch=save_every_kth_epoch) + trainer.fit(n_iterations, load_from_checkpoint=load_from_checkpoint, + save_every_kth_epoch=save_every_kth_epoch) def _derive_key_from_files(files, key): diff --git a/synapse_net/training/transform.py b/synapse_net/training/transform.py index 04bad88e..816fe62d 100644 --- a/synapse_net/training/transform.py +++ b/synapse_net/training/transform.py @@ -1,6 +1,15 @@ +"""Transformations applied to the data and labels during training. +""" + import numpy as np +import torch +import torch_em from torch_em.transform.label import labels_to_binary from scipy.ndimage import distance_transform_edt +from skimage.measure import label as connected_components + +# The voxel size of the cristae training data, in nanometer. Used when no voxel size is given. +CRISTAE_VOXEL_SIZE = (1.74, 1.74, 1.74) class AZDistanceLabelTransform: @@ -16,3 +25,384 @@ def __call__(self, input_): distances = np.clip(distances, 0.0, self.max_distance) distances /= self.max_distance return np.stack([binary_target, distances]) + + +def standardize_channel(raw: np.ndarray, channel: int = 0) -> np.ndarray: + """Standardize a single channel of a multi-channel input, leaving the other channels unchanged. + + This is used for training on data where one channel holds the image data and another one holds a + semantic mask, which has to keep its integer values so that it can be compared against them. + + Args: + raw: The multi-channel input data, with the channels in the first axis. + channel: The channel to standardize. + + Returns: + The input data with the given channel standardized. + + Raises: + ValueError: If the input is not 4D or if the channel does not exist. + """ + if raw.ndim != 4: + raise ValueError(f"Expect input data with 4 dimensions (channels, z, y, x), got {raw.ndim}.") + if channel >= raw.shape[0]: + raise ValueError(f"Invalid channel {channel} for input data with {raw.shape[0]} channels.") + raw = np.float32(raw) + raw[channel] = torch_em.transform.raw.standardize(raw[channel]) + return raw + + +def normalize_channel(raw: np.ndarray, channel: int = 0, lower: float = 1.0, upper: float = 99.0) -> np.ndarray: + """Percentile-normalize a single channel of a multi-channel input, leaving the others unchanged. + + See `standardize_channel` for why only one channel is normalized. + + Args: + raw: The multi-channel input data, with the channels in the first axis. + channel: The channel to normalize. + lower: The lower percentile. + upper: The upper percentile. + + Returns: + The input data with the given channel normalized to the range [0, 1]. + + Raises: + ValueError: If the input is not 4D or if the channel does not exist. + """ + if raw.ndim != 4: + raise ValueError(f"Expect input data with 4 dimensions (channels, z, y, x), got {raw.ndim}.") + if channel >= raw.shape[0]: + raise ValueError(f"Invalid channel {channel} for input data with {raw.shape[0]} channels.") + raw = np.float32(raw) + normalized = torch_em.transform.raw.normalize_percentile(raw[channel], lower=lower, upper=upper) + raw[channel] = np.clip(normalized, 0, 1) + return raw + + +def remove_white_patches(raw: np.ndarray, min_size: int = 20, value: int = 255) -> np.ndarray: + """Set the connected components of saturated voxels in a volume EM block to zero. + + Volume EM blocks that were cut out of a larger FIB-SEM volume carry white filler borders where the + cutout extends past the imaged region. They take up a large part of some blocks, so they skew the + percentile normalization of the data around them. Electron tomograms do not have these borders, + this is only meant for volume EM data. + + Small saturated components are kept, because they are saturated sample structure rather than + filler. The components are determined with full connectivity, so that filler that is only + connected diagonally still counts as one component. + + Args: + raw: The volume, with the values of the filler given by `value`. + min_size: The minimal size (in voxels) of a component that is removed. + value: The value that marks the filler. + + Returns: + The volume with the filler set to zero. The input is returned unchanged if it holds no filler. + + Raises: + ValueError: If the volume is not 2d or 3d, or if it does not have an integer dtype. + """ + if raw.ndim not in (2, 3): + raise ValueError(f"Expect 2d or 3d input data, got {raw.ndim} dimensions.") + # The filler is identified by an exact value, so this has to run on the unnormalized data. + # Normalized or standardized data would silently contain no filler at all. + if not np.issubdtype(raw.dtype, np.integer): + raise ValueError( + f"Expect input data with an integer dtype, got {raw.dtype}. This transform has to be applied " + "to the raw data, before it is normalized." + ) + + # This runs once per patch in every dataloader worker, and most patches do not touch the border. + mask = raw == value + if not mask.any(): + return raw + + components = connected_components(mask) + sizes = np.bincount(components.ravel()) + filler_ids = np.where(sizes >= min_size)[0] + filler_ids = filler_ids[filler_ids != 0] + if filler_ids.size == 0: + return raw + + raw = raw.copy() + raw[np.isin(components, filler_ids)] = 0 + return raw + + +class RemoveWhitePatchesAndNormalize: + """Remove the white filler borders of a volume EM block, then percentile-normalize it. + + This is the **training** raw transform of SynapseNet's volume EM mitochondria model, see + `synapse_net.training.mitochondria_vol_em`. It is applied per patch, to the unnormalized data that + the dataset reads. It is a class rather than a function so that it can be pickled by reference: + the dataloader workers and the training checkpoint both require that. + + Do not pass it as the `preprocess` of an inference function. `synapse_net.inference.util.get_prediction` + standardizes a numpy input volume before `preprocess` runs, so the filler would no longer have the + value that identifies it. At inference, call `remove_white_patches` on the whole volume first and + then normalize with `torch_em.transform.raw.normalize_percentile`. + + Args: + min_size: The minimal size (in voxels) of a filler component that is removed. + lower: The lower percentile for the normalization. + upper: The upper percentile for the normalization. + """ + + def __init__(self, min_size: int = 20, lower: float = 1.0, upper: float = 99.0): + self.min_size = min_size + self.lower = lower + self.upper = upper + + def __call__(self, raw: np.ndarray) -> np.ndarray: + raw = remove_white_patches(raw, min_size=self.min_size) + return torch_em.transform.raw.normalize_percentile(raw, lower=self.lower, upper=self.upper) + + +def _fast_distance_transform(mask, sampling): + # This runs once per patch in every dataloader worker, so it is worth using the faster + # implementation from bioimage-cpp when it is available (measured ~5x faster than scipy on a + # padded 32 x 256 x 256 patch, agreeing to within float32 rounding). + # It is called single-threaded on purpose: the dataloader workers are already parallel. + try: + from bioimage_cpp.distance import distance_transform + except ImportError: + return distance_transform_edt(mask, sampling=sampling) + return distance_transform(mask, sampling=list(sampling), number_of_threads=1) + + +def proximity_band( + mito_mask: np.ndarray, + voxel_size=CRISTAE_VOXEL_SIZE, + band_nm: float = 12.0, + offset_nm: float = 0.0, +) -> np.ndarray: + """Compute the shell inside `mito_mask` with a distance to its surface in the given band. + + The shell contains the voxels with `offset_nm < distance <= offset_nm + band_nm`. Since the + mitochondrial membrane is not annotated in the training data, the surface of the mitochondria + mask is used to approximate it. + + This is meant to run on a training patch, so it corrects for the patch borders: a mitochondrion + that is cut by a patch face would otherwise be seen as ending there, and the distance transform + would introduce a membrane band through the lumen of the mitochondrion. Padding the mask in + 'edge' mode before the transform lets such a mitochondrion continue instead. Where the face is + background, the padding does not change anything. + + Args: + mito_mask: The binary mitochondria mask. + voxel_size: The voxel size in nanometer, in the order (z, y, x). + band_nm: The thickness of the shell in nanometer. + offset_nm: Start the shell this far inside the membrane. Zero starts it at the mask surface. + + Returns: + The binary shell, with the same shape as `mito_mask`. + """ + mito_mask = np.asarray(mito_mask, dtype=bool) + if not mito_mask.any(): + return np.zeros(mito_mask.shape, dtype=bool) + + lower = float(offset_nm) + upper = lower + float(band_nm) + pad = int(np.ceil(upper / float(min(voxel_size)))) + 2 + distances = _fast_distance_transform(np.pad(mito_mask, pad, mode="edge"), voxel_size) + distances = distances[tuple(slice(pad, -pad) for _ in range(mito_mask.ndim))] + return mito_mask & (distances > lower) & (distances <= upper) + + +def membrane_proximity_weight( + mito_state: np.ndarray, + gt_foreground: np.ndarray, + voxel_size=CRISTAE_VOXEL_SIZE, + band_nm: float = 12.0, + offset_nm: float = 0.0, + w_pos: float = 1.0, + w_neg: float = 1.0, + annotated_state: int = 1, + exclude_state_value: float = 2.0, +) -> np.ndarray: + """Compute a per-voxel loss weight that emphasizes the cristae close to the mitochondria membrane. + + This generalizes the binary loss mask used for cristae training:: + + w = 0.0 where the state is `exclude_state_value`, i.e. an unannotated mitochondrion + w = 1.0 everywhere else + w = w_pos inside the membrane band and in the cristae ground-truth + w = w_neg inside the membrane band and outside of the cristae ground-truth + + Splitting the weight into `w_pos` and `w_neg` is the point of this function: a single band weight + would amplify the cristae and the membrane equally. `w_pos` pushes the model towards predicting + cristae at the junction with the membrane, `w_neg` separately controls how strongly the membrane + itself is pushed down. + + With `w_pos == w_neg == 1.0` the result is exactly the binary loss mask, so the unweighted recipe + is not affected by this function and no distance transform is computed. + + Args: + mito_state: The semantic mitochondria state, where 0 is background, 1 is a mitochondrion with + cristae annotations, and 2 is a mitochondrion without cristae annotations. + gt_foreground: The binary cristae ground-truth, with the same shape as `mito_state`. + voxel_size: The voxel size in nanometer, in the order (z, y, x). + band_nm: The thickness of the membrane band in nanometer. + offset_nm: Start the band this far inside the membrane. + w_pos: The weight for cristae voxels inside the band. + w_neg: The weight for non-cristae voxels inside the band. + annotated_state: The state value of mitochondria that carry cristae annotations. + exclude_state_value: The state value that is excluded from the loss. + + Returns: + The per-voxel weight, with the same shape as `mito_state`. + """ + mito_state = np.asarray(mito_state) + # One where the state is NOT excluded from the loss. Note that this keeps the background, so that + # the model still learns that there are no cristae outside of mitochondria. + weight = (np.abs(mito_state - float(exclude_state_value)) >= 0.5).astype(np.float32) + + w_pos, w_neg = float(w_pos), float(w_neg) + if w_pos == 1.0 and w_neg == 1.0: + return weight + + band = proximity_band(mito_state == annotated_state, voxel_size, band_nm, offset_nm) + if not band.any(): + return weight + + foreground = np.asarray(gt_foreground) > 0 + # Multiply rather than assign, so that excluded voxels keep a weight of zero. + if w_pos != 1.0: + weight[band & foreground] *= np.float32(w_pos) + if w_neg != 1.0: + weight[band & ~foreground] *= np.float32(w_neg) + return weight + + +class MitoStateMaskTransform: + """Joint transform that appends the cristae loss mask to the labels. + + After the label transform has produced labels of shape (n_channels, z, y, x), this appends + n_channels mask channels that encode where the loss should be computed, resulting in labels of + shape (2 * n_channels, z, y, x) as expected by + `synapse_net.training.loss.MaskedDiceLossPerSample`. + + Voxels where the mitochondria state channel of the raw data equals `exclude_state_value` are + masked out, so that the network is not penalized for predicting cristae inside mitochondria that + carry no cristae annotations. All other voxels, including the background, stay active. + + The mask can also become a weight map that emphasizes the membrane proximity band, see + `membrane_proximity_weight`. With the default `w_pos == w_neg == 1.0` it is exactly the binary + mask. Note that this transform replaces the augmentations in the dataloader; use + `AugmentedMitoStateMaskTransform` to keep them. + + Args: + mito_channel: The channel of the raw data that holds the mitochondria state. + exclude_state_value: The state value that is excluded from the loss. + band_nm: The thickness of the membrane band in nanometer. + offset_nm: Start the band this far inside the membrane. + w_pos: The weight for cristae voxels inside the band. + w_neg: The weight for non-cristae voxels inside the band. + voxel_size: The voxel size in nanometer, in the order (z, y, x). + """ + + def __init__( + self, + mito_channel: int = 1, + exclude_state_value: float = 2.0, + band_nm: float = 12.0, + offset_nm: float = 0.0, + w_pos: float = 1.0, + w_neg: float = 1.0, + voxel_size=CRISTAE_VOXEL_SIZE, + ): + self.mito_channel = mito_channel + self.exclude_state_value = exclude_state_value + self.band_nm = float(band_nm) + self.offset_nm = float(offset_nm) + self.w_pos = float(w_pos) + self.w_neg = float(w_neg) + self.voxel_size = tuple(float(vs) for vs in voxel_size) + + @property + def weighted(self) -> bool: + """Whether a membrane proximity weight is applied instead of the plain binary mask.""" + return self.w_pos != 1.0 or self.w_neg != 1.0 + + def _weight(self, mito_state, gt_foreground): + return membrane_proximity_weight( + mito_state, gt_foreground, voxel_size=self.voxel_size, band_nm=self.band_nm, + offset_nm=self.offset_nm, w_pos=self.w_pos, w_neg=self.w_neg, + exclude_state_value=self.exclude_state_value, + ) + + def __call__(self, raw, labels): + """Append the loss mask to the labels. + + Args: + raw: The raw data, with the mitochondria state in `mito_channel`. + labels: The labels produced by the label transform. + + Returns: + The unchanged raw data. + The labels with the loss mask appended to the channel axis. + """ + mito_state = raw[self.mito_channel] + if self.weighted: + # The first label channel is the binary cristae channel from the boundary transform. + mask = self._weight(mito_state, labels[0]) + else: + mask = (np.abs(mito_state - self.exclude_state_value) >= 0.5).astype(np.float32) + masks = np.stack([mask] * labels.shape[0], axis=0) + return raw, np.concatenate([labels, masks], axis=0) + + +class AugmentedMitoStateMaskTransform(MitoStateMaskTransform): + """`MitoStateMaskTransform` that applies an augmentation before computing the loss mask. + + A joint transform occupies the slot of the dataloader that otherwise holds the augmentations, so + `MitoStateMaskTransform` on its own trains without any augmentation. This class restores them: it + augments the raw data and labels first and then computes the loss mask from the augmented state + channel, so that the mask matches the data the network sees. + + Args: + augmentation: The joint augmentation for the raw data and labels. Pass + `torch_em.transform.get_augmentations(ndim)`, which is what the dataloader uses by default. + kwargs: Additional arguments for `MitoStateMaskTransform`. + """ + + def __init__(self, augmentation, **kwargs): + super().__init__(**kwargs) + self.augmentation = augmentation + + def _weight_torch(self, mito_state, gt_foreground, channel_axis): + # The default augmentations are flips only, so the state channel stays integral and the + # distance transform runs on the same geometry the network sees. + state = mito_state.detach().cpu().numpy() + foreground = gt_foreground.detach().cpu().numpy() + if channel_axis == 1: # There is a leading batch axis, so the transform has to run per sample. + return np.stack([self._weight(state[i], foreground[i]) for i in range(state.shape[0])], axis=0) + return self._weight(state, foreground) + + def __call__(self, raw, labels): + """Augment the raw data and labels, then append the loss mask to the labels. + + Args: + raw: The raw data, with the mitochondria state in `mito_channel`. + labels: The labels produced by the label transform. + + Returns: + The augmented raw data. + The augmented labels with the loss mask appended to the channel axis. + """ + raw, labels = self.augmentation(raw, labels) + raw, labels = torch.as_tensor(raw), torch.as_tensor(labels) + # The augmentations return the data with a leading batch axis, so the channel axis is one. + # Support data without it as well, where the channel axis is zero. + channel_axis = 1 if raw.dim() >= 5 else 0 + mito_state = raw.narrow(channel_axis, self.mito_channel, 1).squeeze(channel_axis).float() + if self.weighted: + # The first label channel is the binary cristae channel from the boundary transform. + foreground = labels.narrow(channel_axis, 0, 1).squeeze(channel_axis) + mask = torch.as_tensor(self._weight_torch(mito_state, foreground, channel_axis)).to(labels.dtype) + else: + mask = (torch.abs(mito_state - self.exclude_state_value) >= 0.5).to(labels.dtype) + repeats = [1] * labels.dim() + repeats[channel_axis] = labels.shape[channel_axis] + masks = mask.unsqueeze(channel_axis).repeat(*repeats) + return raw, torch.cat([labels, masks], dim=channel_axis) diff --git a/test/test_cristae_training.py b/test/test_cristae_training.py new file mode 100644 index 00000000..845822d0 --- /dev/null +++ b/test/test_cristae_training.py @@ -0,0 +1,337 @@ +import json +import os +import platform +import unittest +from shutil import rmtree +from subprocess import run + +import h5py +import numpy as np +import torch +from skimage.data import binary_blobs +from unittest import mock + + +class TestCristaeTraining(unittest.TestCase): + tmp_folder = "./tmp_cristae_training" + roots = {"ds0": os.path.join(tmp_folder, "ds0"), "ds1": os.path.join(tmp_folder, "ds1")} + + def setUp(self): + rng = np.random.default_rng(42) + for i in range(8): + root = self.roots["ds0"] if i % 2 == 0 else self.roots["ds1"] + path = os.path.join(root, "sub", f"tomo-{i}_combined.h5") + os.makedirs(os.path.dirname(path), exist_ok=True) + + raw = rng.normal(size=(32, 64, 64)).astype("float32") + # The mitochondria state: 0 = background, 1 = annotated mito, 2 = unannotated mito. + state = np.zeros((32, 64, 64), dtype="float32") + state[:, :32] = 1 + state[:, 32:48] = 2 + cristae = binary_blobs(length=64, n_dim=3, volume_fraction=0.2)[:32].astype("uint8") + + with h5py.File(path, "w") as f: + f.create_dataset("raw_mitos_combined", data=np.stack([raw, state])) + f.create_dataset("labels/cristae", data=cristae) + + def tearDown(self): + try: + rmtree(self.tmp_folder) + except OSError: + pass + + # Data discovery and splitting. + + def test_get_cristae_paths(self): + from synapse_net.training.cristae import get_cristae_paths + + train_paths, val_paths = get_cristae_paths(self.roots) + self.assertEqual(len(train_paths), 7) + self.assertEqual(len(val_paths), 1) + self.assertEqual(len(set(train_paths) & set(val_paths)), 0) + for path in train_paths + val_paths: + self.assertTrue(os.path.exists(path)) + + # The split has to be deterministic for a fixed seed. + self.assertEqual(get_cristae_paths(self.roots), (train_paths, val_paths)) + + def test_get_cristae_paths_with_exclude(self): + from synapse_net.training.cristae import get_cristae_paths + + train_paths, val_paths = get_cristae_paths(self.roots, exclude=["tomo-0_", "tomo-1_"]) + self.assertEqual(len(train_paths) + len(val_paths), 6) + self.assertFalse(any("tomo-0_" in path or "tomo-1_" in path for path in train_paths + val_paths)) + + def test_get_cristae_paths_from_split_file(self): + from synapse_net.training.cristae import get_cristae_paths + + split = {"train": ["ds0/sub/tomo-0_combined.h5", "ds1/sub/tomo-1_combined.h5"], + "val": ["ds0/sub/tomo-2_combined.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + train_paths, val_paths = get_cristae_paths(self.roots, split_file=split_file) + self.assertEqual(train_paths, [os.path.join(self.tmp_folder, *name.split("/")) for name in split["train"]]) + self.assertEqual(val_paths, [os.path.join(self.tmp_folder, *name.split("/")) for name in split["val"]]) + + def test_get_cristae_paths_from_split_file_with_roots(self): + from synapse_net.training.cristae import get_cristae_paths + + # A split file that carries its own roots can be used without passing them. + split = {"roots": self.roots, + "train": ["ds0/sub/tomo-0_combined.h5"], "val": ["ds1/sub/tomo-1_combined.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + train_paths, val_paths = get_cristae_paths(split_file=split_file) + self.assertEqual(train_paths, [os.path.join(self.roots["ds0"], "sub", "tomo-0_combined.h5")]) + self.assertEqual(val_paths, [os.path.join(self.roots["ds1"], "sub", "tomo-1_combined.h5")]) + + # Roots passed explicitly take precedence, so that the data can be moved. + moved = os.path.join(self.tmp_folder, "moved") + os.makedirs(os.path.join(moved, "sub"), exist_ok=True) + open(os.path.join(moved, "sub", "tomo-0_combined.h5"), "w").close() + train_paths, _ = get_cristae_paths({"ds0": moved}, split_file=split_file) + self.assertEqual(train_paths, [os.path.join(moved, "sub", "tomo-0_combined.h5")]) + + def test_get_cristae_paths_without_roots(self): + from synapse_net.training.cristae import get_cristae_paths + + with self.assertRaisesRegex(ValueError, "No data roots"): + get_cristae_paths() + + def test_get_cristae_paths_from_invalid_split_file(self): + from synapse_net.training.cristae import get_cristae_paths + + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump({"train": ["ds0/sub/tomo-0_combined.h5"], "val": ["ds0/does-not-exist.h5"]}, f) + with self.assertRaisesRegex(ValueError, "do not exist"): + get_cristae_paths(self.roots, split_file=split_file) + + with open(split_file, "w") as f: + json.dump({"train": ["nope/tomo-0_combined.h5"], "val": []}, f) + with self.assertRaisesRegex(ValueError, "unknown data root"): + get_cristae_paths(self.roots, split_file=split_file) + + def test_get_cristae_paths_without_data(self): + from synapse_net.training.cristae import get_cristae_paths + + with self.assertRaisesRegex(ValueError, "Did not find any files"): + get_cristae_paths(self.roots, file_pattern="*.mrc") + + def test_get_cristae_test_paths(self): + from synapse_net.training.cristae import get_cristae_paths, get_cristae_test_paths + + split = {"train": ["ds0/sub/tomo-0_combined.h5"], "val": ["ds1/sub/tomo-1_combined.h5"], + "test": ["ds0/sub/tomo-2_combined.h5", "ds1/sub/tomo-3_combined.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + test_paths = get_cristae_test_paths(split_file, self.roots) + self.assertEqual(test_paths, [os.path.join(self.roots[e.split("/")[0]], *e.split("/")[1:]) + for e in split["test"]]) + + # A random split must be able to keep those volumes out, otherwise it trains on them. + train_paths, val_paths = get_cristae_paths(self.roots) + self.assertTrue(set(test_paths) & set(train_paths + val_paths)) + + train_paths, val_paths = get_cristae_paths(self.roots, exclude=test_paths) + self.assertFalse(set(test_paths) & set(train_paths + val_paths)) + + def test_get_cristae_test_paths_without_test_key(self): + from synapse_net.training.cristae import get_cristae_test_paths + + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump({"train": ["ds0/sub/tomo-0_combined.h5"], "val": ["ds1/sub/tomo-1_combined.h5"]}, f) + self.assertEqual(get_cristae_test_paths(split_file, self.roots), []) + + # The masked losses. + + def test_masked_dice_loss_per_sample(self): + from synapse_net.training.loss import MaskedDiceLossPerSample + + loss = MaskedDiceLossPerSample() + prediction = torch.rand(3, 2, 4, 8, 8) + target = (torch.rand(3, 2, 4, 8, 8) > 0.5).float() + + # With an all-one mask this is the plain per-sample dice. + full_mask = torch.cat([target, torch.ones_like(target)], dim=1) + expected = 0.0 + for channel in range(2): + dice = 0.0 + for sample in range(3): + pred, tgt = prediction[sample, channel], target[sample, channel] + dice += 2 * (pred * tgt).sum() / ((pred * pred).sum() + (tgt * tgt).sum()) + expected += 1.0 - dice / 3 + self.assertAlmostEqual(loss(prediction, full_mask).item(), expected.item(), places=5) + + # Masked-out voxels must not influence the loss: changing the prediction there is a no-op. + mask = torch.ones_like(target) + mask[..., :4] = 0 + masked_target = torch.cat([target, mask], dim=1) + reference = loss(prediction, masked_target) + perturbed = prediction.clone() + perturbed[..., :4] = 1.0 - perturbed[..., :4] + self.assertAlmostEqual(loss(perturbed, masked_target).item(), reference.item(), places=5) + + # A target with the wrong number of channels is rejected. + with self.assertRaisesRegex(ValueError, "expects a target"): + loss(prediction, target) + + def test_masked_dice_loss_weighting(self): + from synapse_net.training.loss import MaskedDiceLossPerSample + + # A constant weight w has to cancel out of the dice, since it scales the numerator and the + # denominator alike. This is what the sqrt of the mask is for. + loss = MaskedDiceLossPerSample() + prediction = torch.rand(2, 1, 4, 8, 8) + target = (torch.rand(2, 1, 4, 8, 8) > 0.5).float() + unweighted = loss(prediction, torch.cat([target, torch.ones_like(target)], dim=1)) + weighted = loss(prediction, torch.cat([target, 3.0 * torch.ones_like(target)], dim=1)) + self.assertAlmostEqual(unweighted.item(), weighted.item(), places=5) + + # The membrane weighting and the mask transforms. + + def test_membrane_proximity_weight(self): + from synapse_net.training.transform import membrane_proximity_weight + + state = np.zeros((16, 32, 32), dtype="float32") + state[4:12, 8:24, 8:24] = 1 # an annotated mitochondrion + state[:2] = 2 # an unannotated mitochondrion + foreground = np.zeros_like(state) + foreground[4:12, 10:22, 10:22] = 1 + + # Unit weights give exactly the binary mask, without computing a distance transform. + unweighted = membrane_proximity_weight(state, foreground, w_pos=1.0, w_neg=1.0) + np.testing.assert_array_equal(unweighted, (state != 2).astype("float32")) + + weighted = membrane_proximity_weight(state, foreground, w_pos=3.0, w_neg=2.0, band_nm=12.0) + # Excluded voxels stay at zero whatever the weights are. + self.assertTrue((weighted[state == 2] == 0).all()) + # Only the weights 1, w_pos and w_neg occur, and both weights are actually applied. + self.assertEqual(set(np.unique(weighted).tolist()), {0.0, 1.0, 2.0, 3.0}) + + def test_mito_state_mask_transform(self): + from synapse_net.training.transform import MitoStateMaskTransform + + raw = np.zeros((2, 8, 16, 16), dtype="float32") + raw[1, :, :8] = 1 + raw[1, :, 8:12] = 2 + labels = np.zeros((2, 8, 16, 16), dtype="float32") + labels[0, :, 2:6] = 1 + + transform = MitoStateMaskTransform() + out_raw, out_labels = transform(raw, labels) + # The raw data is unchanged and the labels gain one mask channel per label channel. + np.testing.assert_array_equal(out_raw, raw) + self.assertEqual(out_labels.shape, (4, 8, 16, 16)) + np.testing.assert_array_equal(out_labels[:2], labels) + np.testing.assert_array_equal(out_labels[2], (raw[1] != 2).astype("float32")) + np.testing.assert_array_equal(out_labels[2], out_labels[3]) + + def test_augmented_mito_state_mask_transform(self): + from synapse_net.training.transform import AugmentedMitoStateMaskTransform + + raw = np.zeros((2, 8, 16, 16), dtype="float32") + raw[1, :, :8] = 1 + raw[1, :, 8:12] = 2 + labels = np.zeros((2, 8, 16, 16), dtype="float32") + + # Use the identity as augmentation, so that the output can be compared directly. + transform = AugmentedMitoStateMaskTransform(lambda x, y: (x, y)) + out_raw, out_labels = transform(raw, labels) + self.assertEqual(tuple(out_labels.shape), (4, 8, 16, 16)) + expected = torch.as_tensor((raw[1] != 2).astype("float32")) + torch.testing.assert_close(out_labels[2], expected) + torch.testing.assert_close(out_labels[3], expected) + torch.testing.assert_close(out_raw, torch.as_tensor(raw)) + + def test_standardize_channel(self): + from synapse_net.training.transform import standardize_channel + + raw = np.stack([np.random.rand(4, 8, 8) * 100, np.full((4, 8, 8), 2.0)]).astype("float32") + out = standardize_channel(raw) + self.assertAlmostEqual(float(out[0].mean()), 0.0, places=4) + self.assertAlmostEqual(float(out[0].std()), 1.0, places=3) + # The state channel must keep its exact values, so that it can be compared against them. + np.testing.assert_array_equal(out[1], raw[1]) + with self.assertRaisesRegex(ValueError, "4 dimensions"): + standardize_channel(raw[0]) + + # The recipe. + + def test_recipe(self): + from synapse_net.training.cristae import cristae_training + from synapse_net.training.loss import MaskedDiceLossPerSample + from synapse_net.training.transform import AugmentedMitoStateMaskTransform, standardize_channel + + with mock.patch("synapse_net.training.cristae.supervised_training") as training: + cristae_training(name="test-model", train_paths=["a.h5"], val_paths=["b.h5"]) + + kwargs = training.call_args.kwargs + self.assertEqual(kwargs["raw_key"], "raw_mitos_combined") + self.assertEqual(kwargs["label_key"], "labels/cristae") + self.assertEqual(kwargs["patch_shape"], (32, 256, 256)) + self.assertEqual(kwargs["batch_size"], 24) + self.assertAlmostEqual(kwargs["lr"], 1.7e-4) + self.assertEqual(kwargs["in_channels"], 2) + self.assertEqual(kwargs["out_channels"], 2) + self.assertEqual(kwargs["initial_features"], 32) + self.assertIsNone(kwargs["norm"]) + self.assertIs(kwargs["with_channels"], True) + self.assertEqual(kwargs["early_stopping"], 25) + self.assertEqual(kwargs["log_image_interval"], 50) + self.assertIs(kwargs["mixed_precision"], True) + self.assertIs(kwargs["shuffle"], False) + self.assertEqual(kwargs["num_workers"], 8) + self.assertIs(kwargs["persistent_workers"], True) + self.assertEqual(kwargs["prefetch_factor"], 4) + self.assertIs(kwargs["raw_transform"], standardize_channel) + self.assertEqual(kwargs["sampler"].p_reject, 0.95) + self.assertIsInstance(kwargs["loss_fn"], MaskedDiceLossPerSample) + + transform = kwargs["transform"] + self.assertIsInstance(transform, AugmentedMitoStateMaskTransform) + self.assertEqual(transform.w_pos, 3.0) + self.assertEqual(transform.w_neg, 2.0) + self.assertEqual(transform.band_nm, 12.0) + self.assertEqual(transform.mito_channel, 1) + self.assertEqual(transform.exclude_state_value, 2.0) + + def test_recipe_without_augmentations(self): + from synapse_net.training.cristae import cristae_training + from synapse_net.training.transform import MitoStateMaskTransform + + with mock.patch("synapse_net.training.cristae.supervised_training") as training: + cristae_training(name="test-model", train_paths=["a.h5"], val_paths=["b.h5"], augmentations=False) + transform = training.call_args.kwargs["transform"] + self.assertIsInstance(transform, MitoStateMaskTransform) + self.assertNotIsInstance(transform, type(None)) + + @unittest.skipIf(platform.system() == "Windows", "CLI does not work on Windows") + def test_cristae_training_cli(self): + name = "test-cristae-model" + cmd = [ + "synapse_net.run_cristae_training", + "-n", name, + "-i", self.roots["ds0"], self.roots["ds1"], + "--patch_shape", "16", "32", "32", + "--batch_size", "1", + "--initial_features", "4", + "--n_samples_train", "4", + "--n_samples_val", "1", + "--n_iterations", "4", + "--num_workers", "1", + "--save_root", self.tmp_folder, + ] + run(cmd, check=True) + self.assertTrue(os.path.exists(os.path.join(self.tmp_folder, "checkpoints", name, "latest.pt"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_mitochondria_training.py b/test/test_mitochondria_training.py new file mode 100644 index 00000000..43f7232b --- /dev/null +++ b/test/test_mitochondria_training.py @@ -0,0 +1,120 @@ +import json +import os +import platform +import unittest +from shutil import rmtree +from subprocess import run +from unittest import mock + +import h5py +import torch_em +from skimage.data import binary_blobs +from skimage.measure import label + +from synapse_net.training.mitochondria import get_mitochondria_paths, mitochondria_training + + +class TestMitochondriaTraining(unittest.TestCase): + tmp_folder = "./tmp_mito_training" + data_folder = os.path.join(tmp_folder, "tomograms") + + def setUp(self): + os.makedirs(self.data_folder, exist_ok=True) + # Write the tomograms into two sub-folders, to check that they are found recursively. + for i in range(8): + data = binary_blobs(length=64, n_dim=3, volume_fraction=0.15).astype("uint8") + labels = label(data).astype("uint16") + path = os.path.join(self.data_folder, f"ds{i % 2}", f"tomo-{i}.h5") + os.makedirs(os.path.dirname(path), exist_ok=True) + with h5py.File(path, "w") as f: + f.create_dataset("raw", data=data) + f.create_dataset("labels/mitochondria", data=labels) + + def tearDown(self): + try: + rmtree(self.tmp_folder) + except OSError: + pass + + def test_get_mitochondria_paths(self): + train_paths, val_paths = get_mitochondria_paths(self.data_folder) + + self.assertEqual(len(train_paths), 7) + self.assertEqual(len(val_paths), 1) + self.assertEqual(len(set(train_paths) & set(val_paths)), 0) + for path in train_paths + val_paths: + self.assertTrue(os.path.exists(path)) + + # The split has to be deterministic for a fixed seed, and has to change with the seed. + self.assertEqual(get_mitochondria_paths(self.data_folder), (train_paths, val_paths)) + self.assertNotEqual(get_mitochondria_paths(self.data_folder, seed=1), (train_paths, val_paths)) + + def test_get_mitochondria_paths_from_split_file(self): + split = {"train": ["ds0/tomo-0.h5", "ds1/tomo-1.h5"], "val": ["ds0/tomo-2.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + train_paths, val_paths = get_mitochondria_paths(self.data_folder, split_file=split_file) + self.assertEqual(train_paths, [os.path.join(self.data_folder, name) for name in split["train"]]) + self.assertEqual(val_paths, [os.path.join(self.data_folder, name) for name in split["val"]]) + + def test_get_mitochondria_paths_from_invalid_split_file(self): + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump({"train": ["ds0/tomo-0.h5"], "val": ["does-not-exist.h5"]}, f) + + with self.assertRaisesRegex(ValueError, "do not exist"): + get_mitochondria_paths(self.data_folder, split_file=split_file) + + def test_get_mitochondria_paths_without_data(self): + with self.assertRaisesRegex(ValueError, "Did not find any files"): + get_mitochondria_paths(self.data_folder, file_pattern="*.mrc") + + def test_recipe(self): + # Check that the hyperparameters of the published 'mitochondria2' model are passed on. + with mock.patch("synapse_net.training.mitochondria.supervised_training") as training: + mitochondria_training(name="test-model", train_paths=["a.h5"], val_paths=["b.h5"]) + + kwargs = training.call_args.kwargs + self.assertEqual(kwargs["label_key"], "labels/mitochondria") + self.assertEqual(kwargs["raw_key"], "raw") + self.assertEqual(kwargs["patch_shape"], (32, 256, 256)) + self.assertEqual(kwargs["batch_size"], 8) + self.assertEqual(kwargs["lr"], 1e-4) + self.assertEqual(kwargs["n_samples_train"], 500) + self.assertEqual(kwargs["n_samples_val"], 500) + self.assertEqual(kwargs["initial_features"], 32) + self.assertEqual(kwargs["num_workers"], 8) + self.assertEqual(kwargs["early_stopping"], 20) + self.assertEqual(kwargs["log_image_interval"], 50) + self.assertIs(kwargs["mixed_precision"], False) + self.assertIs(kwargs["shuffle"], False) + self.assertIs(kwargs["raw_transform"], torch_em.transform.raw.normalize_percentile) + self.assertEqual(kwargs["sampler"].p_reject, 0.95) + + @unittest.skipIf(platform.system() == "Windows", "CLI does not work on Windows") + def test_mitochondria_training_cli(self): + name = "test-mito-model" + cmd = [ + "synapse_net.run_mitochondria_training", + "-n", name, + "-i", self.data_folder, + "--patch_shape", "32", "64", "64", + "--batch_size", "1", + "--initial_features", "4", + "--n_samples_train", "5", + "--n_samples_val", "1", + "--n_iterations", "6", + "--num_workers", "1", + "--save_root", self.tmp_folder, + ] + run(cmd, check=True) + + # Check that the checkpoint exists. + ckpt_path = os.path.join(self.tmp_folder, "checkpoints", name, "latest.pt") + self.assertTrue(os.path.exists(ckpt_path)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_training_safeguards.py b/test/test_training_safeguards.py new file mode 100644 index 00000000..bd80d808 --- /dev/null +++ b/test/test_training_safeguards.py @@ -0,0 +1,253 @@ +import importlib +import os +import unittest +from shutil import rmtree +from unittest import mock + +import numpy as np +import torch + + +class TestOverwriteGuard(unittest.TestCase): + """The guard that stops a training run from replacing the checkpoints of a previous one.""" + + tmp_folder = "./tmp_training_safeguards" + name = "a-previous-run" + + def setUp(self): + self.checkpoint_folder = os.path.join(self.tmp_folder, "checkpoints", self.name) + os.makedirs(self.checkpoint_folder, exist_ok=True) + + def tearDown(self): + try: + rmtree(self.tmp_folder) + except OSError: + pass + + def _write_checkpoint(self, name="latest", iteration=10, with_scaler=True): + save_dict = {"iteration": iteration, "epoch": 1, "best_epoch": 1, "best_metric": 0.5, + "current_metric": 0.5, "model_state": {}, "optimizer_state": {}, "train_time": 1.0} + if with_scaler: + save_dict["scaler_state"] = {} + torch.save(save_dict, os.path.join(self.checkpoint_folder, f"{name}.pt")) + + def test_refuses_to_overwrite_a_previous_run(self): + from synapse_net.training.supervised_training import _check_overwrite + self._write_checkpoint() + with self.assertRaisesRegex(ValueError, "already has checkpoints"): + _check_overwrite(self.tmp_folder, self.name, overwrite=False, resume=False) + + def test_overwrite_and_resume_bypass_the_guard(self): + from synapse_net.training.supervised_training import _check_overwrite + self._write_checkpoint() + _check_overwrite(self.tmp_folder, self.name, overwrite=True, resume=False) + _check_overwrite(self.tmp_folder, self.name, overwrite=False, resume=True) + + def test_guard_allows_a_new_run(self): + from synapse_net.training.supervised_training import _check_overwrite + # No checkpoint yet, and no save_root at all. + _check_overwrite(self.tmp_folder, self.name, overwrite=False, resume=False) + _check_overwrite(None, self.name, overwrite=False, resume=False) + + +class TestResume(unittest.TestCase): + """Resuming has to continue the previous run, not start a new one from its weights.""" + + tmp_folder = "./tmp_training_resume" + name = "a-previous-run" + + def setUp(self): + self.checkpoint_folder = os.path.join(self.tmp_folder, "checkpoints", self.name) + os.makedirs(self.checkpoint_folder, exist_ok=True) + + def tearDown(self): + try: + rmtree(self.tmp_folder) + except OSError: + pass + + def _write_checkpoint(self, name="latest", iteration=10, with_scaler=True): + save_dict = {"iteration": iteration, "epoch": 1, "best_epoch": 1, "best_metric": 0.5, + "current_metric": 0.5, "model_state": {}, "optimizer_state": {}, "train_time": 1.0} + if with_scaler: + save_dict["scaler_state"] = {} + torch.save(save_dict, os.path.join(self.checkpoint_folder, f"{name}.pt")) + + def test_resume_returns_a_stem_and_the_remaining_iterations(self): + from synapse_net.training.supervised_training import _prepare_resume + self._write_checkpoint(iteration=10) + load_from_checkpoint, remaining = _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=True) + + # torch-em appends '.pt', so the stem must NOT carry the extension, and must resolve. + self.assertFalse(load_from_checkpoint.endswith(".pt")) + self.assertTrue(os.path.exists(f"{load_from_checkpoint}.pt")) + self.assertIsInstance(load_from_checkpoint, str) + # 'n_iterations' is the total, so only the missing ones are requested. + self.assertEqual(remaining, 90) + + def test_resume_backs_up_the_previous_best(self): + from synapse_net.training.supervised_training import _prepare_resume + self._write_checkpoint() + self._write_checkpoint(name="best") + _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=True) + self.assertTrue(os.path.exists(os.path.join(self.checkpoint_folder, "best-before-resume.pt"))) + + def test_resume_without_a_previous_run(self): + from synapse_net.training.supervised_training import _prepare_resume + with self.assertRaisesRegex(ValueError, "does not exist"): + _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=True) + + def test_resume_of_a_finished_run(self): + from synapse_net.training.supervised_training import _prepare_resume + self._write_checkpoint(iteration=100) + with self.assertRaisesRegex(ValueError, "already at iteration"): + _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=True) + + def test_resume_with_mismatched_mixed_precision(self): + from synapse_net.training.supervised_training import _prepare_resume + # torch-em reads 'scaler_state' unconditionally when the new trainer has a scaler. + self._write_checkpoint(with_scaler=False) + with self.assertRaisesRegex(ValueError, "without mixed precision"): + _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=True) + # ... but continuing it without mixed precision is fine. + _, remaining = _prepare_resume(self.tmp_folder, self.name, 100, mixed_precision=False) + self.assertEqual(remaining, 90) + + +class TestSeeding(unittest.TestCase): + """'--seed' has to cover the training itself, not just the train / val split.""" + + def test_set_seed_makes_weight_init_reproducible(self): + from synapse_net.training.models import get_3d_model + from synapse_net.training.supervised_training import _set_seed + + _set_seed(42) + first = [p.detach().clone() for p in get_3d_model(out_channels=2, initial_features=4).parameters()] + _set_seed(42) + same = [p.detach().clone() for p in get_3d_model(out_channels=2, initial_features=4).parameters()] + _set_seed(7) + other = [p.detach().clone() for p in get_3d_model(out_channels=2, initial_features=4).parameters()] + + self.assertTrue(all(torch.equal(a, b) for a, b in zip(first, same))) + self.assertFalse(all(torch.equal(a, b) for a, b in zip(first, other))) + + def test_set_seed_covers_random_and_numpy(self): + import random + from synapse_net.training.supervised_training import _set_seed + + _set_seed(42) + values = (random.random(), np.random.rand()) + _set_seed(42) + self.assertEqual(values, (random.random(), np.random.rand())) + + def test_deterministic_disables_cudnn_benchmark(self): + from synapse_net.training.supervised_training import _set_seed + + benchmark = torch.backends.cudnn.benchmark + deterministic = torch.are_deterministic_algorithms_enabled() + try: + torch.backends.cudnn.benchmark = True + _set_seed(42, deterministic=False) + self.assertTrue(torch.backends.cudnn.benchmark) # untouched by default + self.assertFalse(torch.are_deterministic_algorithms_enabled()) + + _set_seed(42, deterministic=True) + self.assertFalse(torch.backends.cudnn.benchmark) + self.assertTrue(torch.are_deterministic_algorithms_enabled()) + finally: + # This is global state, so it must not leak into the other tests. + torch.backends.cudnn.benchmark = benchmark + torch.use_deterministic_algorithms(deterministic, warn_only=True) + + +class TestTrainingArgumentsArePassedOn(unittest.TestCase): + """The new arguments have to reach 'supervised_training' from every task module.""" + + def test_mitochondria_passes_the_new_arguments(self): + from synapse_net.training.mitochondria import mitochondria_training + + with mock.patch("synapse_net.training.mitochondria.supervised_training") as training: + mitochondria_training(name="m", train_paths=["a.h5"], val_paths=["b.h5"], + resume=True, overwrite=True, seed=7, deterministic=True) + kwargs = training.call_args.kwargs + self.assertIs(kwargs["resume"], True) + self.assertIs(kwargs["overwrite"], True) + self.assertEqual(kwargs["seed"], 7) + self.assertIs(kwargs["deterministic"], True) + + def test_cristae_passes_the_new_arguments(self): + from synapse_net.training.cristae import cristae_training + + with mock.patch("synapse_net.training.cristae.supervised_training") as training: + cristae_training(name="c", train_paths=["a.h5"], val_paths=["b.h5"], + resume=True, overwrite=True, seed=7, deterministic=True) + kwargs = training.call_args.kwargs + self.assertIs(kwargs["resume"], True) + self.assertIs(kwargs["overwrite"], True) + self.assertEqual(kwargs["seed"], 7) + self.assertIs(kwargs["deterministic"], True) + + def test_vol_em_mitochondria_passes_the_new_arguments(self): + from synapse_net.training.mitochondria_vol_em import vol_em_mitochondria_training + + with mock.patch("synapse_net.training.mitochondria_vol_em.supervised_training") as training: + vol_em_mitochondria_training(name="v", train_paths=["a.h5"], val_paths=["b.h5"], + resume=True, overwrite=True, seed=7, deterministic=True) + kwargs = training.call_args.kwargs + self.assertIs(kwargs["resume"], True) + self.assertIs(kwargs["overwrite"], True) + self.assertEqual(kwargs["seed"], 7) + self.assertIs(kwargs["deterministic"], True) + + def test_defaults_are_unchanged(self): + # The published recipes must not become seeded or overwriting by default. + from synapse_net.training.mitochondria import mitochondria_training + + with mock.patch("synapse_net.training.mitochondria.supervised_training") as training: + mitochondria_training(name="m", train_paths=["a.h5"], val_paths=["b.h5"]) + kwargs = training.call_args.kwargs + self.assertIs(kwargs["resume"], False) + self.assertIs(kwargs["overwrite"], False) + self.assertIsNone(kwargs["seed"]) + + +class TestScaleFactorsReachTheModel(unittest.TestCase): + """'supervised_training' used to build the model without the scale factors it was given.""" + + def _build_model(self, **kwargs): + # 'synapse_net.training' re-exports the function 'supervised_training' under the name of its + # own module, which shadows it, so the module has to be looked up explicitly. + st = importlib.import_module("synapse_net.training.supervised_training") + models = importlib.import_module("synapse_net.training.models") + + with mock.patch.object(st, "get_supervised_loader"), \ + mock.patch.object(models, "get_3d_model") as get_model, \ + mock.patch.object(st.torch_em, "default_segmentation_trainer"): + st.supervised_training( + name="m", train_paths=("a.h5",), val_paths=("b.h5",), label_key="labels", + patch_shape=(32, 512, 512), **kwargs, + ) + return get_model.call_args.kwargs + + def test_supervised_training_passes_scale_factors_to_the_model(self): + scale_factors = [[1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2]] + kwargs = self._build_model(scale_factors=scale_factors, norm=None) + self.assertEqual(kwargs["scale_factors"], scale_factors) + self.assertIsNone(kwargs["norm"]) + + def test_supervised_training_keeps_the_default_scale_factors(self): + # Not passing them must not change the model that the other recipes build. + self.assertNotIn("scale_factors", self._build_model()) + + def test_two_anisotropic_levels_build_a_usable_model(self): + from synapse_net.training.models import get_3d_model + + model = get_3d_model(out_channels=2, initial_features=4, norm=None, + scale_factors=[[1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2]]) + self.assertEqual(sum(1 for module in model.modules() if "Norm" in type(module).__name__), 0) + with torch.no_grad(): + self.assertEqual(tuple(model(torch.rand(1, 1, 8, 64, 64)).shape), (1, 2, 8, 64, 64)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vol_em_mitochondria_inference.py b/test/test_vol_em_mitochondria_inference.py new file mode 100644 index 00000000..d99f7642 --- /dev/null +++ b/test/test_vol_em_mitochondria_inference.py @@ -0,0 +1,105 @@ +"""The scripts are normally not covered by tests, but the preprocessing contract between the volume EM +mitochondria training and its inference script is the thing that is most likely to rot unnoticed: a +model applied with a different normalization than it was trained with does not predict anything useful, +and it fails silently rather than loudly. +""" + +import argparse +import importlib.util +import os +import unittest +from unittest import mock + +import torch_em + +SCRIPT_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "scripts", "volume_em", "inference", "run_mitochondria_vol_em_segmentation.py", +) + + +def _load_script(): + spec = importlib.util.spec_from_file_location("run_mitochondria_vol_em_segmentation", SCRIPT_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestVolEmMitochondriaInference(unittest.TestCase): + def setUp(self): + self.script = _load_script() + self.args = argparse.Namespace( + input_path="in.h5", output_path="out", model="model.pt", segmentation_key="seg", + data_ext=".h5", mask_path=None, force=False, scale=None, + no_white_patch_fix=False, white_patch_min_size=20, + tile_shape=None, halo=None, + min_size=self.script.MIN_SIZE, seed_distance=self.script.SEED_DISTANCE, + boundary_threshold=self.script.BOUNDARY_THRESHOLD, area_threshold=self.script.AREA_THRESHOLD, + ws_block_shape=[128, 256, 256], ws_halo=[48, 48, 48], verbose=False, + ) + + def _run(self): + with mock.patch.object(self.script, "inference_helper") as helper: + self.script.run_mitochondria_segmentation(self.args) + return helper.call_args + + def test_preprocessing_matches_the_training(self): + keywords = self._run().args[2].keywords + # The filler removal must not go into 'preprocess': the volume is standardized before that + # runs, and the filler is identified by its literal value. + self.assertIs(keywords["preprocess"], torch_em.transform.raw.normalize_percentile) + self.assertEqual(keywords["white_patch_min_size"], 20) + + def test_preprocessing_without_the_white_patch_fix(self): + self.args.no_white_patch_fix = True + keywords = self._run().args[2].keywords + self.assertIsNone(keywords["white_patch_min_size"]) + self.assertIs(keywords["preprocess"], torch_em.transform.raw.normalize_percentile) + + def test_the_filler_is_removed_before_the_segmentation(self): + # The filler has to be gone before 'segment_mitochondria' rescales the volume, which would + # interpolate it away from its literal value. + import numpy as np + + raw = np.full((4, 16, 16), 100, dtype="uint8") + raw[:, :, :4] = 255 + with mock.patch.object(self.script, "segment_mitochondria") as segment: + self.script.segment_mitochondria_vol_em(raw, white_patch_min_size=20) + passed = segment.call_args.args[0] + self.assertTrue((passed[:, :, :4] == 0).all()) + + with mock.patch.object(self.script, "segment_mitochondria") as segment: + self.script.segment_mitochondria_vol_em(raw, white_patch_min_size=None) + self.assertIs(segment.call_args.args[0], raw) + + def test_post_processing_defaults(self): + keywords = self._run().args[2].keywords + self.assertEqual(keywords["min_size"], 1000) + self.assertEqual(keywords["seed_distance"], 1) + self.assertEqual(keywords["boundary_threshold"], 0.12) + self.assertEqual(keywords["area_threshold"], 200) + self.assertEqual(keywords["model_path"], "model.pt") + + def test_the_data_is_not_rescaled_by_default(self): + # The model is applied to data at its own training resolution, and the voxel size cannot be + # derived from hdf5 files, so the scale has to stay unset unless it is passed explicitly. + kwargs = self._run().kwargs + self.assertIsNone(kwargs["scale"]) + self.assertEqual(kwargs["data_ext"], ".h5") + self.assertEqual(kwargs["output_key"], "seg") + + def test_cli_defaults_match_the_script_constants(self): + with mock.patch("sys.argv", ["script", "-i", "in.h5", "-o", "out", "-m", "model.pt"]), \ + mock.patch.object(self.script, "run_mitochondria_segmentation") as run_segmentation: + self.script.main() + args = run_segmentation.call_args.args[0] + self.assertEqual(args.min_size, self.script.MIN_SIZE) + self.assertEqual(args.seed_distance, self.script.SEED_DISTANCE) + self.assertEqual(args.boundary_threshold, self.script.BOUNDARY_THRESHOLD) + self.assertEqual(args.area_threshold, self.script.AREA_THRESHOLD) + self.assertEqual(args.segmentation_key, "seg") + self.assertIs(args.no_white_patch_fix, False) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vol_em_mitochondria_training.py b/test/test_vol_em_mitochondria_training.py new file mode 100644 index 00000000..1b18aff9 --- /dev/null +++ b/test/test_vol_em_mitochondria_training.py @@ -0,0 +1,260 @@ +import json +import os +import pickle +import platform +import unittest +from shutil import rmtree +from subprocess import run +from unittest import mock + +import h5py +import numpy as np +import torch_em +from skimage.data import binary_blobs +from skimage.measure import label + + +class TestVolEmMitochondriaTraining(unittest.TestCase): + tmp_folder = "./tmp_vol_em_mito_training" + roots = {"4007": os.path.join(tmp_folder, "4007"), "4009": os.path.join(tmp_folder, "4009")} + + def setUp(self): + for i in range(8): + # The 4007 blocks sit in sub-folders and the 4009 blocks do not, as in the real data. + if i % 2 == 0: + path = os.path.join(self.roots["4007"], f"cutout_{i}", f"block-{i}.h5") + else: + path = os.path.join(self.roots["4009"], f"block-{i}.h5") + os.makedirs(os.path.dirname(path), exist_ok=True) + + raw = (binary_blobs(length=64, n_dim=3, volume_fraction=0.15) * 200).astype("uint8") + # A white filler border, so that the raw transform has something to remove. + raw[:, :, :8] = 255 + labels = label(raw > 0).astype("uint16") + + with h5py.File(path, "w") as f: + f.create_dataset("raw", data=raw) + f.create_dataset("labels/mitochondria", data=labels) + + def tearDown(self): + try: + rmtree(self.tmp_folder) + except OSError: + pass + + # Data discovery and splitting. + + def test_get_vol_em_mitochondria_paths(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + train_paths, val_paths = get_vol_em_mitochondria_paths(self.roots) + self.assertEqual(len(train_paths), 7) + self.assertEqual(len(val_paths), 1) + self.assertEqual(len(set(train_paths) & set(val_paths)), 0) + for path in train_paths + val_paths: + self.assertTrue(os.path.exists(path)) + + # The split has to be deterministic for a fixed seed, and has to change with the seed. + self.assertEqual(get_vol_em_mitochondria_paths(self.roots), (train_paths, val_paths)) + self.assertNotEqual(get_vol_em_mitochondria_paths(self.roots, seed=1), (train_paths, val_paths)) + + def test_get_vol_em_mitochondria_paths_with_exclude(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + train_paths, val_paths = get_vol_em_mitochondria_paths(self.roots, exclude=["block-0.", "block-1."]) + self.assertEqual(len(train_paths) + len(val_paths), 6) + self.assertFalse(any("block-0." in path or "block-1." in path for path in train_paths + val_paths)) + + def test_get_vol_em_mitochondria_paths_from_split_file(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + split = {"train": ["4007/cutout_0/block-0.h5", "4009/block-1.h5"], "val": ["4007/cutout_2/block-2.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + train_paths, val_paths = get_vol_em_mitochondria_paths(self.roots, split_file=split_file) + self.assertEqual(train_paths, [os.path.join(self.tmp_folder, *name.split("/")) for name in split["train"]]) + self.assertEqual(val_paths, [os.path.join(self.tmp_folder, *name.split("/")) for name in split["val"]]) + + def test_get_vol_em_mitochondria_paths_from_split_file_with_roots(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + # A split file that carries its own roots can be used without passing them. + split = {"roots": self.roots, "train": ["4007/cutout_0/block-0.h5"], "val": ["4009/block-1.h5"]} + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump(split, f) + + train_paths, val_paths = get_vol_em_mitochondria_paths(split_file=split_file) + self.assertEqual(train_paths, [os.path.join(self.roots["4007"], "cutout_0", "block-0.h5")]) + self.assertEqual(val_paths, [os.path.join(self.roots["4009"], "block-1.h5")]) + + # Roots passed explicitly take precedence, so that the data can be moved. + moved = os.path.join(self.tmp_folder, "moved") + os.makedirs(os.path.join(moved, "cutout_0"), exist_ok=True) + open(os.path.join(moved, "cutout_0", "block-0.h5"), "w").close() + train_paths, _ = get_vol_em_mitochondria_paths({"4007": moved}, split_file=split_file) + self.assertEqual(train_paths, [os.path.join(moved, "cutout_0", "block-0.h5")]) + + def test_get_vol_em_mitochondria_paths_without_roots(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + with self.assertRaisesRegex(ValueError, "No data roots"): + get_vol_em_mitochondria_paths() + + def test_get_vol_em_mitochondria_paths_from_invalid_split_file(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + split_file = os.path.join(self.tmp_folder, "split.json") + with open(split_file, "w") as f: + json.dump({"train": ["4007/cutout_0/block-0.h5"], "val": ["4007/does-not-exist.h5"]}, f) + with self.assertRaisesRegex(ValueError, "do not exist"): + get_vol_em_mitochondria_paths(self.roots, split_file=split_file) + + with open(split_file, "w") as f: + json.dump({"train": ["nope/block-0.h5"], "val": []}, f) + with self.assertRaisesRegex(ValueError, "unknown data root"): + get_vol_em_mitochondria_paths(self.roots, split_file=split_file) + + def test_get_vol_em_mitochondria_paths_without_data(self): + from synapse_net.training.mitochondria_vol_em import get_vol_em_mitochondria_paths + + with self.assertRaisesRegex(ValueError, "Did not find any files"): + get_vol_em_mitochondria_paths(self.roots, file_pattern="*.mrc") + + # The white patch removal. + + def test_remove_white_patches(self): + from synapse_net.training.transform import remove_white_patches + + raw = np.full((8, 16, 16), 100, dtype="uint8") + raw[:, :, :4] = 255 # The filler border, well above the minimal size. + raw[0, 0, 10] = 255 # A single saturated voxel, which is sample structure. + + out = remove_white_patches(raw) + self.assertTrue((out[:, :, :4] == 0).all()) + self.assertEqual(out[0, 0, 10], 255) + self.assertTrue((out[:, :, 4:10] == 100).all()) + self.assertEqual(out.dtype, raw.dtype) + # The input must not be modified in place. + self.assertTrue((raw[:, :, :4] == 255).all()) + + def test_remove_white_patches_size_threshold(self): + from synapse_net.training.transform import remove_white_patches + + # The threshold is inclusive: a component of exactly 'min_size' voxels is removed. + raw = np.zeros((4, 8, 8), dtype="uint8") + raw[0, 0, :4] = 255 + self.assertTrue((remove_white_patches(raw, min_size=4) == 0).all()) + self.assertEqual((remove_white_patches(raw, min_size=5) == 255).sum(), 4) + + def test_remove_white_patches_connectivity(self): + from synapse_net.training.transform import remove_white_patches + + # Two blocks that only touch diagonally form one component under full connectivity, so both + # are removed. With the 6-connectivity of 'scipy.ndimage.label' neither would be. + raw = np.zeros((4, 8, 8), dtype="uint8") + raw[0, 0:2, 0:2] = 255 + raw[0, 2:4, 2:4] = 255 + self.assertEqual((remove_white_patches(raw, min_size=8) == 255).sum(), 0) + self.assertEqual((remove_white_patches(raw, min_size=9) == 255).sum(), 8) + + def test_remove_white_patches_without_filler(self): + from synapse_net.training.transform import remove_white_patches + + raw = np.full((4, 8, 8), 100, dtype="uint8") + self.assertIs(remove_white_patches(raw), raw) + + def test_remove_white_patches_rejects_normalized_data(self): + from synapse_net.training.transform import remove_white_patches + + raw = np.random.rand(4, 8, 8).astype("float32") + with self.assertRaisesRegex(ValueError, "integer dtype"): + remove_white_patches(raw) + + with self.assertRaisesRegex(ValueError, "2d or 3d"): + remove_white_patches(np.zeros((2, 4, 8, 8), dtype="uint8")) + + def test_remove_white_patches_and_normalize(self): + from synapse_net.training.transform import RemoveWhitePatchesAndNormalize, remove_white_patches + + raw = (binary_blobs(length=32, n_dim=3, volume_fraction=0.2) * 200).astype("uint8") + raw[:, :, :4] = 255 + + transform = RemoveWhitePatchesAndNormalize() + expected = torch_em.transform.raw.normalize_percentile(remove_white_patches(raw)) + self.assertTrue(np.allclose(transform(raw), expected)) + + def test_remove_white_patches_and_normalize_is_picklable(self): + from synapse_net.training.transform import RemoveWhitePatchesAndNormalize + + # The dataloader workers and the training checkpoint both pickle the raw transform. + transform = pickle.loads(pickle.dumps(RemoveWhitePatchesAndNormalize(min_size=7))) + self.assertEqual(transform.min_size, 7) + + # The recipe. + + def test_recipe(self): + # Check that the hyperparameters of the published volume EM mitochondria model are passed on. + from synapse_net.training.mitochondria_vol_em import vol_em_mitochondria_training + from synapse_net.training.transform import RemoveWhitePatchesAndNormalize + + with mock.patch("synapse_net.training.mitochondria_vol_em.supervised_training") as training: + vol_em_mitochondria_training(name="test-model", train_paths=["a.h5"], val_paths=["b.h5"]) + + kwargs = training.call_args.kwargs + self.assertEqual(kwargs["label_key"], "labels/mitochondria") + self.assertEqual(kwargs["raw_key"], "raw") + self.assertEqual(kwargs["patch_shape"], (32, 512, 512)) + self.assertEqual(kwargs["batch_size"], 4) + self.assertEqual(kwargs["lr"], 1e-4) + self.assertEqual(kwargs["n_iterations"], 50000) + self.assertEqual(kwargs["n_samples_train"], 500) + self.assertEqual(kwargs["n_samples_val"], 500) + self.assertEqual(kwargs["initial_features"], 32) + self.assertEqual(kwargs["scale_factors"], [[1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2]]) + self.assertIsNone(kwargs["norm"]) + self.assertEqual(kwargs["num_workers"], 8) + self.assertEqual(kwargs["early_stopping"], 20) + self.assertEqual(kwargs["log_image_interval"], 50) + self.assertIs(kwargs["mixed_precision"], False) + self.assertIs(kwargs["shuffle"], False) + self.assertIsInstance(kwargs["raw_transform"], RemoveWhitePatchesAndNormalize) + self.assertEqual(kwargs["raw_transform"].min_size, 20) + self.assertEqual(kwargs["sampler"].p_reject, 0.95) + + def test_recipe_without_white_patch_fix(self): + from synapse_net.training.mitochondria_vol_em import vol_em_mitochondria_training + + with mock.patch("synapse_net.training.mitochondria_vol_em.supervised_training") as training: + vol_em_mitochondria_training(name="test-model", train_paths=["a.h5"], val_paths=["b.h5"], + fix_white_patches=False) + raw_transform = training.call_args.kwargs["raw_transform"] + self.assertIs(raw_transform, torch_em.transform.raw.normalize_percentile) + + @unittest.skipIf(platform.system() == "Windows", "CLI does not work on Windows") + def test_vol_em_mitochondria_training_cli(self): + name = "test-vol-em-mito-model" + cmd = [ + "synapse_net.run_vol_em_mitochondria_training", + "-n", name, + "-i", self.roots["4007"], self.roots["4009"], + "--patch_shape", "8", "64", "64", + "--batch_size", "1", + "--initial_features", "4", + "--n_samples_train", "5", + "--n_samples_val", "1", + "--n_iterations", "6", + "--num_workers", "1", + "--save_root", self.tmp_folder, + ] + run(cmd, check=True) + + # Check that the checkpoint exists. + ckpt_path = os.path.join(self.tmp_folder, "checkpoints", name, "latest.pt") + self.assertTrue(os.path.exists(ckpt_path)) + + +if __name__ == "__main__": + unittest.main()