Merge branch 'master' of github.com:PytorchConnectomics/pytorch_connectomics

Author: Tian-hao Zhang

mito_lucchi++/20260210_132007/results/logs/events.out.tfevents.1770841281.gdw001.1316099.0

@@ -39,9 +39,9 @@ def main():
         )
         print("")
         print("Examples:")
-        print('  python scripts/images_to_h5.py "datasets/images/*.tiff" output.h5')
+        print('  python scripts/images_to_h5.py "datasets/images/*.tiff" output.h5 seg')
         print('  python scripts/images_to_h5.py "/path/to/images/*.png" output.h5')
-        print('  python scripts/images_to_h5.py "/path/to/labels/*.png" labels.h5 labels')
+        print('  python scripts/images_to_h5.py "/path/to/labels/*.png" labels.h5 seg')
         print("")
         print("Note: Use quotes around the input pattern to prevent shell expansion.")
         sys.exit(1)

scripts/images_to_h5.py

@@ -9,7 +9,10 @@ python scripts/main.py --config tutorials/<config>.yaml
 ## Active top-level configs
 
 - `tutorials/mito_lucchi++.yaml`: Lucchi++ mitochondria segmentation (MONAI UNet).
-- `tutorials/mito_mitoEM.yaml`: MitoEM mitochondria instance segmentation (MedNeXt, SDT).
+- `tutorials/mito_mitoEM.yaml`: Backward-compatible alias to `mito_mitoEM_30h.yaml`.
+- `tutorials/mito_mitoEM_30h.yaml`: MitoEM-Human (EM30-H) instance segmentation (MedNeXt, SDT).
+- `tutorials/mito_mitoEM_30r.yaml`: MitoEM-Rat (EM30-R) instance segmentation (MedNeXt, SDT).
+- `tutorials/mito_mitoEM_30hr.yaml`: Joint EM30-H + EM30-R training (MedNeXt, SDT).
 - `tutorials/mito_mitolab.yaml`: CEM-MitoLab 2D mitochondria segmentation (MedNeXt).
 - `tutorials/mito_betaseg.yaml`: BetaSeg mitochondria instance segmentation (MedNeXt, affinity+SDT).
 - `tutorials/neuron_snemi.yaml`: SNEMI3D neuron segmentation (RSUNet, affinities).

tutorials/README.md

@@ -45,106 +45,40 @@ model:
     - sdt
     - - 1
 data:
-  train_path: /projects/weilab/qiongwang/datasets/betaseg/tif
+  train_path: /projects/weilab/dataset/mito/mitoEM/
   train_image:
-  - high_c3_im.tiff
-  - low_c1_im.tiff
-  - low_c2_im.tiff
+  - EM30-H/im_train_val.h5
   train_label:
-  - high_c3_mito.tiff
-  - low_c1_mito.tiff
-  - low_c2_mito.tiff
+  - EM30-H/mito_train_val.h5
   train_resolution:
-  - 16
-  - 16
-  - 16
-  val_path: /projects/weilab/qiongwang/datasets/betaseg/tif
-  val_image:
-  - high_c1_im.tiff
-  val_label:
-  - high_c1_mito.tiff
-  use_preloaded_cache: false
-  use_cache: true
+  - 30
+  - 8
+  - 8
+  use_preloaded_cache: true
   cache_rate: 1.0
   persistent_workers: true
   patch_size:
-  - 128
-  - 128
-  - 128
-  pad_size:
   - 16
+  - 256
+  - 256
+  pad_size:
+  - 4
   - 16
   - 16
   pad_mode: reflect
-  iter_num_per_epoch: 790
+  iter_num_per_epoch: 2000
   image_transform:
     clip_percentile_low: 0.005
     clip_percentile_high: 0.995
   label_transform:
     targets:
-    - name: affinity
-      kwargs:
-        offsets:
-        - 0-0-1
-        - 0-1-0
-        - 1-0-0
-        - 0-0-10
-        - 0-10-0
-        - 10-0-0
-    - name: skeleton_aware_edt
+    - name: instance_edt
       kwargs:
-        resolution:
-        - 16
-        - 16
-        - 16
-        alpha: 0.8
-        bg_value: -1.0
-        relabel: true
+        mode: 3d
+        quantize: false
+        normalize: true
   augmentation:
