Skip to content

Commit 49515e7

Browse files
Junjun2016Junjun2016
authored
Add more UNet-based medical segmentation benchmark (#324)
* Add UNet as backbone and FCN PSPNet DeepLabV3 as decode_head benchmark on 4 retinal vessel segmentation datasets * adjust README of UNet
1 parent da0af50 commit 49515e7

16 files changed

+185
-10
lines changed

‎configs/_base_/models/deeplabv3_unet_s5-d16.py

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# model settings
2+
norm_cfg = dict(type='SyncBN', requires_grad=True)
3+
model = dict(
4+
type='EncoderDecoder',
5+
pretrained=None,
6+
backbone=dict(
7+
type='UNet',
8+
in_channels=3,
9+
base_channels=64,
10+
num_stages=5,
11+
strides=(1, 1, 1, 1, 1),
12+
enc_num_convs=(2, 2, 2, 2, 2),
13+
dec_num_convs=(2, 2, 2, 2),
14+
downsamples=(True, True, True, True),
15+
enc_dilations=(1, 1, 1, 1, 1),
16+
dec_dilations=(1, 1, 1, 1),
17+
with_cp=False,
18+
conv_cfg=None,
19+
norm_cfg=norm_cfg,
20+
act_cfg=dict(type='ReLU'),
21+
upsample_cfg=dict(type='InterpConv'),
22+
norm_eval=False),
23+
decode_head=dict(
24+
type='ASPPHead',
25+
in_channels=64,
26+
in_index=4,
27+
channels=16,
28+
dilations=(1, 12, 24, 36),
29+
dropout_ratio=0.1,
30+
num_classes=2,
31+
norm_cfg=norm_cfg,
32+
align_corners=False,
33+
loss_decode=dict(
34+
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
35+
auxiliary_head=dict(
36+
type='FCNHead',
37+
in_channels=128,
38+
in_index=3,
39+
channels=64,
40+
num_convs=1,
41+
concat_input=False,
42+
dropout_ratio=0.1,
43+
num_classes=2,
44+
norm_cfg=norm_cfg,
45+
align_corners=False,
46+
loss_decode=dict(
47+
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)))
48+
# model training and testing settings
49+
train_cfg = dict()
50+
test_cfg = dict(mode='slide', crop_size=256, stride=170)

‎configs/_base_/models/unet_s5-d16.py ‎configs/_base_/models/fcn_unet_s5-d16.py

File renamed without changes.

‎configs/_base_/models/pspnet_unet_s5-d16.py

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# model settings
2+
norm_cfg = dict(type='SyncBN', requires_grad=True)
3+
model = dict(
4+
type='EncoderDecoder',
5+
pretrained=None,
6+
backbone=dict(
7+
type='UNet',
8+
in_channels=3,
9+
base_channels=64,
10+
num_stages=5,
11+
strides=(1, 1, 1, 1, 1),
12+
enc_num_convs=(2, 2, 2, 2, 2),
13+
dec_num_convs=(2, 2, 2, 2),
14+
downsamples=(True, True, True, True),
15+
enc_dilations=(1, 1, 1, 1, 1),
16+
dec_dilations=(1, 1, 1, 1),
17+
with_cp=False,
18+
conv_cfg=None,
19+
norm_cfg=norm_cfg,
20+
act_cfg=dict(type='ReLU'),
21+
upsample_cfg=dict(type='InterpConv'),
22+
norm_eval=False),
23+
decode_head=dict(
24+
type='PSPHead',
25+
in_channels=64,
26+
in_index=4,
27+
channels=16,
28+
pool_scales=(1, 2, 3, 6),
29+
dropout_ratio=0.1,
30+
num_classes=2,
31+
norm_cfg=norm_cfg,
32+
align_corners=False,
33+
loss_decode=dict(
34+
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
35+
auxiliary_head=dict(
36+
type='FCNHead',
37+
in_channels=128,
38+
in_index=3,
39+
channels=64,
40+
num_convs=1,
41+
concat_input=False,
42+
dropout_ratio=0.1,
43+
num_classes=2,
44+
norm_cfg=norm_cfg,
45+
align_corners=False,
46+
loss_decode=dict(
47+
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)))
48+
# model training and testing settings
49+
train_cfg = dict()
50+
test_cfg = dict(mode='slide', crop_size=256, stride=170)

‎configs/unet/README.md

+31-6
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,34 @@
1717

1818
## Results and models
1919

20-
| Backbone | Head | Dataset | Image Size | Crop Size | Stride | Lr schd | Mem (GB) | Inf time (fps) | Dice | download |
21-
|--------|----------|----------|----------|-----------|--------:|----------|----------------|------:|--------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
22-
| UNet-S5-D16 | FCN | DRIVE | 584x565 | 64x64 | 42x42 | 40000 | 0.680 | - | 78.67 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_64x64_40k_drive/unet_s5-d16_64x64_40k_drive_20201223_191051-9cd163b8.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_64x64_40k_drive/unet_s5-d16_64x64_40k_drive-20201223_191051.log.json) |
23-
| UNet-S5-D16 | FCN | STARE | 605x700 | 128x128 | 85x85 | 40000 | 0.968 | - | 81.02 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_128x128_40k_stare/unet_s5-d16_128x128_40k_stare_20201223_191051-e5439846.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_128x128_40k_stare/unet_s5-d16_128x128_40k_stare-20201223_191051.log.json) |
24-
| UNet-S5-D16 | FCN | CHASE_DB1 | 960x999 | 128x128 | 85x85 | 40000 | 0.968 | - | 80.24 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_128x128_40k_chase_db1/unet_s5-d16_128x128_40k_chase_db1_20201223_191051-8b16ca0b.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_128x128_40k_chase_db1/unet_s5-d16_128x128_40k_chase_db1-20201223_191051.log.json) |
25-
| UNet-S5-D16 | FCN | HRF | 2336x3504 | 256x256 | 170x170 | 40000 | 2.525 | - | 79.45 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_256x256_40k_hrf/unet_s5-d16_256x256_40k_hrf_20201223_173724-d89cf1ed.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/unet_s5-d16_256x256_40k_hrf/unet_s5-d16_256x256_40k_hrf-20201223_173724.log.json) |
20+
### DRIVE
21+
22+
| Backbone | Head | Image Size | Crop Size | Stride | Lr schd | Mem (GB) | Inf time (fps) | Dice | download |
23+
|--------|----------|----------|-----------|--------:|----------|----------------|------:|--------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
24+
| UNet-S5-D16 | FCN | 584x565 | 64x64 | 42x42 | 40000 | 0.680 | - | 78.67 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_64x64_40k_drive/fcn_unet_s5-d16_64x64_40k_drive_20201223_191051-26cee593.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_64x64_40k_drive/fcn_unet_s5-d16_64x64_40k_drive-20201223_191051.log.json) |
25+
| UNet-S5-D16 | PSPNet | 584x565 | 64x64 | 42x42 | 40000 | 0.599 | - | 78.62 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_64x64_40k_drive/pspnet_unet_s5-d16_64x64_40k_drive_20201227_181818-aac73387.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_64x64_40k_drive/pspnet_unet_s5-d16_64x64_40k_drive-20201227_181818.log.json) |
26+
| UNet-S5-D16 | DeepLabV3 | 584x565 | 64x64 | 42x42 | 40000 | 0.596 | - | 78.69 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_64x64_40k_drive/deeplabv3_unet_s5-d16_64x64_40k_drive_20201226_094047-0671ff20.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_64x64_40k_drive/deeplabv3_unet_s5-d16_64x64_40k_drive-20201226_094047.log.json) |
27+
28+
### STARE
29+
30+
| Backbone | Head | Image Size | Crop Size | Stride | Lr schd | Mem (GB) | Inf time (fps) | Dice | download |
31+
|--------|----------|----------|-----------|--------:|----------|----------------|------:|--------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
32+
| UNet-S5-D16 | FCN | 605x700 | 128x128 | 85x85 | 40000 | 0.968 | - | 81.02 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_128x128_40k_stare/fcn_unet_s5-d16_128x128_40k_stare_20201223_191051-6ea7cfda.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_128x128_40k_stare/fcn_unet_s5-d16_128x128_40k_stare-20201223_191051.log.json) |
33+
| UNet-S5-D16 | PSPNet | 605x700 | 128x128 | 85x85 | 40000 | 0.982 | - | 81.22 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_128x128_40k_stare/pspnet_unet_s5-d16_128x128_40k_stare_20201227_181818-3c2923c4.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_128x128_40k_stare/pspnet_unet_s5-d16_128x128_40k_stare-20201227_181818.log.json) |
34+
| UNet-S5-D16 | DeepLabV3 | 605x700 | 128x128 | 85x85 | 40000 | 0.999 | - | 80.93 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_128x128_40k_stare/deeplabv3_unet_s5-d16_128x128_40k_stare_20201226_094047-93dcb93c.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_128x128_40k_stare/deeplabv3_unet_s5-d16_128x128_40k_stare-20201226_094047.log.json) |
35+
36+
### CHASE_DB1
37+
38+
| Backbone | Head | Image Size | Crop Size | Stride | Lr schd | Mem (GB) | Inf time (fps) | Dice | download |
39+
|--------|----------|----------|-----------|--------:|----------|----------------|------:|--------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
40+
| UNet-S5-D16 | FCN | 960x999 | 128x128 | 85x85 | 40000 | 0.968 | - | 80.24 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_128x128_40k_chase_db1/fcn_unet_s5-d16_128x128_40k_chase_db1_20201223_191051-95852f45.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_128x128_40k_chase_db1/fcn_unet_s5-d16_128x128_40k_chase_db1-20201223_191051.log.json) |
41+
| UNet-S5-D16 | PSPNet | 960x999 | 128x128 | 85x85 | 40000 | 0.982 | - | 80.36 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_128x128_40k_chase_db1/pspnet_unet_s5-d16_128x128_40k_chase_db1_20201227_181818-68d4e609.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_128x128_40k_chase_db1/pspnet_unet_s5-d16_128x128_40k_chase_db1-20201227_181818.log.json) |
42+
| UNet-S5-D16 | DeepLabV3 | 960x999 | 128x128 | 85x85 | 40000 | 0.999 | - | 80.47 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_128x128_40k_chase_db1/deeplabv3_unet_s5-d16_128x128_40k_chase_db1_20201226_094047-4c5aefa3.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_128x128_40k_chase_db1/deeplabv3_unet_s5-d16_128x128_40k_chase_db1-20201226_094047.log.json) |
43+
44+
### HRF
45+
46+
| Backbone | Head | Image Size | Crop Size | Stride | Lr schd | Mem (GB) | Inf time (fps) | Dice | download |
47+
|--------|----------|----------|-----------|--------:|----------|----------------|------:|--------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
48+
| UNet-S5-D16 | FCN | 2336x3504 | 256x256 | 170x170 | 40000 | 2.525 | - | 79.45 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_256x256_40k_hrf/fcn_unet_s5-d16_256x256_40k_hrf_20201223_173724-df3ec8c4.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_256x256_40k_hrf/fcn_unet_s5-d16_256x256_40k_hrf-20201223_173724.log.json) |
49+
| UNet-S5-D16 | PSPNet | 2336x3504 | 256x256 | 170x170 | 40000 | 2.588 | - | 80.07 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_256x256_40k_hrf/pspnet_unet_s5-d16_256x256_40k_hrf_20201227_181818-fdb7e29b.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/pspnet_unet_s5-d16_256x256_40k_hrf/pspnet_unet_s5-d16_256x256_40k_hrf-20201227_181818.log.json) |
50+
| UNet-S5-D16 | DeepLabV3 | 2336x3504 | 256x256 | 170x170 | 40000 | 2.604 | - | 80.21 | [model](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_256x256_40k_hrf/deeplabv3_unet_s5-d16_256x256_40k_hrf_20201226_094047-3a1fdf85.pth) | [log](https://download.openmmlab.com/mmsegmentation/v0.5/unet/deeplabv3_unet_s5-d16_256x256_40k_hrf/deeplabv3_unet_s5-d16_256x256_40k_hrf-20201226_094047.log.json) |

‎configs/unet/deeplabv3_unet_s5-d16_128x128_40k_chase_db1.py

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
_base_ = [
2+
'../_base_/models/deeplabv3_unet_s5-d16.py',
3+
'../_base_/datasets/chase_db1.py', '../_base_/default_runtime.py',
4+
'../_base_/schedules/schedule_40k.py'
5+
]
6+
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))
7+
evaluation = dict(metric='mDice')

