Merge pull request #9 from BrainLesion/fix/weights-download-bug

- fix buggy weights download
BrainLesion · Jan 9, 2024 · 8aa8b01 · 8aa8b01
2 parents f7667ca + 1add1e5
commit 8aa8b01
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -129,4 +129,6 @@ dmypy.json
 .pyre/
 
 .vscode
-poetry.lock
+poetry.lock
+
+.DS_Store
diff --git a/brainles_aurora/download.py b/brainles_aurora/download.py
@@ -33,8 +33,12 @@ def download_file(repo: Repository, folder: str, out: str):
 def download_model_weights(target_folder):
     # dl
     g = Github()
-    repo = g.get_repo("neuronflow/BrainLes")
-    dl_folder = "AURORA/brainles_aurora/model_weights"
+    repo = g.get_repo("BrainLesion/AURORA")
+    dl_folder = "brainles_aurora/model_weights"
+    # get parent of parget folder to prevent /AURORA/brainles_aurora/brainles_aurora/model_weights
+    if "brainles_aurora" in target_folder:
+        target_folder = os.path.dirname(target_folder)
+
     download_folder(
         repo=repo,
         folder=dl_folder,