Add HF hub mixin

.flake8

@@ -1,5 +1,5 @@
 [flake8]
 max-line-length = 119
 exclude =.git,__pycache__,docs/conf.py,build,dist,setup.py,tests,.venv
-ignore = I101,I201,F401,F403,S001,D100,D101,D102,D103,D104,D105,D106,D107,D200,D205,D400,W504,D202,E203,W503,B006,D412
+ignore = I101,I201,F401,F403,S001,D100,D101,D102,D103,D104,D105,D106,D107,D200,D205,D400,W504,D202,E203,W503,B006,D412,F821,E501
 inline-quotes = "

.gitignore

@@ -4,6 +4,8 @@ __pycache__/
 *$py.class
 .idea/
 .venv*
+examples/images*
+examples/annotations*
 
 # C extensions
 *.so

segmentation_models_pytorch/__init__.py

@@ -12,6 +12,7 @@
 from .decoders.pspnet import PSPNet
 from .decoders.deeplabv3 import DeepLabV3, DeepLabV3Plus
 from .decoders.pan import PAN
+from .base.hub_mixin import from_pretrained
 
 from .__version__ import __version__
 

segmentation_models_pytorch/__version__.py

@@ -1,3 +1,3 @@
-VERSION = (0, 3, 3)
+VERSION = (0, 3, "4dev0")
 
 __version__ = ".".join(map(str, VERSION))

segmentation_models_pytorch/base/hub_mixin.py

