
# 3D MRI Segmentation using Zoom-in&out training strategy
This is PyTorch implementation of Deep Residual 3D U-Net for MRI segmentation. 

This is also the source code for *[Automatic post-stroke lesion segmentation on MR images using 3D residual convolutional neural network](https://doi.org/10.1016/j.nicl.2020.102276)*


![Visualization of Prediction Results](https://github.com/BMIRDS/3dMRISegmentation/blob/master/assets/fig2.png?raw=true)


## Highlights
+ We developed a deep learning model for MRI 3D lesion segmentation of chronic stroke.
+ Our novel zoom-in&out strategy increases performance and accelerates optimization.
+ High performance was achieved in both volumetric and surface-based metrics.

![Zoom-in&out training strategy for volumetric segmentation](https://github.com/BMIRDS/3dMRISegmentation/blob/master/assets/graphical_abstract.png?raw=true)

## Zoom-in&out training strategy for efficient volumetric segmentation
We used a two-stage zoom-in&out strategy to first train them on small volumes, and then we finetuned the models on larger volumes. Advantages are:
- Training models with smaller volumes can have a regularizing effect due to performing data augmentation by randomly extracting diverse sub-volumes from original volumes.
- A “zoom-in” step is a computationally inexpensive option and can utilize sub-optimal graphic processing units (GPUs) for the task. By feeding smaller volumes to older but more accessible GPUs, models can be trained in parallel, and, as a result, are faster. 
- The “zoom-out” stage involves showing models larger volumes to learn from the broader context of input images and improves the robustness of the model.

##  The details for the model optimization in our experiments.
<table>
<thead>
  <tr>
    <th>Optimization Stage</th>
    <th>Zoom-In Stage</th>
    <th>Zoom-Out Stage</th>
  </tr>
</thead>
<tbody>
  <tr>
    <td>Input volume size (mm3) </td>
    <td>128 × 128 × 128 (24% sub-volume)  </td>
    <td>144 × 172 × 168 (48% sub-volume)  </td>
  </tr>
  <tr>
    <td>Training Length</td>
    <td>1200 epochs</td>
    <td>150 epochs</td>
  </tr>
  <tr>
    <td>Initial learning rate</td>
    <td>1.00E−03</td>
    <td>1.00E−04 </td>
  </tr>
  <tr>
    <td>Optimizer</td>
    <td colspan="2">Adam optimizer and cosine annealing with warm restart scheduler</td>
  </tr>
  <tr>
    <td>GPU</td>
    <td>Nvidia Titan Xp with 12 GB memory</td>
    <td>Nvidia Titan RTX with 24 GB memory</td>
  </tr>
</tbody>
</table>

##  Effectiveness of zoom-in&out strategy
<table>
<thead>
  <tr>
    <th>Methods</th>
    <th>microDSC</th>
    <th>DSC</th>
    <th>HD (mm)</th>
    <th>ASSD (mm)</th>
    <th>TPR</th>
    <th>Precision</th>
  </tr>
</thead>
<tbody>
  <tr>
    <td>3D-ResU-Net w/o zoom-out</td>
    <td>0.73</td>
    <td>0.60 (0.47–0.73)</td>
    <td>35.1 (20.4–51.3)</td>
    <td>7.6 (3.7–12.3)</td>
    <td><b>0.83 (0.71–0.91)</b></td>
    <td>0.54 (0.39–0.67)</td>
  </tr>
  <tr>
    <td>3D-ResU-Net w/ zoom-in&out</td>
    <td><b>0.79</b></td>
    <td><b>0.64 (0.51–0.76)</b></td>
    <td><b>20.4 (10.0–33.3)</b></td>
    <td><b>3.6 (1.7–6.2)</b></td>
    <td>0.81 (0.68–0.89)</td>
    <td><b>0.62 (0.48–0.74)</b></td>
  </tr>
  <tr>
    <td>Δ</td>
    <td>+0.06</td>
    <td>+0.04</td>
    <td>−14.7</td>
    <td>−4.0</td>
    <td>−0.02</td>
    <td>+0.08</td>
  </tr>
</tbody>
</table>
Best scores are marked in bold.

## Visualization of segmentation results
![Visualization of segmentation results1](assets/video1.gif)
![Visualization of segmentation results2](assets/video2.gif)

## Dependencies
- Python 3.6
- [PyTorch 1.3](https://pytorch.org/)
- [NiBabel](https://nipy.org/nibabel/)
- [scikit-image](https://scikit-image.org/)
- [scikit-learn](https://scikit-learn.org/)


## Usage
## Training a 3D segmentation model
- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots>`

## Applying zoom-in&out strategy
Zoom-in stage
- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xi yi zi`
where xi, yi, zi are the size of input volume in the zoom-in stage.

Zoom-out stage
- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xo yo zo --resume_model <path_to_a_snapshot>`
where xo, yo, zo are the size of input volume in the zoom-out stage.

## Volume size selection in Zoom-In&Out training strategy
Here are our suggestions for selecting an optimal volume size in Zoom-In&Out training strategy
- Visualizing typical images from the dataset, and then reviewing the spatial distribution of target lesions to estimate the possible effective input volume size for a specific application. (In our study, we observed in early experiments that using volume sizes significantly smaller than those in our final configuration does not achieve a competitive performance. This performance gap may indicate that volumes cropped for training during these experiments were overwhelmed by regions without positive stroke lesions, which could have negatively biased our model optimization in the early experiments.)
- Performing a systematic search on proportions of original volume size (e.g., 25%, 50%, 75% ) to find an optimal configuration.


## Testing a trained model
- `python test.py --data_dir <path_to_dataset> --model_path <path_to_a_snapshot>`


# Citations
3dMRISegmentation is an open-source library and is licensed under the GNU General Public License (v3). 

If you are using this library, please cite:

```Naofumi Tomita, Steven Jiang, Matthew E. Maeder, and Saeed Hassanpour. "Automatic Post-Stroke Lesion Segmentation on MR Images using 3D Residual Convolutional Neural Network." NeuroImage: Clinical (2020): 102276.```

	a/README.md		b/README.md
1	# 3D MRI Segmentation using Zoom-in&out training strategy	1	# 3D MRI Segmentation using Zoom-in&out training strategy
2	This is PyTorch implementation of Deep Residual 3D U-Net for MRI segmentation.	2	This is PyTorch implementation of Deep Residual 3D U-Net for MRI segmentation.
3		3
4	This is also the source code for [Automatic post-stroke lesion segmentation on MR images using 3D residual convolutional neural network](https://doi.org/10.1016/j.nicl.2020.102276)	4	This is also the source code for [Automatic post-stroke lesion segmentation on MR images using 3D residual convolutional neural network](https://doi.org/10.1016/j.nicl.2020.102276)
5		5
6		6
7	![Visualization of Prediction Results](assets/fig2.png)	7	![Visualization of Prediction Results](https://github.com/BMIRDS/3dMRISegmentation/blob/master/assets/fig2.png?raw=true)
8		8
9		9
10	## Highlights	10	## Highlights
11	+ We developed a deep learning model for MRI 3D lesion segmentation of chronic stroke.	11	+ We developed a deep learning model for MRI 3D lesion segmentation of chronic stroke.
12	+ Our novel zoom-in&out strategy increases performance and accelerates optimization.	12	+ Our novel zoom-in&out strategy increases performance and accelerates optimization.
13	+ High performance was achieved in both volumetric and surface-based metrics.	13	+ High performance was achieved in both volumetric and surface-based metrics.
14		14
15	![Zoom-in&out training strategy for volumetric segmentation](assets/graphical_abstract.png)	15	![Zoom-in&out training strategy for volumetric segmentation](https://github.com/BMIRDS/3dMRISegmentation/blob/master/assets/graphical_abstract.png?raw=true)
16		16
17	## Zoom-in&out training strategy for efficient volumetric segmentation	17	## Zoom-in&out training strategy for efficient volumetric segmentation
18	We used a two-stage zoom-in&out strategy to first train them on small volumes, and then we finetuned the models on larger volumes. Advantages are:	18	We used a two-stage zoom-in&out strategy to first train them on small volumes, and then we finetuned the models on larger volumes. Advantages are:
19	- Training models with smaller volumes can have a regularizing effect due to performing data augmentation by randomly extracting diverse sub-volumes from original volumes.	19	- Training models with smaller volumes can have a regularizing effect due to performing data augmentation by randomly extracting diverse sub-volumes from original volumes.
20	- A “zoom-in” step is a computationally inexpensive option and can utilize sub-optimal graphic processing units (GPUs) for the task. By feeding smaller volumes to older but more accessible GPUs, models can be trained in parallel, and, as a result, are faster.	20	- A “zoom-in” step is a computationally inexpensive option and can utilize sub-optimal graphic processing units (GPUs) for the task. By feeding smaller volumes to older but more accessible GPUs, models can be trained in parallel, and, as a result, are faster.
21	- The “zoom-out” stage involves showing models larger volumes to learn from the broader context of input images and improves the robustness of the model.	21	- The “zoom-out” stage involves showing models larger volumes to learn from the broader context of input images and improves the robustness of the model.
22		22
23	## The details for the model optimization in our experiments.	23	## The details for the model optimization in our experiments.
24	<table>	24	<table>
25	<thead>	25	<thead>
26	<tr>	26	<tr>
27	<th>Optimization Stage</th>	27	<th>Optimization Stage</th>
28	<th>Zoom-In Stage</th>	28	<th>Zoom-In Stage</th>
29	<th>Zoom-Out Stage</th>	29	<th>Zoom-Out Stage</th>
30	</tr>	30	</tr>
31	</thead>	31	</thead>
32	<tbody>	32	<tbody>
33	<tr>	33	<tr>
34	<td>Input volume size (mm3) </td>	34	<td>Input volume size (mm3) </td>
35	<td>128 × 128 × 128 (24% sub-volume) </td>	35	<td>128 × 128 × 128 (24% sub-volume) </td>
36	<td>144 × 172 × 168 (48% sub-volume) </td>	36	<td>144 × 172 × 168 (48% sub-volume) </td>
37	</tr>	37	</tr>
38	<tr>	38	<tr>
39	<td>Training Length</td>	39	<td>Training Length</td>
40	<td>1200 epochs</td>	40	<td>1200 epochs</td>
41	<td>150 epochs</td>	41	<td>150 epochs</td>
42	</tr>	42	</tr>
43	<tr>	43	<tr>
44	<td>Initial learning rate</td>	44	<td>Initial learning rate</td>
45	<td>1.00E−03</td>	45	<td>1.00E−03</td>
46	<td>1.00E−04 </td>	46	<td>1.00E−04 </td>
47	</tr>	47	</tr>
48	<tr>	48	<tr>
49	<td>Optimizer</td>	49	<td>Optimizer</td>
50	<td colspan="2">Adam optimizer and cosine annealing with warm restart scheduler</td>	50	<td colspan="2">Adam optimizer and cosine annealing with warm restart scheduler</td>
51	</tr>	51	</tr>
52	<tr>	52	<tr>
53	<td>GPU</td>	53	<td>GPU</td>
54	<td>Nvidia Titan Xp with 12 GB memory</td>	54	<td>Nvidia Titan Xp with 12 GB memory</td>
55	<td>Nvidia Titan RTX with 24 GB memory</td>	55	<td>Nvidia Titan RTX with 24 GB memory</td>
56	</tr>	56	</tr>
57	</tbody>	57	</tbody>
58	</table>	58	</table>
59		59
60	## Effectiveness of zoom-in&out strategy	60	## Effectiveness of zoom-in&out strategy
61	<table>	61	<table>
62	<thead>	62	<thead>
63	<tr>	63	<tr>
64	<th>Methods</th>	64	<th>Methods</th>
65	<th>microDSC</th>	65	<th>microDSC</th>
66	<th>DSC</th>	66	<th>DSC</th>
67	<th>HD (mm)</th>	67	<th>HD (mm)</th>
68	<th>ASSD (mm)</th>	68	<th>ASSD (mm)</th>
69	<th>TPR</th>	69	<th>TPR</th>
70	<th>Precision</th>	70	<th>Precision</th>
71	</tr>	71	</tr>
72	</thead>	72	</thead>
73	<tbody>	73	<tbody>
74	<tr>	74	<tr>
75	<td>3D-ResU-Net w/o zoom-out</td>	75	<td>3D-ResU-Net w/o zoom-out</td>
76	<td>0.73</td>	76	<td>0.73</td>
77	<td>0.60 (0.47–0.73)</td>	77	<td>0.60 (0.47–0.73)</td>
78	<td>35.1 (20.4–51.3)</td>	78	<td>35.1 (20.4–51.3)</td>
79	<td>7.6 (3.7–12.3)</td>	79	<td>7.6 (3.7–12.3)</td>
80	<td><b>0.83 (0.71–0.91)</b></td>	80	<td><b>0.83 (0.71–0.91)</b></td>
81	<td>0.54 (0.39–0.67)</td>	81	<td>0.54 (0.39–0.67)</td>
82	</tr>	82	</tr>
83	<tr>	83	<tr>
84	<td>3D-ResU-Net w/ zoom-in&out</td>	84	<td>3D-ResU-Net w/ zoom-in&out</td>
85	<td><b>0.79</b></td>	85	<td><b>0.79</b></td>
86	<td><b>0.64 (0.51–0.76)</b></td>	86	<td><b>0.64 (0.51–0.76)</b></td>
87	<td><b>20.4 (10.0–33.3)</b></td>	87	<td><b>20.4 (10.0–33.3)</b></td>
88	<td><b>3.6 (1.7–6.2)</b></td>	88	<td><b>3.6 (1.7–6.2)</b></td>
89	<td>0.81 (0.68–0.89)</td>	89	<td>0.81 (0.68–0.89)</td>
90	<td><b>0.62 (0.48–0.74)</b></td>	90	<td><b>0.62 (0.48–0.74)</b></td>
91	</tr>	91	</tr>
92	<tr>	92	<tr>
93	<td>Δ</td>	93	<td>Δ</td>
94	<td>+0.06</td>	94	<td>+0.06</td>
95	<td>+0.04</td>	95	<td>+0.04</td>
96	<td>−14.7</td>	96	<td>−14.7</td>
97	<td>−4.0</td>	97	<td>−4.0</td>
98	<td>−0.02</td>	98	<td>−0.02</td>
99	<td>+0.08</td>	99	<td>+0.08</td>
100	</tr>	100	</tr>
101	</tbody>	101	</tbody>
102	</table>	102	</table>
103	Best scores are marked in bold.	103	Best scores are marked in bold.
104		104
105	## Visualization of segmentation results	105	## Visualization of segmentation results
106	![Visualization of segmentation results1](assets/video1.gif)	106	![Visualization of segmentation results1](assets/video1.gif)
107	![Visualization of segmentation results2](assets/video2.gif)	107	![Visualization of segmentation results2](assets/video2.gif)
108		108
109	## Dependencies	109	## Dependencies
110	- Python 3.6	110	- Python 3.6
111	- [PyTorch 1.3](https://pytorch.org/)	111	- [PyTorch 1.3](https://pytorch.org/)
112	- [NiBabel](https://nipy.org/nibabel/)	112	- [NiBabel](https://nipy.org/nibabel/)
113	- [scikit-image](https://scikit-image.org/)	113	- [scikit-image](https://scikit-image.org/)
114	- [scikit-learn](https://scikit-learn.org/)	114	- [scikit-learn](https://scikit-learn.org/)
115		115
116		116
117	## Usage	117	## Usage
118	## Training a 3D segmentation model	118	## Training a 3D segmentation model
119	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots>`	119	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots>`
120		120
121	## Applying zoom-in&out strategy	121	## Applying zoom-in&out strategy
122	Zoom-in stage	122	Zoom-in stage
123	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xi yi zi`	123	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xi yi zi`
124	where xi, yi, zi are the size of input volume in the zoom-in stage.	124	where xi, yi, zi are the size of input volume in the zoom-in stage.
125		125
126	Zoom-out stage	126	Zoom-out stage
127	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xo yo zo --resume_model <path_to_a_snapshot>`	127	- `python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xo yo zo --resume_model <path_to_a_snapshot>`
128	where xo, yo, zo are the size of input volume in the zoom-out stage.	128	where xo, yo, zo are the size of input volume in the zoom-out stage.
129		129
130	## Volume size selection in Zoom-In&Out training strategy	130	## Volume size selection in Zoom-In&Out training strategy
131	Here are our suggestions for selecting an optimal volume size in Zoom-In&Out training strategy	131	Here are our suggestions for selecting an optimal volume size in Zoom-In&Out training strategy
132	- Visualizing typical images from the dataset, and then reviewing the spatial distribution of target lesions to estimate the possible effective input volume size for a specific application. (In our study, we observed in early experiments that using volume sizes significantly smaller than those in our final configuration does not achieve a competitive performance. This performance gap may indicate that volumes cropped for training during these experiments were overwhelmed by regions without positive stroke lesions, which could have negatively biased our model optimization in the early experiments.)	132	- Visualizing typical images from the dataset, and then reviewing the spatial distribution of target lesions to estimate the possible effective input volume size for a specific application. (In our study, we observed in early experiments that using volume sizes significantly smaller than those in our final configuration does not achieve a competitive performance. This performance gap may indicate that volumes cropped for training during these experiments were overwhelmed by regions without positive stroke lesions, which could have negatively biased our model optimization in the early experiments.)
133	- Performing a systematic search on proportions of original volume size (e.g., 25%, 50%, 75% ) to find an optimal configuration.	133	- Performing a systematic search on proportions of original volume size (e.g., 25%, 50%, 75% ) to find an optimal configuration.
134		134
135		135
136	## Testing a trained model	136	## Testing a trained model
137	- `python test.py --data_dir <path_to_dataset> --model_path <path_to_a_snapshot>`	137	- `python test.py --data_dir <path_to_dataset> --model_path <path_to_a_snapshot>`
138		138
139		139
140	# Citations	140	# Citations
141	3dMRISegmentation is an open-source library and is licensed under the GNU General Public License (v3).	141	3dMRISegmentation is an open-source library and is licensed under the GNU General Public License (v3).
142		142
143	If you are using this library, please cite:	143	If you are using this library, please cite:
144		144
145	```Naofumi Tomita, Steven Jiang, Matthew E. Maeder, and Saeed Hassanpour. "Automatic Post-Stroke Lesion Segmentation on MR Images using 3D Residual Convolutional Neural Network." NeuroImage: Clinical (2020): 102276.```	145	```Naofumi Tomita, Steven Jiang, Matthew E. Maeder, and Saeed Hassanpour. "Automatic Post-Stroke Lesion Segmentation on MR Images using 3D Residual Convolutional Neural Network." NeuroImage: Clinical (2020): 102276.```