40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser(description="Export trained RF-DETR to ONNX for deployment.")
|
||
|
|
parser.add_argument("--weights", type=Path, required=True, help="Path to trained checkpoint")
|
||
|
|
parser.add_argument("--output-onnx", type=Path, default=Path("model.onnx"), help="Output ONNX file")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
if not args.weights.exists():
|
||
|
|
raise SystemExit(f"Weights not found: {args.weights}")
|
||
|
|
|
||
|
|
from rfdetr import RFDETRBase
|
||
|
|
|
||
|
|
model = RFDETRBase(pretrain_weights=str(args.weights))
|
||
|
|
|
||
|
|
# Export to ONNX. This saves to the current directory by default, but we can specify.
|
||
|
|
# RF-DETR's export() method saves to 'output/model.onnx' I think, but let's check docs.
|
||
|
|
# From earlier fetch: model.export() saves to 'output' dir.
|
||
|
|
# But to make it flexible, perhaps run it and then move.
|
||
|
|
|
||
|
|
model.export() # This should create output/model.onnx
|
||
|
|
|
||
|
|
# Move to desired location
|
||
|
|
onnx_path = Path("output/model.onnx")
|
||
|
|
if onnx_path.exists():
|
||
|
|
onnx_path.rename(args.output_onnx)
|
||
|
|
print(f"Exported ONNX: {args.output_onnx}")
|
||
|
|
else:
|
||
|
|
raise SystemExit("ONNX export failed - check output/ dir")
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|