Wei Huang, Chang Chen, Zhiwei Xiong(*), Yueyi Zhang, Xuejin Chen, Xiaoyan Sun, Feng Wu
*Corresponding Author
University of Science and Technology of China (USTC)
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.
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.4| 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 |
Take the training on the AC3 dataset as an example.
python pre_training.py -c=pretraining_snemi3dWeight Sharing (WS)
python main.py -c=seg_snemi3d_d5_u200EMA
python main_ema.py -c=seg_snemi3d_d5_1024_u200_emaTake the validation on the AC3 dataset as an example.
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-ac3python2 evaluate_waterz.py -mn=seg_ac3_d5_1024_u200_WS -id=seg_ac3_d5_1024_u200_WS -m=snemi3d-ac3python evaluate_lmc.py -mn=seg_ac3_d5_1024_u200_WS -id=seg_ac3_d5_1024_u200_WS -m=snemi3d-ac3We 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 |
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.
If you have any problem with the released code, please do not hesitate to contact me by email (weih527@mail.ustc.edu.cn).