import yaml
import onnx
from pathlib import Path
from ultralytics import YOLO

CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"


def load_config() -> dict:
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f)["pt_to_onnx"]


def check_onnx_opset(onnx_path):
    model  = onnx.load(str(onnx_path), load_external_data=False)
    opsets = {op.domain: op.version for op in model.opset_import}
    return opsets.get("", opsets.get("ai.onnx", 0))


def main():
    cfg = load_config()

    print(f"Loading PyTorch model from '{cfg['model']}'...")
    yolo_model = YOLO(cfg["model"])

    print("Exporting to ONNX (with simplify=True)...")
    exported_onnx_path = yolo_model.export(
        format="onnx",
        imgsz=cfg["imgsz"],
        dynamic=False,
        simplify=True,
        opset=cfg["opset"],
        batch=1,
        nms=False # export without NMS for better compatibility with KL720's post-processing
    )

    print(f"Exported ONNX opset version: {check_onnx_opset(exported_onnx_path)}")


if __name__ == "__main__":
    main()