import os

import cv2
import numpy as np
import albumentations as A
import onnxruntime as ort


class Detection:
    def __init__(self):
        self.CLASS_NAMES = ["background", "Hollowing", "Water seepage", "Cracking"]
        self.PALETTE = np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]], dtype=np.uint8)
        self.tfm = A.Compose([
            A.Resize(512, 512),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        real_path = os.path.join(os.path.dirname(__file__), "segformer.onnx")
        self.model = ort.InferenceSession(
            real_path,
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )
        print("ONNX 模型加载完成！")

    def get_contours(self, mask_np):
        res = []
        for idx in range(1, len(self.CLASS_NAMES)):
            name = self.CLASS_NAMES[idx]
            binary = (mask_np == idx).astype(np.uint8) * 255
            cnts, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for c in cnts:
                if cv2.contourArea(c) < 20:
                    continue
                eps = 0.002 * cv2.arcLength(c, True)
                poly = cv2.approxPolyDP(c, eps, True)
                if len(poly) >= 3:
                    res.append((name, poly.reshape(-1, 2).tolist()))
        return res

    def detect(self, img_input):
        # 读取图片
        if isinstance(img_input, str):
            img_bgr = cv2.imdecode(np.fromfile(img_input, dtype=np.uint8), cv2.IMREAD_COLOR)
        else:
            img_bgr = img_input

        h, w = img_bgr.shape[:2]
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        # 数据预处理
        x = self.tfm(image=img_rgb)["image"]
        # 将 HWC -> CHW 并增加 batch 维度
        x = np.transpose(x, (2, 0, 1))  # HWC -> CHW
        x = np.expand_dims(x, axis=0).astype(np.float32)  # batch x C x H x W

        # ONNX 推理
        inp_name = self.model.get_inputs()[0].name
        out_name = self.model.get_outputs()[0].name
        out = self.model.run([out_name], {inp_name: x})[0]

        # 获取预测结果
        pred = out.argmax(axis=1)[0].astype(np.uint8)

        mask = cv2.resize(pred, (w, h), interpolation=cv2.INTER_NEAREST)
        mask_rgb = self.PALETTE[mask]
        coords = self.get_contours(mask)
        # print(coords)
        return mask_rgb, coords


class DetectionMock:
    def __init__(self):
        self.onnx_path = "segformer.onnx"

    def detect(self, img_input):
        if isinstance(img_input, str):
            img_bgr = cv2.imdecode(np.fromfile(img_input, dtype=np.uint8), cv2.IMREAD_COLOR)
        else:
            img_bgr = img_input

        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        return img_rgb, [
            ('Cracking', [[771, 928], [757, 935]]),
            ('Cracking', [[254, 740], [251, 942]]),
            ('Cracking', [[764, 420], [764, 424]]),
            ('Cracking', [[257, 238], [251, 245]]),
            ('Cracking', [[1436, 145], [1401, 145]])
        ]


if __name__ == "__main__":
    img_path = "test.jpg"

    detection = Detection()
    mask_res, coords_res = detection.detect(img_path)

    print("Mask Shape:", mask_res.shape)
    print("Detections:", coords_res)
    cv2.imwrite("res_onnx.png", cv2.cvtColor(mask_res, cv2.COLOR_RGB2BGR))