Semi-Supervised Neuron Segmentation via Reinforced Consistency Learning
Wei Huang, Chang Chen, Zhiwei Xiong(*), Yueyi Zhang, Xuejin Chen, Xiaoyan Sun, Feng Wu
*Corresponding Author
University of Science and Technology of China (USTC)
Introduction
This repository is the official implementation of the paper, "Semi-Supervised Neuron Segmentation via Reinforced Consistency Learning", where more visual results and implementation details are presented.
Installation
This code was tested with Pytorch 1.0.1 (later versions may work), CUDA 9.0, Python 3.7.4 and Ubuntu 16.04. It is worth mentioning that, besides some commonly used image processing packages, you also need to install some special post-processing packages for neuron segmentation, such as waterz and elf.
If you have a Docker environment, we strongly recommend you to pull our image as follows,
docker pull registry.cn-hangzhou.aliyuncs.com/renwu527/auto-emseg:v5.4or
docker pull renwu527/auto-emseg:v5.4Dataset
| Datasets | Sizes | Resolutions | Species | Download (Processed) |
|---|---|---|---|---|
| AC3/AC4 | 1024x1024x256, 1024x1024x100 | 6x6x30 nm^3 | Mouse | BaiduYun (Access code: weih) or GoogleDrive |
| CREMI | 1250x1250x125 (x3) | 4x4x40 nm^3 | Drosophila | BaiduYun (Access code: weih) or GoogleDrive |
| Kasthuri15 | 10747x12895x1850 | 6x6x30 nm^3 | Mouse | BaiduYun (Access code: weih) or GoogleDrive |
Training stage
Take the training on the AC3 dataset as an example.
1. Pre-training
python pre_training.py -c=pretraining_snemi3d2. Consistency learning
Weight Sharing (WS)
python main.py -c=seg_snemi3d_d5_u200EMA
python main_ema.py -c=seg_snemi3d_d5_1024_u200_emaValidation stage
Take the validation on the AC3 dataset as an example.
1. Predict affinities
python inference.py -c=seg_snemi3d_d5_1024_u200 -mn=seg_ac3_d5_1024_u200_WS -id=seg_ac3_d5_1024_u200_WS -m=snemi3d-ac32. Evaluate on Waterz
python2 evaluate_waterz.py -mn=seg_ac3_d5_1024_u200_WS -id=seg_ac3_d5_1024_u200_WS -m=snemi3d-ac33. Evaluate on LMC
python evaluate_lmc.py -mn=seg_ac3_d5_1024_u200_WS -id=seg_ac3_d5_1024_u200_WS -m=snemi3d-ac3Model Zoo
We provide the trained models on the AC3 dataset at BaiduYun and GoogleDrive, including the pre-trained model and the segmentation models on different numbers of labeled (*L) and unlabeled (*U) sections (1024x1024).
| Methods | Models | Download |
|---|---|---|
| pre-training | pretraining_snemi3d.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 5L+200U (WS) | seg_ac3_d5_1024_u200_WS.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 5L+200U (EMA) | seg_ac3_d5_1024_u200_EMA.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 10L | seg_kasthuri15_d10.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 10L+200U | seg_kasthuri15_d10_u200.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 50L+200U | seg_kasthuri15_d50_u200.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 100L | seg_kasthuri15_d100.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 100L+200U | seg_kasthuri15_d100_u200.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
| 100L+1000U | seg_kasthuri15_d100_u1000.ckpt | BaiduYun (Access code: weih) or GoogleDrive |
More visual results on the Kasthuri15 dataset
To demonstrate the generalizability performance of our method on the large-scale EM data, we test our models on the Kasthuri15 dataset. The quantitative results can be found in our paper. Here, we provide more visual results on the Subset3 dataset to qualitatively demonstrate the superiority of our semi-supervised method compared with the existing supervised method with full labeled data (100L).
Left images are the results of the supervised method (100L), while right images are the results of our semi-supervised method (100L+1000U), where blue and red arrows represent split and merge errors, respectively.
Related Projects
Contact
If you have any problem with the released code, please do not hesitate to contact me by email (weih527@mail.ustc.edu.cn).