avoid download on import

imagededup/__init__.py

@@ -1 +1 @@
-__version__ = '0.3.2'
+__version__ = '0.3.2'

imagededup/methods/cnn.py

@@ -44,9 +44,7 @@ class CNN:
     def __init__(
         self,
         verbose: bool = True,
-        model_config: CustomModel = CustomModel(
-            model=MobilenetV3(), transform=MobilenetV3.transform, name=MobilenetV3.name
-        ),
+        model_config: Optional[CustomModel] = None
     ) -> None:
         """
         Initialize a pytorch MobileNet model v3 that is sliced at the last convolutional layer.
@@ -56,7 +54,9 @@ def __init__(
             verbose: Display progress bar if True else disable it. Default value is True.
             model_config: A CustomModel that can be used to initialize a custom PyTorch model along with the corresponding transform.
         """
-        self.model_config = model_config
+        self.model_config = model_config if model_config is not None else CustomModel(
+            model=MobilenetV3(), transform=MobilenetV3.transform, name=MobilenetV3.name
+        )
         self._validate_model_config()
 
         self.logger = return_logger(

tests/test_cnn.py

@@ -46,6 +46,44 @@ def mocker_save_json(mocker):
     return mocker.patch('imagededup.methods.cnn.save_json')
 
 
+def test_import_defaults():
+    """Ensure that MobileNet does not get downloaded on import"""
+    from torch.hub import get_dir
+    from pathlib import Path
+
+    checkpoint_dir = Path(get_dir()) / "checkpoints"
+
+    # Clear cached MobileNet model
+    for model_path in checkpoint_dir.iterdir():
+        if model_path.name.startswith("mobilenet_v3_small"):
+            os.remove(model_path)
+    assert not [
+        m_path
+        for m_path in checkpoint_dir.iterdir()
+        if m_path.name.startswith("mobilenet_v3_small")
+    ]
+
+    # Re-import cnn and assert model is not downloaded
+    import sys
+
+    del sys.modules["imagededup.methods.cnn"]
+    from imagededup.methods.cnn import CNN
+
+    assert not [
+        m_path
+        for m_path in checkpoint_dir.iterdir()
+        if m_path.name.startswith("mobilenet_v3_small")
+    ]
+
+    # Instantiate CNN class and assert model was downloaded
+    _ = CNN()
+    assert [
+        m_path
+        for m_path in checkpoint_dir.iterdir()
+        if m_path.name.startswith("mobilenet_v3_small")
+    ]
+
+
 def test__init_defaults(cnn):
     assert cnn.batch_size == TEST_BATCH_SIZE
     assert cnn.model_config.name == MobilenetV3.name