‎configs/unet/deeplabv3_unet_s5-d16_128x128_40k_stare.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/deeplabv3_unet_s5-d16.py', '../_base_/datasets/stare.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))
6+
evaluation = dict(metric='mDice')

‎configs/unet/deeplabv3_unet_s5-d16_256x256_40k_hrf.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/deeplabv3_unet_s5-d16.py', '../_base_/datasets/hrf.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(256, 256), stride=(170, 170))
6+
evaluation = dict(metric='mDice')

‎configs/unet/deeplabv3_unet_s5-d16_64x64_40k_drive.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/deeplabv3_unet_s5-d16.py', '../_base_/datasets/drive.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(64, 64), stride=(42, 42))
6+
evaluation = dict(metric='mDice')

‎configs/unet/fcn_unet_s5-d16_128x128_40k_chase_db1.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/fcn_unet_s5-d16.py', '../_base_/datasets/chase_db1.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))
6+
evaluation = dict(metric='mDice')

‎configs/unet/unet_s5-d16_128x128_40k_chase_db1.py ‎configs/unet/fcn_unet_s5-d16_128x128_40k_stare.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
_base_ = [
2-
'../_base_/models/unet_s5-d16.py', '../_base_/datasets/chase_db1.py',
2+
'../_base_/models/fcn_unet_s5-d16.py', '../_base_/datasets/stare.py',
33
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
44
]
55
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))

