src/model_client.py

from typing import List, Optional, Tuple, Union

from django.core.files.uploadedfile import InMemoryUploadedFile, TemporaryUploadedFile

import logging

from huggingface_client.predict import predict_file, predict_url



def get_color_labels(guesses: List[str], actual_label: Optional[str]) -> List[str]:
    if not actual_label:
        return ["white"] * len(guesses)
    return ["lime" if x == actual_label else "white" for x in guesses]


def url_image_vars(
    input_img: Union[str, InMemoryUploadedFile, TemporaryUploadedFile], label: str
) -> Tuple[List[str], List[str]]:
    actual_label = label.title()
    if isinstance(input_img, str):
        top_guesses = predict_url(input_img)
    elif isinstance(input_img, (InMemoryUploadedFile, TemporaryUploadedFile)):
        input_img.seek(0)
        top_guesses = predict_file(input_img.read())
    else:
        logging.error(f"Unknown input type: {type(input_img)=}")
        top_guesses = ["Unknown Input Type", "", "", "", ""]
    color_labels = get_color_labels(top_guesses, actual_label)
    return top_guesses, color_labels