image_classification.py

from huggingface_hub import InferenceClient
from os import path, unlink, getenv
from PIL.Image import Image
import pandas as pd
from pandas import DataFrame
from utils import save_image_to_temp_file


def image_classification(client: InferenceClient, image: Image) -> DataFrame:
    try:
        temp_file_path = save_image_to_temp_file(image) # Needed because InferenceClient does not accept PIL Images directly.
        classifications = client.image_classification(temp_file_path, model=getenv("IMAGE_CLASSIFICATION_MODEL"))
        return pd.DataFrame({
                                "Label": classification.label,
                                "Probability": f"{classification.score:.2%}"
                            }
                            for classification
                            in classifications)
    finally:
        if temp_file_path and path.exists(temp_file_path): # Clean up temporary file.
            try:
                unlink(temp_file_path)
            except Exception:
                pass # Ignore clean-up errors.