setup_model: raising err on failed download; cleanup

juglab · Jun 3, 2024 · 65deee0 · 65deee0
1 parent 2ea98c9
commit 65deee0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/src/featureforest/SAM/setup_model.py b/src/featureforest/SAM/setup_model.py
@@ -12,10 +12,14 @@ def setup_mobile_sam_model():
     # sam model (light hq sam)
     model = MobileSAM.setup_model().to(device)
     # download model's weights
+    model_url = "https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt"
     model_file = download_model(
-        model_url="https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt",
+        model_url=model_url,
         model_name="mobile_sam.pt"
     )
+    if model_file is None:
+        raise ValueError(f"Could not download the model from {model_url}.")
+
     # load weights
     weights = torch.load(model_file, map_location=device)
     model.load_state_dict(weights, strict=True)

diff --git a/src/featureforest/utils/downloader.py b/src/featureforest/utils/downloader.py
@@ -43,11 +43,3 @@ def download_model(
     except Exception as err:
         print(f"\nError while downloading the model:\n{err}")
         return None
-
-
-if __name__ == "__main__":
-    model_file = download_model(
-        model_url="https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt",
-        model_name="mobile_sam.pt"
-    )
-    print(model_file)