Merge pull request #1 from zudi-lin/master

new pull

.gitignore

@@ -123,6 +123,7 @@ dmypy.json
 # vim
 *.swp
 *.swo
+.vscode/
 .github/
 *~
 

README.md

@@ -16,26 +16,26 @@ If you want new features that are relatively easy to implement (e.g., loss funct
 
 The code is developed and tested under the following configurations.
 - Hardware: 1-8 Nvidia GPUs (with at least 12G GPU memories) (change ```SYSTEM.NUM_GPU``` accordingly)
-- Software: CentOS Linux 7.4 (Core), ***CUDA>=9.0, Python>=3.7, PyTorch>=1.4.0***
+- Software: CentOS Linux 7.4 (Core), ***CUDA>=10.2, Python>=3.7, PyTorch>=1.5.0***
 
 ## Installation
 
 Create a new conda environment:
 ```
 conda create -n py3_torch python=3.7
 source activate py3_torch
-conda install pytorch torchvision cudatoolkit=9.2 -c pytorch
+conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
 ```
-Please note that this package is developed on the Harvard [FASRC](https://www.rc.fas.harvard.edu) cluster, where the current version of the Nvidia driver installed is 396.26 that supports Cuda version 9. More information about GPU computing on the FASRC cluster can be found [here](https://www.rc.fas.harvard.edu/resources/documentation/gpgpu-computing-on-the-cluster/).
+Please note that this package is mainly developed on the Harvard [FASRC](https://www.rc.fas.harvard.edu) cluster. More information about GPU computing on the FASRC cluster can be found [here](https://www.rc.fas.harvard.edu/resources/documentation/gpgpu-computing-on-the-cluster/).
 
 Download and install the package:
 ```
-git clone git@github.com:zudi-lin/pytorch_connectomics.git
+git clone https://github.com/zudi-lin/pytorch_connectomics.git
 cd pytorch_connectomics
 pip install -r requirements.txt
 pip install --editable .
 ```
-For more information and frequently asked questions about installation, please check the [installation guide](https://zudi-lin.github.io/pytorch_connectomics/build/html/notes/installation.html). If you meet compilation errors, please check [TROUBLESHOOTING.md](https://github.com/zudi-lin/pytorch_connectomics/blob/master/TROUBLESHOOTING.md).
+For more information and frequently asked questions about installation, please check the [installation guide](https://zudi-lin.github.io/pytorch_connectomics/build/html/notes/installation.html).
 
 ## Visualization
 
@@ -44,7 +44,7 @@ For more information and frequently asked questions about installation, please c
 * Use TensorBoard with `tensorboard --logdir runs`  (needs to install TensorFlow).
 
 ### Test
-* Visualize the affinity graph and segmentation using Neuroglancer.
+* Visualize the affinity graph and segmentation using [Neuroglancer](https://github.com/google/neuroglancer). 
 
 ## Notes
 

configs/CREMI-Synaptic-Cleft.yaml

@@ -22,25 +22,6 @@ DATASET:
   REJECT_SAMPLING:
     SIZE_THRES: 1000
     P: 0.95
-AUGMENTOR:
-  ROTATE: 
-    ENABLED: True
-  RESCALE:
-    ENABLED: True
-  FLIP:
-    ENABLED: True
-  ELASTIC:
-    ENABLED: True
-  GRAYSCALE:
-    ENABLED: True
-  MISSINGPARTS:
-    ENABLED: True
-  MISSINGSECTION:
-    ENABLED: True
-  MISALIGNMENT:
-    ENABLED: True
-  MOTIONBLUR:
-    ENABLED: True
 SOLVER:
   LR_SCHEDULER_NAME: "WarmupMultiStepLR"
   BASE_LR: 0.001

configs/Lucchi-Mitochondria.yaml

@@ -20,27 +20,13 @@ DATASET:
   OUTPUT_PATH: 'outputs/Lucchi_mito_baseline'
   PAD_SIZE: [56, 56, 56]
 AUGMENTOR:
-  ROTATE: 
-    ENABLED: True
-  RESCALE:
-    ENABLED: True
   # Since the Lucchi dataset is isotropic (each voxel is cubic), we
   # apply flip augmentation also to x-z and z-y axes by default.
   FLIP:
     ENABLED: True
     DO_ZTRANS: 1
-  ELASTIC:
-    ENABLED: True
-  GRAYSCALE:
-    ENABLED: True
-  MISSINGPARTS:
-    ENABLED: True
-  MISSINGSECTION:
-    ENABLED: True
-  MISALIGNMENT:
-    ENABLED: True
-  MOTIONBLUR:
-    ENABLED: True
+  CUTBLUR:
+    DOWNSAMPLE_Z: True
 SOLVER:
   LR_SCHEDULER_NAME: "WarmupMultiStepLR"
   BASE_LR: 0.01

configs/MitoEM/MitoEM-R-A.yaml

@@ -0,0 +1,44 @@
+# For affinity prediction
+# All other configurations are set by default. If you want to add new config options,
+# please modify ../connectomics/config/config.py
+# 30x8x8 nm in (z,y,x), 1000x4096x4096 voxel
+SYSTEM:
+  NUM_GPUS: 4
+  NUM_CPUS: 4
+MODEL:
+  ARCHITECTURE: 'unet_residual_3d'
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IN_PLANES: 1
+  OUT_PLANES: 3
+  LOSS_OPTION: [['WeightedBCE']]
+  TARGET_OPT: ['2']
+  WEIGHT_OPT: [['1']]
+DATASET:
+  IMAGE_NAME: 'im_train.json'
+  LABEL_NAME: 'mito_train.json'
+  INPUT_PATH: '<path/to/data>'
+  OUTPUT_PATH: 'outputs/MitoEM_R_A/'
+  PAD_SIZE: [16, 128, 128]
+  DO_CHUNK_TITLE: 1
+  DATA_CHUNK_NUM: [8, 2, 2]
+  DATA_CHUNK_ITER: 2500
+  LABEL_EROSION: 1
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupMultiStepLR"
+  BASE_LR: 5e-04
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 5000
+  ITERATION_TOTAL: 50000
+  SAMPLES_PER_BATCH: 4
+INFERENCE:
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IMAGE_NAME: 'im_test.json'
+  OUTPUT_PATH: 'outputs/MitoEM_R_A/test/'
+  OUTPUT_NAME: 'result.h5'
+  PAD_SIZE: [16, 128, 128]
+  AUG_MODE: 'mean'
+  AUG_NUM: 4
+  STRIDE: [16, 128, 128]
+  SAMPLES_PER_BATCH: 16

configs/MitoEM/MitoEM-R-AC.yaml

@@ -0,0 +1,45 @@
+# For affinity & instance contour prediction.
+# All other configurations are set by default. If you want to add new config options,
+# please modify ../connectomics/config/config.py
+# 30x8x8 nm in (z,y,x), 1000x4096x4096 voxel
+SYSTEM:
+  NUM_GPUS: 4
+  NUM_CPUS: 4
+MODEL:
+  ARCHITECTURE: 'unet_residual_3d'
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IN_PLANES: 1
+  OUT_PLANES: 4
+  LOSS_OPTION: [['WeightedBCE'],['WeightedBCE','DiceLoss']]
+  LOSS_WEIGHT: [[2.0], [1.0, 1.0]]
+  TARGET_OPT: ['2','4-2-1']
+  WEIGHT_OPT: [['1'],['1','0']]
+DATASET:
+  IMAGE_NAME: 'im_train.json'
+  LABEL_NAME: 'mito_train.json'
+  INPUT_PATH: '<path/to/data>'
+  OUTPUT_PATH: 'outputs/MitoEM_R_AC/'
+  PAD_SIZE: [16, 128, 128]
+  DO_CHUNK_TITLE: 1
+  DATA_CHUNK_NUM: [8, 2, 2]
+  DATA_CHUNK_ITER: 2500
+  LABEL_EROSION: 1
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupMultiStepLR"
+  BASE_LR: 5e-04
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 2500
+  ITERATION_TOTAL: 100000
+  SAMPLES_PER_BATCH: 4
+INFERENCE:
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IMAGE_NAME: 'im_test.json'
+  OUTPUT_PATH: 'outputs/MitoEM_R_AC/test/'
+  OUTPUT_NAME: 'result.h5'
+  PAD_SIZE: [16, 128, 128]
+  AUG_MODE: 'mean'
+  AUG_NUM: 4
+  STRIDE: [16, 128, 128]
+  SAMPLES_PER_BATCH: 16

configs/MitoEM/MitoEM-R-BC.yaml

@@ -0,0 +1,45 @@
+# For binary mask & instance contour prediction.
+# All other configurations are set by default. If you want to add new config options,
+# please modify ../connectomics/config/config.py
+# 30x8x8 nm in (z,y,x), 1000x4096x4096 voxel
+SYSTEM:
+  NUM_GPUS: 8
+  NUM_CPUS: 8
+MODEL:
+  ARCHITECTURE: 'unet_residual_3d'
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IN_PLANES: 1
+  OUT_PLANES: 2
+  LOSS_OPTION: [['WeightedBCE'], ['WeightedBCE']]
+  LOSS_WEIGHT: [[1.0], [1.0]]
+  TARGET_OPT: ['0','4-2-1']
+  WEIGHT_OPT: [['1'],['1']]
+DATASET:
+  IMAGE_NAME: 'im_train.json'
+  LABEL_NAME: 'mito_train.json'
+  INPUT_PATH: '<path/to/data>'
+  OUTPUT_PATH: 'outputs/MitoEM_R_BC/'
+  PAD_SIZE: [16, 128, 128]
+  DO_CHUNK_TITLE: 1
+  DATA_CHUNK_NUM: [8, 2, 2]
+  DATA_CHUNK_ITER: 2500
+  LABEL_EROSION: 1
+SOLVER:
+  LR_SCHEDULER_NAME: "WarmupMultiStepLR"
+  BASE_LR: 1e-03
+  ITERATION_STEP: 1
+  ITERATION_SAVE: 2500
+  ITERATION_TOTAL: 100000
+  SAMPLES_PER_BATCH: 4
+INFERENCE:
+  INPUT_SIZE: [32, 256, 256]
+  OUTPUT_SIZE: [32, 256, 256]
+  IMAGE_NAME: 'im_test.json'
+  OUTPUT_PATH: 'outputs/MitoEM_R_BC/test/'
+  OUTPUT_NAME: 'result.h5'
+  PAD_SIZE: [16, 128, 128]
+  AUG_MODE: 'mean'
+  AUG_NUM: 4
+  STRIDE: [16, 128, 128]
+  SAMPLES_PER_BATCH: 16

configs/MitoEM/README.md

@@ -0,0 +1,9 @@
+## MitoEM Dataset: Large-scale 3D Mitochondria Instance Segmentation from EM Images
+
+### Introduction
+
+Serial electron microscopy (EM) allows identification of intracellular organelles such as mitochondria, which provides novel insights for both clinical and scientific studies. However, the mitochondria reconstruction benchmark only contains around 100 instances that are well-separated and exhibit simple morphologies. Therefore, existing automatic methods that have achieved almost human-level performance on the small dataset usually fail to produce preferred results due to object diversity in appearance and morphologies.
+
+To enable the development of robust models for large-scale biomedical analysis, we introduce **MitoEM**, a 3D mitochondria instance segmentation dataset consisting of two 30μm cubic volumes from human and rat cortices respectively, which are **3,600x** larger than the previous benchmark dataset. Our new dataset posts new challenges for existing state-of-the-art segmentation approaches as they consistently fail to generate object masks with quality on par with expert annotators. With approximately 40k mitochondria in our new dataset, we provide in-depth analysis of the dataset properties, as well as the performance of different combinations of deep learning models and post-processing methods. The MitoEM dataset and our comprehensive analysis will enable further researches in large-scale instance segmentation and a better understanding of mammalian brains.
+
+The configuration files provided in this folder are used to reproduce the results in the paper. We will add detailed usage soon!

configs/SNEMI-Neuron.yaml

@@ -13,30 +13,11 @@ MODEL:
   TARGET_OPT: ['2']
   WEIGHT_OPT: [['1']]
 DATASET:
-  IMAGE_NAME: 'img/train-input.tif'
-  LABEL_NAME: 'label/train_labels.tif'
+  IMAGE_NAME: 'train_image.h5'
+  LABEL_NAME: 'train_label.h5'
   INPUT_PATH: '/path/to/SNEMI/'
   OUTPUT_PATH: 'outputs/SNEMI/'
   PAD_SIZE: [4, 128, 128]
-AUGMENTOR:
-  ROTATE: 
-    ENABLED: True
-  RESCALE:
-    ENABLED: True
-  FLIP:
-    ENABLED: True
-  ELASTIC:
-    ENABLED: True
-  GRAYSCALE:
-    ENABLED: True
-  MISSINGPARTS:
-    ENABLED: True
-  MISSINGSECTION:
-    ENABLED: True
-  MISALIGNMENT:
-    ENABLED: True
-  MOTIONBLUR:
-    ENABLED: True
 SOLVER:
   LR_SCHEDULER_NAME: "MultiStepLR"
   BASE_LR: 0.001

configs/Synaptic-Partner-Segmentation.yaml

@@ -23,29 +23,9 @@ DATASET:
   INPUT_PATH: '<path/to/data>'
   OUTPUT_PATH: 'outputs/synaptic_polarity/'
   PAD_SIZE: [4, 128, 128]
-  DATA_SCALE: [1.0, 0.5, 0.5]
   REJECT_SAMPLING:
     SIZE_THRES: 1000
     P: 0.95
-AUGMENTOR:
-  ROTATE: 
-    ENABLED: True
-  RESCALE:
-    ENABLED: True
-  FLIP:
-    ENABLED: True
-  ELASTIC:
-    ENABLED: True
-  GRAYSCALE:
-    ENABLED: True
-  MISSINGPARTS:
-    ENABLED: True
-  MISSINGSECTION:
-    ENABLED: True
-  MISALIGNMENT:
-    ENABLED: True
-  MOTIONBLUR:
-    ENABLED: True
 SOLVER:
   LR_SCHEDULER_NAME: "WarmupMultiStepLR"
   BASE_LR: 0.001
@@ -57,11 +37,13 @@ SOLVER:
 MONITOR:
   VIS_OPT: [1, 16]
 INFERENCE:
+  IMAGE_NAME: '<name/of/image (h5py volumes)>'
+  OUTPUT_PATH: 'outputs/synaptic_polarity/test'
   AUG_MODE: 'mean'
   AUG_NUM: 4
   OUTPUT_NAME: 'syn_polarity_pred.h5'
   STRIDE: [4, 128, 128]
-  SAMPLES_PER_BATCH: 32
+  SAMPLES_PER_BATCH: 16
   INPUT_SIZE: [8, 256, 256]
   OUTPUT_SIZE: [8, 256, 256]
   PAD_SIZE: [4, 128, 128]