import argparse

from ultralytics import YOLO


def main(checkpoint_path):
    model = YOLO(checkpoint_path)
    model.export(
        format="engine",
        imgsz=(640, 480),
        dynamic=True,
        simplify=True,
        half=True,
        workspace=0.8,
        batch=1,
        device=0
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("checkpoint_path", type=str)
    args = parser.parse_args()
    main(args.checkpoint_path)