"""Exports RF-DETR Large (2026) and Seg Large to ONNX from Roboflow's own
checkpoints, each at the size its variant is trained for.

    pip install rfdetr onnx onnxruntime onnxsim
    python scripts/export.py <output directory>

The checkpoints are fetched by the rfdetr package into ~/.roboflow/models.
"""

import glob
import hashlib
import os
import sys

from rfdetr import RFDETRLarge, RFDETRSegLarge


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main(out):
    for name, build in [("detect", RFDETRLarge), ("segment", RFDETRSegLarge)]:
        directory = os.path.join(out, name)
        os.makedirs(directory, exist_ok=True)
        build().export(format="onnx", output_dir=directory)
        for path in glob.glob(os.path.join(directory, "*.onnx")):
            print(path, os.path.getsize(path), sha256(path))


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "onnx")