imgutils.generic.classify

Generic tools for classification models.

This module provides utilities and classes for working with classification models, particularly those stored in Hugging Face repositories. It includes functions for image encoding, model loading, and prediction, as well as a main ClassifyModel class that manages the interaction with classification models.

Key components:

Image encoding and preprocessing
ClassifyModel: A class for managing and using classification models
Utility functions for making predictions with classification models

The module is designed to work with ONNX models and supports various image input formats. It also handles token-based authentication for accessing private Hugging Face repositories.

ClassifyModel

class imgutils.generic.classify.ClassifyModel(repo_id: str, hf_token: str | None = None)[source]

A class for managing and using classification models.

This class provides methods for loading classification models from a Hugging Face repository, making predictions, and managing model resources. It supports multiple models within a single repository and handles token-based authentication.

Parameters:

repo_id (str) – The Hugging Face repository ID containing the classification models.
hf_token (Optional[str], optional) – The Hugging Face API token for accessing private repositories, defaults to None.

Variables:

repo_id – The Hugging Face repository ID.
_model_names – Cached list of available model names in the repository.
_models – Dictionary of loaded ONNX models.
_labels – Dictionary of labels for each model.
_hf_token – The Hugging Face API token.

Usage:

>>> model = ClassifyModel("username/repo_name")
>>> image = Image.open("path/to/image.jpg")
>>> prediction, score = model.predict(image, "model_name")
>>> print(f"Predicted class: {prediction}, Score: {score}")

__init__(repo_id: str, hf_token: str | None = None)[source]

Initialize the ClassifyModel instance.

Parameters:

repo_id (str) – The repository ID containing the models.
hf_token (Optional[str], optional) – The Hugging Face API token, defaults to None.

clear()[source]

Clear the cached models and labels.

This method frees up memory by removing all loaded models and labels from the cache.

launch_demo(default_model_name: str | None = None, server_name: str | None = None, server_port: int | None = None, **kwargs)[source]

Launch the Gradio demo for the classifier model.

This method creates a Gradio Blocks interface, sets up the UI components using make_ui(), and launches the demo server.

Parameters:

default_model_name (Optional[str]) – The name of the default model to be selected in the dropdown.
server_name (Optional[str]) – The name of the server to run the demo on. Defaults to None.
server_port (Optional[int]) – The port number to run the demo on. Defaults to None.
kwargs – Additional keyword arguments to pass to the Gradio launch method.

Raises:

ImportError – If Gradio is not installed or properly configured.

Example:

>>> model = ClassifyModel("username/repo_name")
>>> model.launch_demo(default_model_name="model_v1", server_name="0.0.0.0", server_port=7860)

make_ui(default_model_name: str | None = None)[source]

Create the user interface components for the classifier model demo.

This method sets up the Gradio UI components including an image input, model selection dropdown, submit button, and output label. It also configures the interaction between these components.

Parameters:: default_model_name (Optional[str]) – The name of the default model to be selected in the dropdown. If None, the most recently updated model will be selected.
Raises:: ImportError – If Gradio is not installed or properly configured.
Example:

>>> model = ClassifyModel("username/repo_name")
>>> model.make_ui(default_model_name="model_v1")

Predict the class with the highest score for the given image.

This method runs the image through the model and returns the predicted class and its score.

Parameters:

image (ImageTyping) – The input image to classify.
model_name (str) – The name of the model to use for prediction.

Returns:

A tuple containing the predicted class label and its score.

Return type:

Tuple[str, float]

Raises:

ValueError – If the model name is invalid.
RuntimeError – If there’s an error during prediction.

Predict the scores for each class using the specified model.

This method runs the image through the model and returns a dictionary of class scores.

Parameters:

image (ImageTyping) – The input image to classify.
model_name (str) – The name of the model to use for prediction.

Returns:

A dictionary mapping class labels to their predicted scores.

Return type:

Dict[str, float]

Raises:

ValueError – If the model name is invalid.
RuntimeError – If there’s an error during prediction.

classify_predict_score

Predict the scores for each class using the specified model and repository.

This function is a convenience wrapper around ClassifyModel’s predict_score method.

Parameters:

image (ImageTyping) – The input image to classify.
repo_id (str) – The repository ID containing the models.
model_name (str) – The name of the model to use for prediction.
hf_token (Optional[str]) – Optional Hugging Face authentication token.

Returns:

A dictionary mapping class labels to their predicted scores.

Return type:

Dict[str, float]

Raises:

ValueError – If the model name or repository ID is invalid.
RuntimeError – If there’s an error during prediction.

classify_predict

Predict the class with the highest score using the specified model and repository.

This function is a convenience wrapper around ClassifyModel’s predict method.

Parameters:

image (ImageTyping) – The input image to classify.
repo_id (str) – The repository ID containing the models.
model_name (str) – The name of the model to use for prediction.
hf_token (Optional[str]) – Optional Hugging Face authentication token.

Returns:

A tuple containing the predicted class label and its score.

Return type:

Tuple[str, float]

Raises:

ValueError – If the model name or repository ID is invalid.
RuntimeError – If there’s an error during prediction.