Add vol em mitochondria training

.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
+scripts/cooper/revision/count_AZ.py
+tmp_*/

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.

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.