-    preset: some
-    affine:
-      enabled: true
-      prob: 0.5
-      rotate_range:
-      - 0.2
-      - 0.2
-      - 0.2
-      scale_range:
-      - 0.2
-      - 0.2
-      - 0.2
-      shear_range:
-      - 0.5
-      - 0.5
-      - 0.5
-    intensity:
-      enabled: true
-      gaussian_noise_prob: 0.3
-      gaussian_noise_std: 0.5
-      shift_intensity_prob: 0.3
-      shift_intensity_offset: 0.1
-      contrast_prob: 0.3
-      contrast_range:
-      - 0.7
-      - 1.4
-    missing_section:
-      enabled: true
-      prob: 0.05
-      num_sections: 2
-    misalignment:
-      enabled: true
-      prob: 0.05
-      displacement: 10
-      rotate_ratio: 0.0
-    flip:
-      enabled: true
-      prob: 0.5
-    rotate:
-      enabled: true
-      prob: 0.5
-    elastic:
-      enabled: true
-      prob: 0.3
+    preset: all
 optimization:
   max_steps: 1000000
   accumulate_grad_batches: 1
@@ -186,23 +120,13 @@ monitor:
     divergence_threshold: 100.0
 test:
   data:
-    test_image:
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/high_c2_im.tiff
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/high_c4_im.tiff
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/low_c3_im.tiff
-    test_label:
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/high_c2_mito.tiff
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/high_c4_mito.tiff
-    - /projects/weilab/qiongwang/datasets/betaseg/tif/low_c3_mito.tiff
+    test_image: /projects/weilab/dataset/mito/mitoEM/EM30-H/im_test.h5
+    test_label: /projects/weilab/dataset/mito/mitoEM/EM30-H/mito_test.h5
     test_resolution:
-    - 16
-    - 16
-    - 16
-    output_path: outputs/betaseg_mednext_affinity_sdt/results/
-    image_transform:
-      normalize: 0-1
-      clip_percentile_low: 0.005
-      clip_percentile_high: 0.995
+    - 30
+    - 8
+    - 8
+    output_path: outputs/mitoem_mednext_sdt/results/
   decoding:
   - name: decode_distance_watershed
     kwargs:

tutorials/mito_betaseg.yaml

@@ -0,0 +1,20 @@
+_base_: mito_mitoEM_common.yaml
+
+experiment_name: mitoem30h_mednext_sdt
+data:
+  train_image:
+  - EM30-H/im_train.h5
+  train_label:
+  - EM30-H/mito_train.h5
+  val_image:
+  - EM30-H/im_val.h5
+  val_label:
+  - EM30-H/mito_val.h5
+monitor:
+  checkpoint:
+    dirpath: outputs/mitoem30h_mednext_sdt/checkpoints/
+test:
+  data:
+    test_image: /projects/weilab/dataset/mito/mitoEM/EM30-H/im_test.h5
+    test_label: /projects/weilab/dataset/mito/mitoEM/EM30-H/mito_test.h5
+    output_path: outputs/mitoem30h_mednext_sdt/results/

tutorials/mito_mitoEM_30h.yaml

@@ -0,0 +1,28 @@
+_base_: mito_mitoEM_common.yaml
+
+experiment_name: mitoem30hr_mednext_sdt
+data:
+  train_image:
+  - EM30-H/im_train.h5
+  - EM30-R/im_train.h5
+  train_label:
+  - EM30-H/mito_train.h5
+  - EM30-R/mito_train.h5
+  val_image:
+  - EM30-H/im_val.h5
+  - EM30-R/im_val.h5
+  val_label:
+  - EM30-H/mito_val.h5
+  - EM30-R/mito_val.h5
+monitor:
+  checkpoint:
+    dirpath: outputs/mitoem30hr_mednext_sdt/checkpoints/
+test:
+  data:
+    test_image:
+    - /projects/weilab/dataset/mito/mitoEM/EM30-H/im_test.h5
+    - /projects/weilab/dataset/mito/mitoEM/EM30-R/im_test.h5
+    test_label:
+    - /projects/weilab/dataset/mito/mitoEM/EM30-H/mito_test.h5
+    - /projects/weilab/dataset/mito/mitoEM/EM30-R/mito_test.h5
+    output_path: outputs/mitoem30hr_mednext_sdt/results/

