add LV segmentation

Author: linok-bc

ves_seg/README.md

@@ -0,0 +1,41 @@
+# Large vesicle segmentation
+
+This module contains code for training and infering on a new dataset, as well as pretrained model checkpoints and a sample for validation. For setup, install [Pytorch Connectomics](https://connectomics.readthedocs.io/en/latest/notes/installation.html), [em_util](https://github.com/PytorchConnectomics/em_util), [neuroglancer](https://pypi.org/project/neuroglancer), and [cloud-volume](https://pypi.org/project/cloud-volume).
+
+## Training
+Download your data. This model was built using the Pytorch Connectomics package, and thus all configuration for training is done within the yaml files within `configs`. Refer to the file `configs/00_base.yaml` as a template. Most fields can be left untouched, but to adapt to your dataset, change the following parameters at the very minimum:
+
+* `DATASET.INPUT_PATH`: points towards directory containing training data. PyTC's dataloader will automatically seperate this into training/validation partitions.
+* `DATASET.IMAGE_NAME`: points to a single image or a list of images. Lists can be represented using standard YAML notation or by concatenating lists of images with `@` as a delimiter
+* `DATASET.OUTPUT_PATH`: points to an output directory that will contain model checkpoints
+
+After editing the config file, you can execute training with:
+
+`python scripts/main.py --config-base configs/<base.yaml> --config-file configs/bcd_config.yaml`
+
+To continue retraining from a checkpoint, use:
+
+`python scripts/main.py --config-base configs/<base.yaml> --config-file configs/bcd_config.yaml --checkpoint <path-to-checkpoint> SOLVER.ITERATION_RESTART True`
+
+## Inference
+
+Inference tools are available in `tools/process.py`. To do inference, in a Python shell or file, run:
+
+`> do_inference(<im_path>, <pred_path>, [<base_config.yaml>, <bcd_config.yaml>], <checkpoint_path>)`
+
+To also get adapted rand, precison, and recall metrics, run:
+
+`> infer_and_stats([<im_path1>, ...], [<pred_path1>, ...], [<mask_path1, ...>], [<base_config.yaml>, <bcd_config.yaml>], <checkpoint_path>), 
+`
+
+Our final model checkpoint is provided at `outputs/checkpoint_1000000.pth.tar`.
+
+## Visualiztion
+
+Visualing predictions is done with neuroglancer. To load a volume and its segmentation, modify the file `scripts/ng.py`, then run `python -i scripts/ng.py`. A link to view neuroglancer will open. If running on a remote cluster, please note that port forwarding will be necessary to view on your machine. The `screnshot()` function can be used to take high-resolution screenshots of the neuroglancer representation.
+
+Further documentation for both functions is contained in `tools/process.py`
+
+## Sample
+
+A small sample is prepared in the `sample` directory. `11-5` shows a small inter-neuron region with an assortment of large vesicles. Neuroglancer scripts that display the CLAHE-enhanced images, neuron mask, ground truth, and predictions with the final model checkpoints are prepared in `sample/11-5_ng.py`.

ves_seg/configs/00_base.yaml

@@ -0,0 +1,45 @@
+SYSTEM:
+  NUM_GPUS: 1
+  NUM_CPUS: 16
+MODEL:
+  ARCHITECTURE: unet_3d
+  BLOCK_TYPE: residual_se
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  NORM_MODE: gn
+  IN_PLANES: 1
+  MIXED_PRECESION: False
+  FILTERS: [32, 64, 96, 128, 160]
+  ISOTROPY: [False, False, False, True, True]
+DATASET:
+  INPUT_PATH: input_dir
+  IMAGE_NAME: train_img1.h5@train_img2.h5@train_img3.h5
+  LABEL_NAME: train_gt1.h5@train_gt2.h5@train_gt3.h5
+  VALID_MASK_NAME: train_mask1.h5@train_mask2.h5@train_mask3.h5
+  REDUCE_LABEL: False
+  OUTPUT_PATH: output_dir
+  PAD_SIZE: [0, 0, 0]
+  DATA_SCALE: [1.0, 1.0, 1.0]
+  VALID_RATIO: 0.1
+  DISTRIBUTED: False
+SOLVER:
+  LR_SCHEDULER_NAME: WarmupCosineLR
+  BASE_LR: 0.02
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 10000
+  ITERATION_TOTAL: 100000
+  SAMPLES_PER_BATCH: 64
+MONITOR:
+  ITERATION_NUM: [40, 400]
+INFERENCE:
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  INPUT_PATH: input_dir
+  OUTPUT_PATH: output_dir
+  IMAGE_NAME: infer_im.h5
+  OUTPUT_NAME: infer_pred.h5
+  PAD_SIZE: [0, 0, 0]
+  AUG_MODE: "mean"
+  AUG_NUM: 4
+  STRIDE: [4, 32, 32]
+  SAMPLES_PER_BATCH: 24

ves_seg/configs/01_base.yaml

@@ -0,0 +1,45 @@
+SYSTEM:
+  NUM_GPUS: 1
+  NUM_CPUS: 16
+MODEL:
+  ARCHITECTURE: unet_3d
+  BLOCK_TYPE: residual_se
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  NORM_MODE: gn
+  IN_PLANES: 1
+  MIXED_PRECESION: False
+  FILTERS: [32, 64, 96, 128, 160]
+  ISOTROPY: [False, False, False, True, True]
+DATASET:
+  INPUT_PATH: /mmfs1/data/linok/hydra/datasets/large_vesicle
+  IMAGE_NAME: 01-initial/vol0_clahe.h5@01-initial/vol1_clahe.h5@01-initial/vol2_clahe.h5
+  LABEL_NAME: 01-initial/vol0_ves.h5@01-initial/vol1_ves.h5@01-initial/vol2_ves.h5
+  VALID_MASK_NAME: 01-initial/vol0_mask.h5@01-initial/vol1_mask.h5@01-initial/vol2_mask.h5
+  REDUCE_LABEL: False
+  OUTPUT_PATH: /mmfs1/data/linok/hydra/outputs/large_vesicle/01-initial
+  PAD_SIZE: [0, 0, 0]
+  DATA_SCALE: [1.0, 1.0, 1.0]
+  VALID_RATIO: 0.1
+  DISTRIBUTED: False
+SOLVER:
+  LR_SCHEDULER_NAME: WarmupCosineLR
+  BASE_LR: 0.02
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 10000
+  ITERATION_TOTAL: 100000
+  SAMPLES_PER_BATCH: 64
+MONITOR:
+  ITERATION_NUM: [40, 400]
+INFERENCE:
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  INPUT_PATH: /data/projects/weilab/dataset/hydra/big_vesicle/v0-intial_dataset
+  OUTPUT_PATH: outputs/large_vesicle/initial_volumes/results
+  IMAGE_NAME: vol0_im.h5
+  OUTPUT_NAME: vol0_inf.h5
+  PAD_SIZE: [0, 0, 0]
+  AUG_MODE: "mean"
+  AUG_NUM: 4
+  STRIDE: [4, 32, 32]
+  SAMPLES_PER_BATCH: 24

ves_seg/configs/02_base.yaml

@@ -0,0 +1,45 @@
+SYSTEM:
+  NUM_GPUS: 1
+  NUM_CPUS: 16
+MODEL:
+  ARCHITECTURE: unet_3d
+  BLOCK_TYPE: residual_se
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  NORM_MODE: gn
+  IN_PLANES: 1
+  MIXED_PRECESION: False
+  FILTERS: [32, 64, 96, 128, 160]
+  ISOTROPY: [False, False, False, True, True]
+DATASET:
+  INPUT_PATH: /mmfs1/data/linok/hydra/datasets/large_vesicle
+  IMAGE_NAME: 01-initial/vol0_clahe.h5@01-initial/vol1_clahe.h5@01-initial/vol2_clahe.h5@03-finetune/vol_10_im0_clahe.h5@03-finetune/vol_10_im1_clahe.h5@03-finetune/vol_11_im0_clahe.h5@03-finetune/vol_11_im1_clahe.h5@03-finetune/vol_11_im2_clahe.h5@03-finetune/vol_12_im0_clahe.h5@03-finetune/vol_12_im1_clahe.h5@03-finetune/vol_12_im2_clahe.h5@03-finetune/vol_12_im3_clahe.h5@03-finetune/vol_12_im4_clahe.h5@03-finetune/vol_12_im5_clahe.h5@03-finetune/vol_13_im0_clahe.h5@03-finetune/vol_13_im1_clahe.h5@03-finetune/vol_13_im2_clahe.h5@03-finetune/vol_14_im0_clahe.h5@03-finetune/vol_14_im1_clahe.h5@03-finetune/vol_14_im2_clahe.h5@03-finetune/vol_14_im3_clahe.h5@03-finetune/vol_14_im4_clahe.h5@03-finetune/vol_15_im_clahe.h5
+  LABEL_NAME: 01-initial/vol0_ves.h5@01-initial/vol1_ves.h5@01-initial/vol2_ves.h5@03-finetune/vol_10_ves0.h5@03-finetune/vol_10_ves1.h5@03-finetune/vol_11_ves0.h5@03-finetune/vol_11_ves1.h5@03-finetune/vol_11_ves2.h5@03-finetune/vol_12_ves0.h5@03-finetune/vol_12_ves1.h5@03-finetune/vol_12_ves2.h5@03-finetune/vol_12_ves3.h5@03-finetune/vol_12_ves4.h5@03-finetune/vol_12_ves5.h5@03-finetune/vol_13_ves0.h5@03-finetune/vol_13_ves1.h5@03-finetune/vol_13_ves2.h5@03-finetune/vol_14_ves0.h5@03-finetune/vol_14_ves1.h5@03-finetune/vol_14_ves2.h5@03-finetune/vol_14_ves3.h5@03-finetune/vol_14_ves4.h5@03-finetune/vol_15_ves.h5
+  VALID_MASK_NAME: 01-initial/vol0_mask.h5@01-initial/vol1_mask.h5@01-initial/vol2_mask.h5@03-finetune/vol_10_mask0.h5@03-finetune/vol_10_mask1.h5@03-finetune/vol_11_mask0.h5@03-finetune/vol_11_mask1.h5@03-finetune/vol_11_mask2.h5@03-finetune/vol_12_mask0.h5@03-finetune/vol_12_mask1.h5@03-finetune/vol_12_mask2.h5@03-finetune/vol_12_mask3.h5@03-finetune/vol_12_mask4.h5@03-finetune/vol_12_mask5.h5@03-finetune/vol_13_mask0.h5@03-finetune/vol_13_mask1.h5@03-finetune/vol_13_mask2.h5@03-finetune/vol_14_mask0.h5@03-finetune/vol_14_mask1.h5@03-finetune/vol_14_mask2.h5@03-finetune/vol_14_mask3.h5@03-finetune/vol_14_mask4.h5@03-finetune/vol_15_mask.h5
+  REDUCE_LABEL: False
+  OUTPUT_PATH: /mmfs1/data/linok/hydra/outputs/large_vesicle/03-finetune
+  PAD_SIZE: [0, 0, 0]
+  DATA_SCALE: [1.0, 1.0, 1.0]
+  VALID_RATIO: 0.1
+  DISTRIBUTED: False
+SOLVER:
+  LR_SCHEDULER_NAME: WarmupCosineLR
+  BASE_LR: 0.02
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 10000
+  ITERATION_TOTAL: 300000
+  SAMPLES_PER_BATCH: 64
+MONITOR:
+  ITERATION_NUM: [40, 400]
+INFERENCE:
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  INPUT_PATH: /data/projects/weilab/dataset/hydra/big_vesicle/v0-intial_dataset
+  OUTPUT_PATH: outputs/large_vesicle/initial_volumes/results
+  IMAGE_NAME: vol0_im.h5
+  OUTPUT_NAME: vol0_inf.h5
+  PAD_SIZE: [0, 0, 0]
+  AUG_MODE: "mean"
+  AUG_NUM: 4
+  STRIDE: [4, 32, 32]
+  SAMPLES_PER_BATCH: 24

ves_seg/configs/03_base.yaml

@@ -0,0 +1,49 @@
+SYSTEM:
+  NUM_GPUS: 1
+  NUM_CPUS: 64
+  DISTRIBUTED: True
+  PARALLEL: 'DDP'
+MODEL:
+  ARCHITECTURE: unet_3d
+  BLOCK_TYPE: residual_se
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  NORM_MODE: gn
+  IN_PLANES: 1
+  MIXED_PRECESION: False
+  FILTERS: [32, 64, 96, 128]
+  ISOTROPY: [False, False, False, True]
+DATASET:
+  INPUT_PATH: /mmfs1/data/linok/hydra/datasets/large_vesicle
+  IMAGE_NAME: 01-initial/vol0_clahe.h5@01-initial/vol1_clahe.h5@01-initial/vol2_clahe.h5@03-finetune/vol_10_im0_clahe.h5@03-finetune/vol_10_im1_clahe.h5@03-finetune/vol_11_im0_clahe.h5@03-finetune/vol_11_im1_clahe.h5@03-finetune/vol_11_im2_clahe.h5@03-finetune/vol_12_im0_clahe.h5@03-finetune/vol_12_im1_clahe.h5@03-finetune/vol_12_im2_clahe.h5@03-finetune/vol_12_im3_clahe.h5@03-finetune/vol_12_im4_clahe.h5@03-finetune/vol_12_im5_clahe.h5@03-finetune/vol_13_im0_clahe.h5@03-finetune/vol_13_im1_clahe.h5@03-finetune/vol_13_im2_clahe.h5@03-finetune/vol_14_im0_clahe.h5@03-finetune/vol_14_im1_clahe.h5@03-finetune/vol_14_im2_clahe.h5@03-finetune/vol_14_im3_clahe.h5@03-finetune/vol_14_im4_clahe.h5@03-finetune/vol_15_im_clahe.h5@05-tile_0-188/6-12_clahe1.h5@05-tile_0-188/6-12_clahe2.h5@05-tile_0-188/6-12_clahe4.h5@05-tile_0-188/6-13_clahe1.h5@05-tile_0-188/6-13_clahe2.h5@05-tile_0-188/6-13_clahe3.h5@05-tile_0-188/6-13_clahe4.h5@05-tile_0-188/6-13_clahe5.h5@05-tile_0-188/6-13_clahe6.h5@05-tile_0-188/6-13_clahe7.h5
+  LABEL_NAME: 01-initial/vol0_ves.h5@01-initial/vol1_ves.h5@01-initial/vol2_ves.h5@03-finetune/vol_10_ves0.h5@03-finetune/vol_10_ves1.h5@03-finetune/vol_11_ves0.h5@03-finetune/vol_11_ves1.h5@03-finetune/vol_11_ves2.h5@03-finetune/vol_12_ves0.h5@03-finetune/vol_12_ves1.h5@03-finetune/vol_12_ves2.h5@03-finetune/vol_12_ves3.h5@03-finetune/vol_12_ves4.h5@03-finetune/vol_12_ves5.h5@03-finetune/vol_13_ves0.h5@03-finetune/vol_13_ves1.h5@03-finetune/vol_13_ves2.h5@03-finetune/vol_14_ves0.h5@03-finetune/vol_14_ves1.h5@03-finetune/vol_14_ves2.h5@03-finetune/vol_14_ves3.h5@03-finetune/vol_14_ves4.h5@03-finetune/vol_15_ves.h5@05-tile_0-188/6-12_ves1.h5@05-tile_0-188/6-12_ves2.h5@05-tile_0-188/6-12_ves4.h5@05-tile_0-188/6-13_ves1.h5@05-tile_0-188/6-13_ves2.h5@05-tile_0-188/6-13_ves3.h5@05-tile_0-188/6-13_ves4.h5@05-tile_0-188/6-13_ves5.h5@05-tile_0-188/6-13_ves6.h5@05-tile_0-188/6-13_ves7.h5
+  VALID_MASK_NAME: 01-initial/vol0_mask.h5@01-initial/vol1_mask.h5@01-initial/vol2_mask.h5@03-finetune/vol_10_mask0.h5@03-finetune/vol_10_mask1.h5@03-finetune/vol_11_mask0.h5@03-finetune/vol_11_mask1.h5@03-finetune/vol_11_mask2.h5@03-finetune/vol_12_mask0.h5@03-finetune/vol_12_mask1.h5@03-finetune/vol_12_mask2.h5@03-finetune/vol_12_mask3.h5@03-finetune/vol_12_mask4.h5@03-finetune/vol_12_mask5.h5@03-finetune/vol_13_mask0.h5@03-finetune/vol_13_mask1.h5@03-finetune/vol_13_mask2.h5@03-finetune/vol_14_mask0.h5@03-finetune/vol_14_mask1.h5@03-finetune/vol_14_mask2.h5@03-finetune/vol_14_mask3.h5@03-finetune/vol_14_mask4.h5@03-finetune/vol_15_mask.h5@05-tile_0-188/6-12_mask1.h5@05-tile_0-188/6-12_mask2.h5@05-tile_0-188/6-12_mask4.h5@05-tile_0-188/6-13_mask1.h5@05-tile_0-188/6-13_mask2.h5@05-tile_0-188/6-13_mask3.h5@05-tile_0-188/6-13_mask4.h5@05-tile_0-188/6-13_mask5.h5@05-tile_0-188/6-13_mask6.h5@05-tile_0-188/6-13_mask7.h5
+  REDUCE_LABEL: False
+  OUTPUT_PATH: /mmfs1/data/linok/hydra/outputs/large_vesicle/04-more_data
+  PAD_SIZE: [0, 0, 0]
+  DATA_SCALE: [1.0, 1.0, 1.0]
+  VALID_RATIO: 0.1
+  DISTRIBUTED: False
+  #REJECT_SAMPLING:
+  #    NUM_TRIAL: 5
+SOLVER:
+  LR_SCHEDULER_NAME: WarmupCosineLR
+  BASE_LR: 0.02
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 10000
+  ITERATION_TOTAL: 1000000
+  SAMPLES_PER_BATCH: 96
+MONITOR:
+  ITERATION_NUM: [40, 400]
+INFERENCE:
+  INPUT_SIZE: [9, 65, 65]
+  OUTPUT_SIZE: [9, 65, 65]
+  INPUT_PATH: /data/projects/weilab/dataset/hydra/big_vesicle/v0-intial_dataset
+  OUTPUT_PATH: outputs/large_vesicle/initial_volumes/results
+  IMAGE_NAME: vol0_im.h5
+  OUTPUT_NAME: vol0_inf.h5
+  PAD_SIZE: [0, 0, 0]
+  AUG_MODE: "mean"
+  AUG_NUM: 4
+  STRIDE: [4, 32, 32]
+  SAMPLES_PER_BATCH: 24

ves_seg/configs/bcd_config.yaml

@@ -0,0 +1,14 @@
+MODEL:
+  OUT_PLANES: 3
+  TARGET_OPT: ['0','4-0-1', '5-3d-1-0-1.0']
+  LOSS_OPTION:
+    - - WeightedBCEWithLogitsLoss
+      - DiceLoss
+    - - WeightedBCEWithLogitsLoss
+      - DiceLoss
+    - - WeightedMSE
+  LOSS_WEIGHT: [[1.0, 0.5], [1.0, 0.5], [5.0]]
+  WEIGHT_OPT: [['1', '0'], ['1', '0'], ['0']]
+  OUTPUT_ACT: [['none', 'sigmoid'], ['none', 'sigmoid'], ['tanh']]
+INFERENCE:
+  OUTPUT_ACT: ['sigmoid', 'sigmoid', 'tanh']