diff --git a/CHANGELOG.md b/CHANGELOG.md index ce0d49a7..e2984d46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ _Add new changes here_ ### Added +- Added Hugging Face model support to ClassifierContainer - Added `piffle` package as dependency ([#575](https://github.com/maps-as-data/MapReader/pull/575)) ## [v1.8.1](https://github.com/Living-with-machines/MapReader/releases/tag/v1.8.1) (2025-08-11) diff --git a/docs/source/using-mapreader/step-by-step-guide/4-classify/train.rst b/docs/source/using-mapreader/step-by-step-guide/4-classify/train.rst index 40f8df03..f76ae0ee 100644 --- a/docs/source/using-mapreader/step-by-step-guide/4-classify/train.rst +++ b/docs/source/using-mapreader/step-by-step-guide/4-classify/train.rst @@ -344,25 +344,24 @@ There are a number of options for the ``model`` argument: If you use this option, your optimizer, scheduler and loss function will be loaded from last time. - **4. To load a** `hugging face model `__\ **, choose your model, follow the "Use in Transformers" or "Use in timm" instructions to load it and then pass this as the ``model`` argument.** + **4. To load a** `hugging face model `__\ **, pass the model's repository ID as a string and set ``huggingface=True``.** - e.g. `This model `__ is based on our `*gold standard* dataset `__. - It can be loaded using the `transformers `__ library: + MapReader will automatically download the model and its corresponding image processor from the Hugging Face Hub using the `transformers `__ library. + e.g. `This model `__ is based on our `*gold standard* dataset `__. + It can be loaded directly like this: + .. code-block:: python #EXAMPLE import torch - from transformers import AutoFeatureExtractor, AutoModelForImageClassification - from mapreader import ClassifierContainer - extractor = AutoFeatureExtractor.from_pretrained("davanstrien/autotrain-mapreader-5000-40830105612") - my_model = AutoModelForImageClassification.from_pretrained("davanstrien/autotrain-mapreader-5000-40830105612") + my_model = "davanstrien/autotrain-mapreader-5000-40830105612" device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu' - my_classifier = ClassifierContainer(my_model, annotated_images.labels_map, dataloaders, device=device) + my_classifier = ClassifierContainer(my_model, annotated_images.labels_map, dataloaders, device=device, huggingface=True) .. note:: You will need to install the `transformers `__ library to do this (``pip install transformers``). diff --git a/mapreader/classify/classifier.py b/mapreader/classify/classifier.py index 46513edf..5be5a6cd 100644 --- a/mapreader/classify/classifier.py +++ b/mapreader/classify/classifier.py @@ -108,6 +108,7 @@ def __init__( is_inception: bool = False, load_path: str | None = None, force_device: bool = False, + huggingface: bool = False, **kwargs, ): # set up device @@ -149,7 +150,31 @@ def __init__( self.input_size = input_size self.is_inception = is_inception elif isinstance(model, str): - self._initialize_model(model, **kwargs) + if huggingface: + try: + from transformers import AutoModelForImageClassification, AutoImageProcessor + except ImportError: + raise ImportError( + "Hugging Face models require the 'transformers' library: 'pip install transformers'." + ) + print(f"[INFO] Initializing Hugging Face model: {model}") + num_labels = len(self.labels_map) + self.model = AutoModelForImageClassification.from_pretrained( + model, + num_labels=num_labels, + ignore_mismatched_sizes=True + ).to(self.device) + hf_processor = AutoImageProcessor.from_pretrained(model) + size = getattr(hf_processor, "size", {}) + if "height" in size and "width" in size: + size = (size["height"], size["width"]) + elif "shortest_edge" in size: + size = (size["shortest_edge"], size["shortest_edge"]) + else: + size = input_size + self.is_inception = False + else: + self._initialize_model(model, **kwargs) self.optimizer = None self.scheduler = None