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