Fix yolox detector prediction (#125)

* fix export error Co-authored-by: zxy <zxy@B-X0GHLVDL-0414.local>
alibaba · Jul 27, 2022 · de84c54 · de84c54
1 parent e3eaa8b
commit de84c54
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/easycv/predictors/detector.py b/easycv/predictors/detector.py
@@ -94,12 +94,11 @@ def __init__(self,
                 'meta'], 'meta.config is missing from checkpoint'
 
             config_str = checkpoint['meta']['config']
-            config_str = config_str[config_str.find('_base_'):]
             # get config
             basename = os.path.basename(self.model_path)
             fname, _ = os.path.splitext(basename)
             self.local_config_file = os.path.join(CACHE_DIR,
-                                                  f'{fname}_config.py')
+                                                  f'{fname}_config.json')
             if not os.path.exists(CACHE_DIR):
                 os.makedirs(CACHE_DIR)
             with open(self.local_config_file, 'w') as ofile:
@@ -109,6 +108,7 @@ def __init__(self,
 
             # build model
             self.model = build_model(self.cfg.model)
+
             self.traceable = getattr(self.model, 'trace_able', False)
 
             self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

diff --git a/tests/ut_config.py b/tests/ut_config.py
@@ -83,7 +83,7 @@
                                        'pretrained_models/faceid')
 
 PRETRAINED_MODEL_YOLOXS_EXPORT = os.path.join(
-    BASE_LOCAL_PATH, 'pretrained_models/detection/yolox_s/epoch_300.pth')
+    BASE_LOCAL_PATH, 'pretrained_models/detection/yolox_s/epoch_300_export.pt')
 PRETRAINED_MODEL_YOLOXS_END2END_JIT = os.path.join(
     BASE_LOCAL_PATH,
     'pretrained_models/detection/yolox_s/epoch_300_end2end.jit')