from deepsparse import Pipeline
import time
import cv2
import numpy as np
import urllib.request

model_path = "model.onnx"
yolo_pipeline = Pipeline.create(
    task="yolov8",
    model_path=model_path,
)

def resize_with_pad(image,
                    new_shape,
                    padding_color = (128, 128, 128)):
    """
    https://gist.github.com/IdeaKing/11cf5e146d23c5bb219ba3508cca89ec
    """
    original_shape = (image.shape[1], image.shape[0])
    ratio = float(max(new_shape))/max(original_shape)
    new_size = tuple([int(x*ratio) for x in original_shape])
    image = cv2.resize(image, new_size)
    delta_w = new_shape[0] - new_size[0]
    delta_h = new_shape[1] - new_size[1]
    top, bottom = delta_h//2, delta_h-(delta_h//2)
    left, right = delta_w//2, delta_w-(delta_w//2)
    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=padding_color)
    # save padded image
    return image

def get_image_from_url(image_url):
    try:
        req = urllib.request.Request(image_url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req) as url:
            image_contents = url.read()
    except urllib.error.HTTPError as e:
        return {"error": str(e)}
    except ConnectionResetError:
        return {"error": "ConnectionResetError"}
    arr = np.asarray(bytearray(image_contents), dtype=np.uint8)
    img = cv2.imdecode(arr, -1)
    return {"image": img}


image_url = "https://cdn.discordapp.com/attachments/1083610420096532581/1125730189733810239/padded.png"
img = get_image_from_url(image_url)
img = img["image"]
img = resize_with_pad(img, (800, 800))
input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
pipeline_outputs = yolo_pipeline(images=[input_img], conf_thres=0.35, iou_thres=0.6)
print(pipeline_outputs)