save

viwoqu · Sep 19, 2019 · e312ab2 · e312ab2
1 parent 3f682c0
commit e312ab2
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 6 deletions.
diff --git a/configs/estimation/hrnet/demo.yaml b/configs/estimation/hrnet/demo.yaml
@@ -17,7 +17,7 @@ processor_cfg:
     bbox_thre: 0.8
   skeleton_cfg:
     model_cfg: configs/estimation/hrnet/pose_hrnet_w32_256x192_test.yaml
-    checkpint_file: checkpoints/pose_hrnet_w32_256x192_official-e6e09ec4.pth
+    checkpoint_file: checkpoints/pose_hrnet_w32_256x192_official.pth
     data_cfg:
       image_size:
         - 192
@@ -35,4 +35,4 @@ processor_cfg:
       save_video: true
       save_dir: "work_dir/estimation_demo"
 
-  gpus: 4
+  gpus: 8
diff --git a/mmskeleton/processor/apis.py b/mmskeleton/processor/apis.py
@@ -1,6 +1,7 @@
 from mmskeleton.utils import call_obj
 import torch
-from mmskeleton.utils import load_checkpoint
+# from mmskeleton.utils import load_checkpoint
+from mmcv.runner import load_checkpoint
 from .utils.infernce_utils import get_final_preds
 from mmcv.utils import Config
 from collections import OrderedDict

diff --git a/mmskeleton/processor/image2skeleton.py b/mmskeleton/processor/image2skeleton.py
@@ -51,7 +51,7 @@ def worker(video_file, index, detection_cfg, skeleton_cfg, skeleon_data_cfg,
     logger.info('Begin to build estimation model')
     beign_time = time()
     skeleton_model_file = skeleton_cfg.model_cfg
-    skeletion_checkpoint_file = skeleton_cfg.checkpint_file
+    skeletion_checkpoint_file = skeleton_cfg.checkpoint_file
     skeleton_model = init_twodimestimator(skeleton_model_file,
                                           skeletion_checkpoint_file,
                                           device='cpu')

diff --git a/mmskeleton/processor/twodimestimation.py b/mmskeleton/processor/twodimestimation.py
@@ -1,8 +1,7 @@
 import torch
 import numpy as np
 from collections import OrderedDict
-from mmcv.runner import Runner, load_checkpoint
-from mmskeleton.utils import call_obj, import_obj
+from mmskeleton.utils import call_obj, import_obj, load_checkpoint
 from mmcv import ProgressBar
 from mmcv.parallel import MMDataParallel
 from mmskeleton.datasets.utils.coco_transform import flip_back

diff --git a/mmskeleton/utils/checkpoint.py b/mmskeleton/utils/checkpoint.py
@@ -15,3 +15,5 @@ def load_checkpoint(model, filename, *args, **kwargs):
         model_url = (mmskeleton_model_urls[model_name])
         checkpoint = mmcv_load_checkpoint(model, model_url, *args, **kwargs)
         return checkpoint
+    else:
+        return mmcv_load_checkpoint(model, filename, *args, **kwargs)