Fix minimal

Author: qubvel

segmentation_models_pytorch/base/hub_mixin.py

@@ -147,4 +147,3 @@ def wrapper(self, *args, **kwargs):
         return func(self, *args, **kwargs)
 
     return wrapper
-from segmentation_models_pytorch.base.hub_mixin import supports_config_loading

tests/config.py

@@ -1,4 +1,9 @@
 import os
+import timm
+from packaging.version import Version
+
+
+has_timm_test_models = Version(timm.__version__) >= Version("1.0.12")
 
 
 def get_commit_message():

tests/encoders/test_timm_universal.py

@@ -1,6 +1,5 @@
-import timm
 from tests.encoders import base
-from packaging.version import Version
+from tests.config import has_timm_test_models
 
 # check if timm >= 1.0.12
 timm_encoders = [
@@ -9,7 +8,7 @@
     "tu-darknet17",  # for timm universal vgg-like encoder
 ]
 
-if Version(timm.__version__) >= Version("1.0.12"):
+if has_timm_test_models:
     timm_encoders.append("tu-test_resnet.r160_in1k")
 
 

tests/models/base.py

@@ -1,3 +1,4 @@
+import os
 import inspect
 import tempfile
 import unittest
@@ -6,9 +7,13 @@
 import torch
 import segmentation_models_pytorch as smp
 
+from tests.config import has_timm_test_models
+
 
 class BaseModelTester(unittest.TestCase):
-    test_encoder_name = "tu-test_resnet.r160_in1k"
+    test_encoder_name = (
+        "tu-test_resnet.r160_in1k" if has_timm_test_models else "resnet18"
+    )
 
     # should be overriden
     test_model_type = None
@@ -136,8 +141,12 @@ def test_save_load_with_hub_mixin(self):
 
         # save model
         with tempfile.TemporaryDirectory() as tmpdir:
-            model.save_pretrained(tmpdir)
+            model.save_pretrained(
+                tmpdir, dataset="test_dataset", metrics={"my_awesome_metric": 0.99}
+            )
             restored_model = smp.from_pretrained(tmpdir)
+            with open(os.path.join(tmpdir, "README.md"), "r") as f:
+                readme = f.read()
 
         # check inference is correct
         sample = self._get_sample(
@@ -153,3 +162,7 @@ def test_save_load_with_hub_mixin(self):
 
         self.assertEqual(output.shape, restored_output.shape)
         self.assertEqual(output.shape[1], 1)
+
+        # check dataset and metrics are saved in readme
+        self.assertIn("test_dataset", readme)
+        self.assertIn("my_awesome_metric", readme)