utils.py

from transformers import pipeline, Pipeline
from functools import lru_cache
from typing import Optional, Dict, Any
import numpy as np


@lru_cache
def model_initialization(task: str = "image-classification", model_name: str = "microsoft/resnet-18") -> Pipeline:
    """
    Initialize the Hugging Face pipeline for a specified task and model.

    Args:
        task (str): The task type, e.g., "image-classification".
        model_name (str): The name or path of the model to use.

    Returns:
        Pipeline: A Hugging Face pipeline object ready for inference.
    """
    pipe = pipeline(task, model=model_name)

    return pipe


def prediction(pipe: Pipeline, img: np.ndarray) -> Optional[Dict[str, Any]]:
    """
    Perform image classification on the given image using the specified pipeline.

    Args:
        pipe (Pipeline): The initialized hf pipeline object.
        img (np.ndarray): The image to classify.

    Returns:
        Optional[Dict[str, Any]]: A dictionary containing the most promising label and its confidence score,
        or None if no results are returned.
    """
    results = pipe(img)
    results.sort(key=lambda x: x["score"], reverse=True)

    if not results:
        return None
    
    response = {
        "most_promising_label": results[0]["label"],
        "confidence": round(results[0]["score"], 2)
    }

    return response