File size: 1,852 Bytes
72748e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """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()
|