‎configs/unet/unet_s5-d16_256x256_40k_hrf.py ‎configs/unet/fcn_unet_s5-d16_256x256_40k_hrf.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
_base_ = [
2-
'../_base_/models/unet_s5-d16.py', '../_base_/datasets/hrf.py',
2+
'../_base_/models/fcn_unet_s5-d16.py', '../_base_/datasets/hrf.py',
33
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
44
]
55
test_cfg = dict(crop_size=(256, 256), stride=(170, 170))

‎configs/unet/unet_s5-d16_64x64_40k_drive.py ‎configs/unet/fcn_unet_s5-d16_64x64_40k_drive.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
_base_ = [
2-
'../_base_/models/unet_s5-d16.py', '../_base_/datasets/drive.py',
2+
'../_base_/models/fcn_unet_s5-d16.py', '../_base_/datasets/drive.py',
33
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
44
]
55
test_cfg = dict(crop_size=(64, 64), stride=(42, 42))

‎configs/unet/pspnet_unet_s5-d16_128x128_40k_chase_db1.py

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
_base_ = [
2+
'../_base_/models/pspnet_unet_s5-d16.py',
3+
'../_base_/datasets/chase_db1.py', '../_base_/default_runtime.py',
4+
'../_base_/schedules/schedule_40k.py'
5+
]
6+
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))
7+
evaluation = dict(metric='mDice')

‎configs/unet/unet_s5-d16_128x128_40k_stare.py ‎configs/unet/pspnet_unet_s5-d16_128x128_40k_stare.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
_base_ = [
2-
'../_base_/models/unet_s5-d16.py', '../_base_/datasets/stare.py',
2+
'../_base_/models/pspnet_unet_s5-d16.py', '../_base_/datasets/stare.py',
33
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
44
]
55
test_cfg = dict(crop_size=(128, 128), stride=(85, 85))

‎configs/unet/pspnet_unet_s5-d16_256x256_40k_hrf.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/pspnet_unet_s5-d16.py', '../_base_/datasets/hrf.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(256, 256), stride=(170, 170))
6+
evaluation = dict(metric='mDice')

‎configs/unet/pspnet_unet_s5-d16_64x64_40k_drive.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
_base_ = [
2+
'../_base_/models/pspnet_unet_s5-d16.py', '../_base_/datasets/drive.py',
3+
'../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
4+
]
5+
test_cfg = dict(crop_size=(64, 64), stride=(42, 42))
6+
evaluation = dict(metric='mDice')

0 commit comments

Comments
 (0)