File size: 2,381 Bytes
0f5efd8 | 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 54 55 56 57 58 59 60 61 62 | """Export the RF-DETR Medium licence-plate checkpoint to ONNX at 576x576.
The checkpoint is our own fine-tune, made with `rfdetr==1.10.1` from
Roboflow's rf-detr-medium.pth on Open Images V7's "Vehicle registration plate"
boxes (annotations CC BY 4.0), the group-of boxes dropped: 4,497 training
images with 6,554 plates, 724 validation images with 985. Twenty epochs at
576, batch 8 with two steps of gradient accumulation, on one RTX 4090; the
EMA weights with the best validation mAP, epoch 4 at 0.655 mAP 50:95, are
the checkpoint exported here.
The export mirrors scripts/export_face.py: the rfdetr package's own
model.export(format="onnx"), opset 17. num_classes is NOT passed, so the head
is exported as trained. rfdetr 1.10.1 lays a one-class head out as the class
plus one spare slot, so `labels` is 1x300x2: index 0 is the plate, index 1
was never trained toward anything and stays quiet.
"""
import glob
import hashlib
import os
import shutil
import sys
from rfdetr import RFDETRMedium
CKPT = r"E:/projects/model-work/ckpt/rfdetr_medium_plate.pth"
OUTDIR = r"E:/projects/model-work/export-plate"
FINAL = r"E:/projects/rfdetr-onnx/rf-detr-medium-plate.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()
|