14 lines
325 B
Python
14 lines
325 B
Python
|
# encoding: utf-8
|
||
|
|
"""
|
||
|
|
@author: sherlock
|
||
|
|
@contact: sherlockliao01@gmail.com
|
||
|
|
"""
|
||
|
|
|
||
|
|
from .baseline import Baseline
|
||
|
|
|
||
|
|
|
||
|
|
def build_model(cfg, num_classes):
|
||
|
|
if cfg.MODEL.NAME == 'resnet50':
|
||
|
|
model = Baseline(num_classes, cfg.MODEL.LAST_STRIDE, cfg.MODEL.PRETRAIN_PATH, cfg.MODEL.NECK, cfg.TEST.NECK_FEAT)
|
||
|
|
return model
|