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


MODEL_ENDPOINT_URL = os.getenv("MODEL_ENDPOINT_URL", "https://0.0.0.0:2000")


def try_make_request(request_kwargs, request_type: str):
    try:
        request_kwargs["timeout"] = 3
        if request_type.lower() == "get":
            response = requests.get(**request_kwargs)
        elif request_type.lower() == "post":
            response = requests.post(**request_kwargs)
        else:
            raise Exception("Request Type not Supported.  Only get, post supported.")

        return json.loads(response.content)
    except requests.exceptions.ConnectionError:
        logging.warning("Failed Model prediction", exc_info=True)
        return ["Image Failed to Predict", "Try Another Image", "", "", ""]
    except Exception:
        logging.warning("Failed Model prediction", exc_info=True)
        return ["Image Failed to Predict", "Try Another Image", "", "", ""]


def predict_url(url: str) -> List[str]:
    params = {"url": url}
    headers = {"content-type": "application/json", "Accept-Charset": "UTF-8"}
    request_url = urljoin(MODEL_ENDPOINT_URL, "predict_url")
    request_kwargs = dict(url=request_url, params=params, headers=headers)
    return try_make_request(request_kwargs, "get")


def predict_file(image_file) -> List[str]:
    image_file.seek(0)
    file_ob = {
        "upload_file": (image_file.name, image_file.read(), image_file.content_type)
    }
    request_url = urljoin(MODEL_ENDPOINT_URL, "predict_file")
    request_kwargs = dict(url=request_url, files=file_ob)
    return try_make_request(request_kwargs, "post")


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)):
        top_guesses = predict_file(input_img)
    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:
    request_url = urljoin(MODEL_ENDPOINT_URL, "healthcheck")
    try:
        response = requests.get(url=request_url, timeout=1)
    except Exception:
        logging.error("Failed to make healthcheck request")
        return False
    if response.status_code == 200:
        try:
            response_content = json.loads(response.content)
        except Exception:
            logging.error("Failed to load healthcheck content")
            return False
        if response_content == {"status": "alive"}:
            return True
    return False