scripts/export_face.py

"""Export the RF-DETR Medium face-detection checkpoint to ONNX at 576x576.

Mirrors the practice in E:/projects/rfdetr-onnx/scripts/export.py: the rfdetr
package's own model.export(format="onnx"), opset 17, with onnx/onnxruntime/onnxsim
installed. num_classes is deliberately NOT passed so that rfdetr auto-aligns the
detection head to the checkpoint's own class_embed (1 logit).
"""

import glob
import hashlib
import os
import shutil
import sys

from rfdetr import RFDETRMedium

CKPT = r"E:/projects/model-work/ckpt/rfdetr_medium_face.pth"
OUTDIR = r"E:/projects/model-work/export"
FINAL = r"E:/projects/rfdetr-onnx/rf-detr-medium-face.onnx"


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():
    os.makedirs(OUTDIR, exist_ok=True)
    model = RFDETRMedium(pretrain_weights=CKPT, resolution=576)
    print("model_config.num_classes =", model.model_config.num_classes)
    print("model_config.resolution  =", model.model_config.resolution)
    print("class_embed out_features =", model.model.model.class_embed.out_features)
    model.export(format="onnx", output_dir=OUTDIR, opset_version=17)
    produced = sorted(glob.glob(os.path.join(OUTDIR, "*.onnx")))
    for path in produced:
        print("produced", path, os.path.getsize(path), sha256(path))
    if len(produced) != 1:
        print("expected exactly one .onnx, got", produced)
        sys.exit(1)
    if os.path.exists(FINAL):
        print("refusing to overwrite existing", FINAL)
        sys.exit(1)
    shutil.copyfile(produced[0], FINAL)
    print("final", FINAL, os.path.getsize(FINAL), sha256(FINAL))
    print("source_pth", CKPT, os.path.getsize(CKPT), sha256(CKPT))


if __name__ == "__main__":
    main()