yjh0410 hace 1 año
padre
commit
6b6ea743b0

+ 10 - 9
odlab/config/__init__.py

@@ -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))

+ 13 - 0
odlab/config/detr_config.py

@@ -1,5 +1,11 @@
 # DETR
 
+def build_detr_config(args):
+    if   args.model == 'detr_r50':
+        return Detr_R50_Config()
+    else:
+        raise NotImplementedError("No config for model: {}".format(args.model))
+
 
 class DetrBaseConfig(object):
     def __init__(self):
@@ -10,6 +16,13 @@ class DetrBaseConfig(object):
         for k, v in config_dict.items():
             print("{} : {}".format(k, v))
 
+class Detr_R50_Config(DetrBaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        ## Backbone
+        pass
+
+
 detr_cfg = {
     'detr_r50':{
         # ---------------- Model config ----------------

+ 12 - 0
odlab/config/fcos_config.py

@@ -1,5 +1,10 @@
 # Fully Convolutional One-Stage object detector
 
+def build_fcos_config(args):
+    if   args.model == 'fcos_r18_1x':
+        return Fcos_R18_1x_Config()
+    else:
+        raise NotImplementedError("No config for model: {}".format(args.model))
 
 class FcosBaseConfig(object):
     def __init__(self):
@@ -10,6 +15,13 @@ class FcosBaseConfig(object):
         for k, v in config_dict.items():
             print("{} : {}".format(k, v))
 
+class Fcos_R18_1x_Config(FcosBaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        ## Backbone
+        pass
+
+
 fcos_cfg = {
     'fcos_r18_1x':{
         # ----------------- Model-----------------

+ 11 - 0
odlab/config/yolof_config.py

@@ -1,5 +1,10 @@
 # Fully Convolutional One-Stage object detector
 
+def build_yolof_config(args):
+    if   args.model == 'yolof_r18_c5_1x':
+        return Yolof_R18_C5_1x_Config()
+    else:
+        raise NotImplementedError("No config for model: {}".format(args.model))
 
 class YolofBaseConfig(object):
     def __init__(self):
@@ -10,6 +15,12 @@ class YolofBaseConfig(object):
         for k, v in config_dict.items():
             print("{} : {}".format(k, v))
 
+class Yolof_R18_C5_1x_Config(YolofBaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        ## Backbone
+        pass
+
 yolof_cfg = {
     # --------------- C5 level ---------------
     'yolof_r18_c5_1x':{