grubguesser-django / src /model_client.py
jsolow's picture
add-django (#2)
f4da933
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