tutorials/mito_mitoEM_30hr.yaml

@@ -0,0 +1,20 @@
+_base_: mito_mitoEM_common.yaml
+
+experiment_name: mitoem30r_mednext_sdt
+data:
+  train_image:
+  - EM30-R/im_train.h5
+  train_label:
+  - EM30-R/mito_train.h5
+  val_image:
+  - EM30-R/im_val.h5
+  val_label:
+  - EM30-R/mito_val.h5
+monitor:
+  checkpoint:
+    dirpath: outputs/mitoem30r_mednext_sdt/checkpoints/
+test:
+  data:
+    test_image: /projects/weilab/dataset/mito/mitoEM/EM30-R/im_test.h5
+    test_label: /projects/weilab/dataset/mito/mitoEM/EM30-R/mito_test.h5
+    output_path: outputs/mitoem30r_mednext_sdt/results/

tutorials/mito_mitoEM_30r.yaml

@@ -1,21 +1,23 @@
+_base_: bases/mednext.yaml
+
 description: MitoEM 3D mitochondria instance segmentation with MedNeXt using SDT only
 system:
   training:
     num_cpus: 8
     num_workers: 8
-    batch_size: 16
+    batch_size: 4
   inference:
     num_cpus: 1
     num_workers: 1
     batch_size: 1
   seed: 42
 model:
   input_size:
-  - 16
+  - 32
   - 256
   - 256
   output_size:
-  - 16
+  - 32
   - 256
   - 256
   out_channels: 1
@@ -31,10 +33,6 @@ model:
   - tanh: true
 data:
   train_path: /projects/weilab/dataset/mito/mitoEM/
-  train_image:
-  - EM30-H/im_train_val.h5
-  train_label:
-  - EM30-H/mito_train_val.h5
   train_resolution:
   - 30
   - 8
@@ -43,7 +41,7 @@ data:
   cache_rate: 1.0
   persistent_workers: true
   patch_size:
-  - 16
+  - 32
   - 256
   - 256
   pad_size:
@@ -66,48 +64,55 @@ data:
     preset: all
 optimization:
   max_epochs: 1000
-  accumulate_grad_batches: 4
-  precision: 16-mixed
+  accumulate_grad_batches: 2
+  precision: bf16-mixed
+  deterministic: false
+  benchmark: true
   optimizer:
-    lr: 0.001
+    lr: 0.0003
     weight_decay: 1e-4
     eps: 1.0e-08
   scheduler:
-    name: constant
+    name: WarmupCosineLR
+    warmup_epochs: 20
+    warmup_start_lr: 3.0e-05
+    min_lr: 1.0e-06
+    interval: epoch
+    frequency: 1
+  ema:
+    enabled: true
+    decay: 0.999
+    warmup_steps: 500
+    validate_with_ema: true
 monitor:
   logging:
     scalar:
       loss:
       - train_loss_total_epoch
-      loss_every_n_steps: 50
+      - val_loss_total
+      - val_loss_total_epoch
+      loss_every_n_steps: 100
       val_check_interval: 1.0
     images:
       max_images: 4
       num_slices: 8
       log_every_n_epochs: 5
   checkpoint:
+    monitor: val_loss_total_epoch
     save_top_k: 5
-    save_every_n_epochs: 25
-    dirpath: outputs/mitoem_mednext_sdt/checkpoints/
+    save_every_n_epochs: 10
   early_stopping:
-    monitor: train_loss_total_epoch
-    patience: 150
-    min_delta: 1e-6
-    threshold: 0.01
-    divergence_threshold: 100.0
+    enabled: false
 test:
   data:
-    test_image: /projects/weilab/dataset/mito/mitoEM/EM30-H/im_test.h5
-    test_label: /projects/weilab/dataset/mito/mitoEM/EM30-H/mito_test.h5
     test_resolution:
     - 30
     - 8
     - 8
-    output_path: outputs/mitoem_mednext_sdt/results/
 inference:
   sliding_window:
     window_size:
-    - 16
+    - 32
     - 256
     - 256
     sw_batch_size: 4