@@ -0,0 +1,144 @@
+import json
+from pathlib import Path
+from typing import Optional, Union
+from functools import wraps
+from huggingface_hub import PyTorchModelHubMixin, ModelCard, ModelCardData, hf_hub_download
+
+
+MODEL_CARD = """
+---
+{{ card_data }}
+---
+# {{ model_name }} Model Card
+
+Table of Contents:
+- [Load trained model](#load-trained-model)
+- [Model init parameters](#model-init-parameters)
+- [Model metrics](#model-metrics)
+- [Dataset](#dataset)
+
+## Load trained model
+```python
+import segmentation_models_pytorch as smp
+
+model = smp.{{ model_name }}.from_pretrained("{{ save_directory | default("<save-directory-or-repo>", true)}}")
+```
+
+## Model init parameters
+```python
+model_init_params = {{ model_parameters }}
+```
+
+## Model metrics
+{{ metrics | default("[More Information Needed]", true) }}
+
+## Dataset
+Dataset name: {{ dataset | default("[More Information Needed]", true) }}
+
+## More Information
+- Library: {{ repo_url | default("[More Information Needed]", true) }}
+- Docs: {{ docs_url | default("[More Information Needed]", true) }}
+
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin)
+"""
+
+
+def _format_parameters(parameters: dict):
+ params = {k: v for k, v in parameters.items() if not k.startswith("_")}
+ params = [f'"{k}": {v}' if not isinstance(v, str) else f'"{k}": "{v}"' for k, v in params.items()]
+ params = ",\n".join([f" {param}" for param in params])
+ params = "{\n" + f"{params}" + "\n}"
+ return params
+
+
+class SMPHubMixin(PyTorchModelHubMixin):
+ def generate_model_card(self, *args, **kwargs) -> ModelCard:
+
+ model_parameters_json = _format_parameters(self._hub_mixin_config)
+ directory = self._save_directory if hasattr(self, "_save_directory") else None
+ repo_id = self._repo_id if hasattr(self, "_repo_id") else None
+ repo_or_directory = repo_id if repo_id is not None else directory
+
+ metrics = self._metrics if hasattr(self, "_metrics") else None
+ dataset = self._dataset if hasattr(self, "_dataset") else None
+
+ if metrics is not None:
+ metrics = json.dumps(metrics, indent=4)
+ metrics = f"```json\n{metrics}\n```"
+
+ model_card_data = ModelCardData(
+ languages=["python"],
+ library_name="segmentation-models-pytorch",
+ license="mit",
+ tags=["semantic-segmentation", "pytorch", "segmentation-models-pytorch"],
+ pipeline_tag="image-segmentation",
+ )
+ model_card = ModelCard.from_template(
+ card_data=model_card_data,
+ template_str=MODEL_CARD,
+ repo_url="https://github.com/qubvel/segmentation_models.pytorch",
+ docs_url="https://smp.readthedocs.io/en/latest/",
+ model_parameters=model_parameters_json,
+ save_directory=repo_or_directory,
+ model_name=self.__class__.__name__,
+ metrics=metrics,
+ dataset=dataset,
+ )
+ return model_card
+
+ def _set_attrs_from_kwargs(self, attrs, kwargs):
+ for attr in attrs:
+ if attr in kwargs:
+ setattr(self, f"_{attr}", kwargs.pop(attr))
+
+ def _del_attrs(self, attrs):
+ for attr in attrs:
+ if hasattr(self, f"_{attr}"):
+ delattr(self, f"_{attr}")
+
+ @wraps(PyTorchModelHubMixin.save_pretrained)
+ def save_pretrained(self, save_directory: Union[str, Path], *args, **kwargs) -> Optional[str]:
+
+ # set additional attributes to be used in generate_model_card
+ self._save_directory = save_directory
+ self._set_attrs_from_kwargs(["metrics", "dataset"], kwargs)
+
+ # set additional attribute to be used in from_pretrained
+ self._hub_mixin_config["_model_class"] = self.__class__.__name__
+
+ try:
+ # call the original save_pretrained
+ result = super().save_pretrained(save_directory, *args, **kwargs)
+ finally:
+ # delete the additional attributes
+ self._del_attrs(["save_directory", "metrics", "dataset"])
+ self._hub_mixin_config.pop("_model_class")
+
+ return result
+
+ @wraps(PyTorchModelHubMixin.push_to_hub)
+ def push_to_hub(self, repo_id: str, *args, **kwargs):
+ self._repo_id = repo_id
+ self._set_attrs_from_kwargs(["metrics", "dataset"], kwargs)
+ result = super().push_to_hub(repo_id, *args, **kwargs)
+ self._del_attrs(["repo_id", "metrics", "dataset"])
+ return result
+
+ @property
+ def config(self):
+ return self._hub_mixin_config
+
+
+@wraps(PyTorchModelHubMixin.from_pretrained)
+def from_pretrained(pretrained_model_name_or_path: str, *args, **kwargs):
+ config_path = hf_hub_download(
+ pretrained_model_name_or_path, filename="config.json", revision=kwargs.get("revision", None)
+ )
+ with open(config_path, "r") as f:
+ config = json.load(f)
+ model_class_name = config.pop("_model_class")
+
+ import segmentation_models_pytorch as smp
+
+ model_class = getattr(smp, model_class_name)
+ return model_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)

segmentation_models_pytorch/base/model.py

@@ -1,8 +1,13 @@
 import torch
+
 from . import initialization as init
+from .hub_mixin import SMPHubMixin
 
 
-class SegmentationModel(torch.nn.Module):
+class SegmentationModel(
+ torch.nn.Module,
+ SMPHubMixin,
+):
  def initialize(self):
  init.initialize_decoder(self.decoder)
  init.initialize_head(self.segmentation_head)

segmentation_models_pytorch/datasets/oxford_pet.py

@@ -87,7 +87,7 @@ def __getitem__(self, *args, **kwargs):
  sample = super().__getitem__(*args, **kwargs)
 
  # resize images
- image = np.array(Image.fromarray(sample["image"]).resize((256, 256), Image.LINEAR))
+ image = np.array(Image.fromarray(sample["image"]).resize((256, 256), Image.BILINEAR))
  mask = np.array(Image.fromarray(sample["mask"]).resize((256, 256), Image.NEAREST))
  trimap = np.array(Image.fromarray(sample["trimap"]).resize((256, 256), Image.NEAREST))