import onnxruntime
import cv2
import numpy as np
import oloren as olo


class ONNX_Detectron:
    REQUIRED_WIDTH = 800
    REQUIRED_HEIGHT = 1043

    def __init__(self, onnx_path):
        self.onnx_path = onnx_path
        self.model = onnxruntime.InferenceSession(self.onnx_path)

    def preprocess(self, path):
        img = np.array(cv2.imread(path))
        initial_h, initial_w, _ = img.shape
        img = cv2.resize(img, (self.REQUIRED_WIDTH, self.REQUIRED_HEIGHT), interpolation=cv2.INTER_LINEAR)
        img = img.transpose((2, 0, 1)).astype(np.float32)
        return {self.model.get_inputs()[0].name: img}, initial_w, initial_h

    def predict(self, image):
        prepared_input, input_w, input_h = self.preprocess(image)
        bboxes, labels, confidence_scores, _ = self.model.run(None, prepared_input)
        regions = self.postprocess(bboxes, labels, confidence_scores, input_w, input_h)
        return regions

    def postprocess(
        self,
        bboxes: np.ndarray,
        labels: np.ndarray,
        confidence_scores: np.ndarray,
        input_w: float,
        input_h: float,
    ):
        """Process output into Unstructured class. Bounding box coordinates are converted to
        original image resolution."""
        regions = []
        width_conversion = input_w / self.REQUIRED_WIDTH
        height_conversion = input_h / self.REQUIRED_HEIGHT
        for (x1, y1, x2, y2), label, conf in zip(bboxes, labels, confidence_scores):
            if conf < 0.8:
                continue
            region = {
                "x1": x1 * width_conversion,
                "y1": y1 * height_conversion,
                "x2": x2 * width_conversion,
                "y2": y2 * height_conversion,
                "confidence": conf,
                "label": label,
            }

            regions.append(region)

        return regions


@olo.register()
def predict(image=olo.File()):
    model = ONNX_Detectron("./model.onnx")
    return model.predict(image)


if __name__ == "__main__":
    olo.run("detectron")