Move everything to HF hub

Author: qubvel

segmentation_models_pytorch/encoders/__init__.py

@@ -1,8 +1,10 @@
+import json
 import timm
 import copy
 import warnings
 import functools
-import torch.utils.model_zoo as model_zoo
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 from .resnet import resnet_encoders
 from .dpn import dpn_encoders
@@ -101,15 +103,26 @@ def get_encoder(name, in_channels=3, depth=5, weights=None, output_stride=32, **
     encoder = EncoderClass(**params)
 
     if weights is not None:
-        try:
-            settings = encoders[name]["pretrained_settings"][weights]
-        except KeyError:
+        if weights not in encoders[name]["pretrained_settings"]:
+            available_weights = list(encoders[name]["pretrained_settings"].keys())
             raise KeyError(
-                "Wrong pretrained weights `{}` for encoder `{}`. Available options are: {}".format(
-                    weights, name, list(encoders[name]["pretrained_settings"].keys())
-                )
+                f"Wrong pretrained weights `{weights}` for encoder `{name}`. "
+                f"Available options are: {available_weights}"
             )
-        encoder.load_state_dict(model_zoo.load_url(settings["url"]))
+
+        settings = encoders[name]["pretrained_settings"][weights]
+        repo_id = settings["repo_id"]
+        revision = settings["revision"]
+
+        # Load config and model
+        hf_hub_download(repo_id, filename="config.json", revision=revision)
+        model_path = hf_hub_download(
+            repo_id, filename="model.safetensors", revision=revision
+        )
+
+        # Load model weights
+        state_dict = load_file(model_path, device="cpu")
+        encoder.load_state_dict(state_dict)
 
     encoder.set_in_channels(in_channels, pretrained=weights is not None)
     if output_stride != 32:
@@ -136,7 +149,16 @@ def get_preprocessing_params(encoder_name, pretrained="imagenet"):
             raise ValueError(
                 "Available pretrained options {}".format(all_settings.keys())
             )
-        settings = all_settings[pretrained]
+
+        repo_id = all_settings[pretrained]["repo_id"]
+        revision = all_settings[pretrained]["revision"]
+
+        # Load config and model
+        config_path = hf_hub_download(
+            repo_id, filename="config.json", revision=revision
+        )
+        with open(config_path, "r") as f:
+            settings = json.load(f)
 
     formatted_settings = {}
     formatted_settings["input_space"] = settings.get("input_space", "RGB")

segmentation_models_pytorch/encoders/densenet.py

@@ -110,92 +110,65 @@ def load_state_dict(self, state_dict):
         super().load_state_dict(state_dict)
 
 
-pretrained_settings = {
-    "densenet121": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/densenet121-fbdb23505.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [0.485, 0.456, 0.406],
-            "std": [0.229, 0.224, 0.225],
-            "num_classes": 1000,
-        }
-    },
-    "densenet169": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/densenet169-f470b90a4.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [0.485, 0.456, 0.406],
-            "std": [0.229, 0.224, 0.225],
-            "num_classes": 1000,
-        }
-    },
-    "densenet201": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/densenet201-5750cbb1e.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [0.485, 0.456, 0.406],
-            "std": [0.229, 0.224, 0.225],
-            "num_classes": 1000,
-        }
-    },
-    "densenet161": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/densenet161-347e6b360.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [0.485, 0.456, 0.406],
-            "std": [0.229, 0.224, 0.225],
-            "num_classes": 1000,
-        }
-    },
-}
-
 densenet_encoders = {
     "densenet121": {
         "encoder": DenseNetEncoder,
-        "pretrained_settings": pretrained_settings["densenet121"],
         "params": {
             "out_channels": [3, 64, 256, 512, 1024, 1024],
             "num_init_features": 64,
             "growth_rate": 32,
             "block_config": (6, 12, 24, 16),
         },
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/densenet121-imagenet",
+                "revision": "main",
+            }
+        },
     },
     "densenet169": {
         "encoder": DenseNetEncoder,
-        "pretrained_settings": pretrained_settings["densenet169"],
         "params": {
             "out_channels": [3, 64, 256, 512, 1280, 1664],
             "num_init_features": 64,
             "growth_rate": 32,
             "block_config": (6, 12, 32, 32),
         },
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/densenet169-imagenet",
+                "revision": "main",
+            }
+        },
     },
     "densenet201": {
         "encoder": DenseNetEncoder,
-        "pretrained_settings": pretrained_settings["densenet201"],
         "params": {
             "out_channels": [3, 64, 256, 512, 1792, 1920],
             "num_init_features": 64,
             "growth_rate": 32,
             "block_config": (6, 12, 48, 32),
         },
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/densenet201-imagenet",
+                "revision": "main",
+            }
+        },
     },
     "densenet161": {
         "encoder": DenseNetEncoder,
-        "pretrained_settings": pretrained_settings["densenet161"],
         "params": {
             "out_channels": [3, 96, 384, 768, 2112, 2208],
             "num_init_features": 96,
             "growth_rate": 48,
             "block_config": (6, 12, 36, 24),
         },
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/densenet161-imagenet",
+                "revision": "main",
+            }
+        },
     },
 }

