forked from mindspore-lab/mindocr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
det_psenet.py
117 lines (98 loc) · 3.17 KB
/
det_psenet.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
from ._registry import register_model
from .backbones.mindcv_models.utils import load_pretrained
from .base_model import BaseModel
__all__ = ['PSENet', 'psenet_resnet152', 'psenet_resnet50', 'psenet_mobilenetv3']
def _cfg(url='', **kwargs):
return {
'url': url,
'input_size': (3, 640, 640),
**kwargs
}
default_cfgs = {
'psenet_resnet152': _cfg(
url='https://download.mindspore.cn/toolkits/mindocr/psenet/psenet_resnet152_ic15-6058a798.ckpt'),
'psenet_resnet50': _cfg(
url='https://download.mindspore.cn/toolkits/mindocr/psenet/psenet_resnet50_ic15-7e36cab9.ckpt'),
'psenet_mobilenetv3': _cfg(
url='https://download.mindspore.cn/toolkits/mindocr/psenet/psenet_mobilenetv3_ic15-bf2c1907.ckpt'),
}
class PSENet(BaseModel):
def __init__(self, config):
BaseModel.__init__(self, config)
@register_model
def psenet_resnet152(pretrained=False, **kwargs):
pretrained_backbone = not pretrained
model_config = {
"backbone": {
'name': 'det_resnet152',
'pretrained': pretrained_backbone
},
"neck": {
"name": 'PSEFPN',
"out_channels": 128,
},
"head": {
"name": 'PSEHead',
"hidden_size": 256,
"out_channels": 7
}
}
model = PSENet(model_config)
# load pretrained weights
if pretrained:
default_cfg = default_cfgs['psenet_resnet152']
load_pretrained(model, default_cfg)
return model
@register_model
def psenet_resnet50(pretrained=False, **kwargs):
pretrained_backbone = not pretrained
model_config = {
"backbone": {
'name': 'det_resnet50',
'pretrained': pretrained_backbone
},
"neck": {
"name": 'PSEFPN',
"out_channels": 256,
},
"head": {
"name": 'PSEHead',
"hidden_size": 256,
"out_channels": 7
}
}
model = PSENet(model_config)
# load pretrained weights
if pretrained:
default_cfg = default_cfgs['psenet_resnet50']
load_pretrained(model, default_cfg)
return model
@register_model
def psenet_mobilenetv3(pretrained=False, **kwargs):
pretrained_backbone = 'https://download.mindspore.cn/toolkits/mindcv/mobilenet/mobilenetv3' \
'/mobilenet_v3_large_050_no_scale_se_v2_expand-3c4047ac.ckpt'
model_config = {
"backbone": {
'name': 'det_mobilenet_v3',
'architecture': 'large',
'alpha': 0.5,
'out_stages': [5, 8, 14, 20],
'bottleneck_params': {'se_version': 'SqueezeExciteV2', 'always_expand': True},
'pretrained': pretrained_backbone if not pretrained else False
},
"neck": {
"name": 'PSEFPN',
"out_channels": 96,
},
"head": {
"name": 'PSEHead',
"hidden_size": 96,
"out_channels": 7
}
}
model = PSENet(model_config)
# load pretrained weights
if pretrained:
default_cfg = default_cfgs['psenet_mobilenetv3']
load_pretrained(model, default_cfg)
return model