"docs/modelfile.mdx" did not exist on "55c3efa9006989f3f680e9936857900312af3fa4"
__init__.py 531 Bytes
Newer Older
dengjb's avatar
update  
dengjb committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from __future__ import absolute_import

from .ResNet import *
from .DenseNet import *
from .ShuffleNet import *
from .InceptionV4 import *

__factory = {
    'resnet50': ResNet50,
    'resnet101': ResNet101,
    'densenet121': DenseNet121,
    'shufflenet': ShuffleNet,
    'inceptionv4': InceptionV4ReID,
}

def get_names():
    return __factory.keys()

def init_model(name, *args, **kwargs):
    if name not in __factory.keys():
        raise KeyError("Unknown model: {}".format(name))
    return __factory[name](*args, **kwargs)