by Rodney LaLonde and Ulas Bagci
The original paper for SegCaps can be found at https://arxiv.org/abs/1804.04241.
A project page for this work can be found at https://rodneylalonde.wixsite.com/personal/research-blog/capsules-for-object-segmentation.
Convolutional neural networks (CNNs) have shown remarkable results over the last several years for a wide range of computer vision tasks. A new architecture recently introduced by Sabour et al., referred to as a capsule networks with dynamic routing, has shown great initial results for digit recognition and small image classification. Our work expands the use of capsule networks to the task of object segmentation for the first time in the literature. We extend the idea of convolutional capsules with locally-connected routing and propose the concept of deconvolutional capsules. Further, we extend the masked reconstruction to reconstruct the positive input class. The proposed convolutional-deconvolutional capsule network, called SegCaps, shows strong results for the task of object segmentation with substantial decrease in parameter space. As an example application, we applied the proposed SegCaps to segment pathological lungs from low dose CT scans and compared its accuracy and efficiency with other U-Net-based architectures. SegCaps is able to handle large image sizes (512 x 512) as opposed to baseline capsules (typically less than 32 x 32). The proposed SegCaps reduced the number of parameters of U-Net architecture by 95.4% while still providing a better segmentation accuracy.
| Method | Parameters | Split-0 (%) | Split-1 (%) | Split-2 (%) | Split-3 (%) | Average (%) |
|---|---|---|---|---|---|---|
| U-Net | 31.0 M | 98.353 | 98.432 | 98.476 | 98.510 | 98.449 |
| Tiramisu | 2.3 M | 98.394 | 98.358 | 98.543 | 98.339 | 98.410 |
| Baseline Caps | 1.7 M | 82.287 | 79.939 | 95.121 | 83.608 | 83.424 |
| SegCaps (R1) | 1.4 M | 98.471 | 98.444 | 98.401 | 98.362 | 98.419 |
| SegCaps (R3) | 1.4 M | 98.499 | 98.523 | 98.455 | 98.474 | 98.479 |
This repo of code is written for Keras using the TensorFlow backend. Please install all required packages before using this code.
pip install -r requirements.txtInside the data root folder (i.e. where you have your data stored) you should have two folders: one called imgs and one called masks. All models, results, etc. are saved to this same root directory.
From the main file (main.py) you can train, test, and manipulate the segmentation capsules of various networks. Simply set the --train, --test, or --manip flags to 0 or 1 to turn these off or on respectively. The argument --data_root_dir is the only required argument and should be set to the directory containing your imgs and masks folders. There are many more arguments that can be set and these are all explained in the main.py file.
If you use significant portions of our code or ideas from our paper in your research, please cite our work:
@article{lalonde2018capsules,
title={Capsules for Object Segmentation},
author={LaLonde, Rodney and Bagci, Ulas},
journal={arXiv preprint arXiv:1804.04241},
year={2018}
}
Please direct any questions or comments to me; I am happy to help in any way I can. You can either comment on the project page, or email me directly at [email protected].