Forráskód Böngészése

Merge pull request #77 from SunAhong1993/syf0519

fix the vis.py
Jason 5 éve
szülő
commit
c20cdfc169
1 módosított fájl, 8 hozzáadás és 1 törlés
  1. 8 1
      paddlex/interpret/visualize.py

+ 8 - 1
paddlex/interpret/visualize.py

@@ -23,6 +23,13 @@ from .core.interpretation import Interpretation
 from .core.normlime_base import precompute_normlime_weights
 
 
+def gen_user_home():
+    if "HOME" in os.environ:
+        home_path = os.environ["HOME"]
+        if os.path.exists(home_path) and os.path.isdir(home_path):
+            return home_path
+    return os.path.expanduser('~')
+
 def visualize(img_file, 
               model, 
               dataset=None,
@@ -109,7 +116,7 @@ def get_normlime_interpreter(img, model, dataset, num_samples=3000, batch_size=5
     labels_name = None
     if dataset is not None:
         labels_name = dataset.labels
-    root_path = os.environ['HOME']
+    root_path = gen_user_home()
     root_path = osp.join(root_path, '.paddlex')
     pre_models_path = osp.join(root_path, "pre_models")
     if not osp.exists(pre_models_path):