import json
import numpy as np
import cv2
import onnxruntime
import urllib.request


model = "/child-adult-nonhuman.onnx"
class_names = ["adults", "children", "nonhuman"]


def predict_single_image(image_url):
    req = urllib.request.Request(image_url, headers={"User-Agent": "Mozilla/5.0"})
    with urllib.request.urlopen(req) as url:
        image_contents = url.read()

    arr = np.asarray(bytearray(image_contents), dtype=np.uint8)
    img = cv2.imdecode(arr, -1)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    resized_image = cv2.resize(img, (300, 300))
    image_for_prediction = resized_image.transpose(2, 0, 1)
    image_for_prediction = image_for_prediction.astype("float32") / 255.0
    image_for_prediction = image_for_prediction[np.newaxis, :]

    data = json.dumps({"data": image_for_prediction.tolist()})
    data = np.array(json.loads(data)["data"]).astype("float32")
    session = onnxruntime.InferenceSession(model, None)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    result = session.run([output_name], {input_name: data})
    prediction = int(np.argmax(np.array(result).squeeze(), axis=0))

    probabilities = np.exp(result) / np.sum(np.exp(result))
    score = np.max(probabilities)
    score = float(score)

    return {"predicted_class": class_names[prediction], "score": score}