src/model_client.py

from time import time
from typing import List, Optional, Tuple, Union
import json
import os
import requests
from urllib.parse import urljoin

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

import logging

from huggingface.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 not is_healthy():
        logging.error("Model failed healthcheck")
        top_guesses = ["Model Offline", "Try Again Later", "", "", ""]
    elif 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


def is_healthy() -> bool:
    return True