segmentation_models_pytorch/encoders/dpn.py

@@ -101,79 +101,15 @@ def load_state_dict(self, state_dict, **kwargs):
         super().load_state_dict(state_dict, **kwargs)
 
 
-pretrained_settings = {
-    "dpn68": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn68-4af7d88d2.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-    "dpn68b": {
-        "imagenet+5k": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn68b_extra-363ab9c19.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-    "dpn92": {
-        "imagenet+5k": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn92_extra-fda993c95.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-    "dpn98": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn98-722954780.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-    "dpn131": {
-        "imagenet": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn131-7af84be88.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-    "dpn107": {
-        "imagenet+5k": {
-            "url": "http://data.lip6.fr/cadene/pretrainedmodels/dpn107_extra-b7f9f4cc9.pth",
-            "input_space": "RGB",
-            "input_size": [3, 224, 224],
-            "input_range": [0, 1],
-            "mean": [124 / 255, 117 / 255, 104 / 255],
-            "std": [1 / (0.0167 * 255)] * 3,
-            "num_classes": 1000,
-        }
-    },
-}
-
 dpn_encoders = {
     "dpn68": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn68"],
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/dpn68-imagenet",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [4, 8, 20, 24],
             "out_channels": [3, 10, 144, 320, 704, 832],
@@ -189,7 +125,12 @@ def load_state_dict(self, state_dict, **kwargs):
     },
     "dpn68b": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn68b"],
+        "pretrained_settings": {
+            "imagenet+5k": {
+                "repo_id": "smp-hub/dpn68b-imagenet-5k",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [4, 8, 20, 24],
             "out_channels": [3, 10, 144, 320, 704, 832],
@@ -206,7 +147,12 @@ def load_state_dict(self, state_dict, **kwargs):
     },
     "dpn92": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn92"],
+        "pretrained_settings": {
+            "imagenet+5k": {
+                "repo_id": "smp-hub/dpn92-imagenet-5k",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [4, 8, 28, 32],
             "out_channels": [3, 64, 336, 704, 1552, 2688],
@@ -221,7 +167,12 @@ def load_state_dict(self, state_dict, **kwargs):
     },
     "dpn98": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn98"],
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/dpn98-imagenet",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [4, 10, 30, 34],
             "out_channels": [3, 96, 336, 768, 1728, 2688],
@@ -236,7 +187,12 @@ def load_state_dict(self, state_dict, **kwargs):
     },
     "dpn107": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn107"],
+        "pretrained_settings": {
+            "imagenet+5k": {
+                "repo_id": "smp-hub/dpn107-imagenet-5k",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [5, 13, 33, 37],
             "out_channels": [3, 128, 376, 1152, 2432, 2688],
@@ -251,7 +207,12 @@ def load_state_dict(self, state_dict, **kwargs):
     },
     "dpn131": {
         "encoder": DPNEncoder,
-        "pretrained_settings": pretrained_settings["dpn131"],
+        "pretrained_settings": {
+            "imagenet": {
+                "repo_id": "smp-hub/dpn131-imagenet",
+                "revision": "main",
+            }
+        },
         "params": {
             "stage_idxs": [5, 13, 41, 45],
             "out_channels": [3, 128, 352, 832, 1984, 2688],