File size: 4,291 Bytes
9840cf0 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #!/usr/bin/env python3
"""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
# Conversion runs in an existing toolkit environment; do not install packages for the user.
os.environ["YOLO_AUTOINSTALL"] = "false"
from ultralytics import YOLO
from ultralytics.utils import YAML
previous_directory = Path.cwd()
try:
# Exporters write intermediate files beside the checkpoint. Keep the original untouched.
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()
|