| |
| """Export an Ultralytics checkpoint and its metadata for one Rockchip chip.""" |
|
|
| import argparse |
| import os |
| from pathlib import Path |
| import shutil |
| import tempfile |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("model", help="Checkpoint path or official model name, e.g. yolo11n-seg.pt") |
| parser.add_argument("--chip", required=True, choices=("rk3576", "rk3588")) |
| parser.add_argument("--output", required=True, type=Path, help="New output directory") |
| parser.add_argument("--imgsz", type=int, help="Square input size; defaults depend on the task") |
| parser.add_argument("--precision", choices=("fp16", "int8", "mixed"), default="fp16") |
| parser.add_argument("--data", type=Path, help="Calibration YAML, image list, or directory; detection INT8 needs a dataset YAML") |
| parser.add_argument("--calibration-limit", type=int, default=100, |
| help="Maximum calibration images for mixed precision (default: 100)") |
| args = parser.parse_args() |
|
|
| destination = args.output.resolve() |
| if destination.exists(): |
| parser.error("Output already exists; choose a new directory") |
| if args.precision != "fp16" and not args.data: |
| parser.error("Quantization requires --data with representative calibration images") |
| if args.imgsz and (args.imgsz < 32 or args.imgsz > 4096 or args.imgsz % 32): |
| parser.error("--imgsz must be a multiple of 32 between 32 and 4096") |
| if not 1 <= args.calibration_limit <= 10000: |
| parser.error("--calibration-limit must be between 1 and 10000") |
| checkpoint = Path(args.model).expanduser() |
| checkpoint = checkpoint.resolve() if checkpoint.is_file() else None |
| dataset = args.data.resolve() if args.data else None |
|
|
| |
| os.environ["YOLO_AUTOINSTALL"] = "false" |
| from ultralytics import YOLO |
| from ultralytics.utils import YAML |
|
|
| previous_directory = Path.cwd() |
| try: |
| |
| with tempfile.TemporaryDirectory(prefix="nvr-export-") as scratch: |
| os.chdir(scratch) |
| if checkpoint: |
| shutil.copy2(checkpoint, checkpoint.name) |
| model = YOLO(checkpoint.name if checkpoint else args.model) |
| size = args.imgsz or {"classify": 224, "semantic": 512, "depth": 384}.get(model.task, 640) |
| optimized = (args.precision == "mixed" or model.task in {"semantic", "depth"} |
| or (args.precision == "int8" and model.task != "detect")) |
| if optimized: |
| from rknn_export import export_optimized |
| binary, metadata = export_optimized( |
| model, args.chip, size, dataset, args.calibration_limit, |
| args.precision) |
| else: |
| settings = dict(format="rknn", name=args.chip, imgsz=size, batch=1, |
| quantize=8 if args.precision == "int8" else 16, |
| device="cpu", simplify=True, opset=19) |
| if dataset: |
| settings["data"] = str(dataset) |
| exported = Path(model.export(**settings)) |
| binaries = list(exported.glob("*.rknn")) |
| if len(binaries) != 1: |
| raise RuntimeError("Expected one RKNN model in the export directory") |
| binary = binaries[0] |
| metadata = YAML.load(exported / "metadata.yaml") |
| metadata["precision"] = args.precision |
| metadata["target_platform"] = args.chip |
| metadata["coordinates"] = ("normalized" if args.precision == "int8" |
| and model.task == "detect" else "pixels") |
| metadata["preprocess"] = "center_crop" if model.task == "classify" else "letterbox" |
| destination.mkdir(parents=True) |
| shutil.copy2(binary, destination / "model.rknn") |
| YAML.save(destination / "metadata.yaml", metadata) |
| finally: |
| os.chdir(previous_directory) |
| print(f"Exported {model.task} for {args.chip}: {destination}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
