|
|
@@ -1,17 +1,18 @@
|
|
|
# ----------------------- Model Config -----------------------
|
|
|
-from .fcos_config import fcos_cfg
|
|
|
-from .yolof_config import yolof_cfg
|
|
|
-from .detr_config import detr_cfg
|
|
|
+from .fcos_config import build_fcos_config
|
|
|
+from .yolof_config import build_yolof_config
|
|
|
+from .detr_config import build_detr_config
|
|
|
|
|
|
def build_config(args):
|
|
|
# FCOS
|
|
|
- if args.model in fcos_cfg.keys():
|
|
|
- return fcos_cfg[args.model]
|
|
|
+ if "fcos" in args.model:
|
|
|
+ return build_fcos_config(args)
|
|
|
# YOLOF
|
|
|
- elif args.model in yolof_cfg.keys():
|
|
|
- return yolof_cfg[args.model]
|
|
|
+ elif "yolof" in args.model:
|
|
|
+ return build_yolof_config
|
|
|
# DETR
|
|
|
- elif args.model in detr_cfg.keys():
|
|
|
- return detr_cfg[args.model]
|
|
|
+ elif "detr" in args.model:
|
|
|
+ return build_detr_config
|
|
|
+
|
|
|
else:
|
|
|
raise NotImplementedError('Unknown Model: {}'.format(args.model))
|