| |
| """Download the ONNX listed by Rockchip and invoke its unmodified converters.""" |
| import argparse |
| import concurrent.futures |
| import hashlib |
| from importlib.metadata import version |
| import json |
| from pathlib import Path |
| import re |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import yaml |
|
|
| DETECTORS = ("yolov5", "yolov6", "yolov7", "yolov8", "yolov10", "yolo11", |
| "yolo26", "yolox", "ppyoloe", "yolo_world") |
| TASKS = ("yolov5_seg", "yolov8_seg", "yolov8_pose", "yolov8_obb") |
| SUPPLIED_ONNX = { |
| "yolo26n": "https://meta.zbox.filez.com/v/link/view/dddfdaa76107404a8a358673fa31a041", |
| "yolo26m": "https://meta.zbox.filez.com/v/link/view/198e90ee6bc94c71ad2002d5f9897ac4", |
| } |
|
|
|
|
| def inventory(zoo): |
| for family in (*DETECTORS, *TASKS): |
| readme = zoo / "examples" / family / "README.md" |
| urls = re.findall(r"\((https://[^)]+\.onnx)\)", readme.read_text()) |
| assets = {Path(url).stem: url for url in urls} |
| if family in ("yolov8_pose", "yolov8_obb"): |
| assets = {name: url for name, url in assets.items() if "v5n-" in name or "v8n-" in name} |
| if family == "yolo26": |
| assets.update({"yolo26n": None, "yolo26m": None}) |
| for name, url in assets.items(): |
| yield family, name, url |
|
|
|
|
| def download(root, zoo, family, name, url): |
| destination = root / "onnx" / family / f"{name}.onnx" |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| if not destination.is_file(): |
| local = zoo / "examples" / family / "model" / destination.name |
| if not url and local.is_file(): |
| shutil.copy2(local, destination) |
| elif url: |
| partial = destination.with_suffix(".part") |
| subprocess.run(["curl", "-fLsS", "--retry", "2", "--connect-timeout", "12", |
| "--max-time", "180", "-o", partial, url], check=True) |
| partial.replace(destination) |
| else: |
| raise FileNotFoundError( |
| f"Download {SUPPLIED_ONNX[name]} and save its ONNX as {destination}") |
| import onnx |
| onnx.load(str(destination), load_external_data=False) |
| return destination |
|
|
|
|
| def digest(path): |
| checksum = hashlib.sha256() |
| with path.open("rb") as stream: |
| for block in iter(lambda: stream.read(1024 * 1024), b""): |
| checksum.update(block) |
| return checksum.hexdigest() |
|
|
|
|
| def write_metadata(zoo, family, model, destination, chip): |
| if family == "yolo_world": |
| return |
| import onnx |
| inputs = onnx.load(str(model), load_external_data=False).graph.input |
| shape = [dim.dim_value for dim in inputs[0].type.tensor_type.shape.dim] |
| if len(shape) != 4 or shape[:2] != [1, 3] or min(shape[2:]) <= 0: |
| raise RuntimeError(f"Expected static NCHW RGB model, got {shape}") |
| assets = zoo / "examples" / family / "model" |
| labels = next(assets.glob("*labels_list.txt")) |
| tasks = {"yolov5_seg": "segment", "yolov8_seg": "segment", |
| "yolov8_pose": "pose", "yolov8_obb": "obb"} |
| task = tasks.get(family, "detect") |
| info = { |
| "task": task, "batch": 1, "channels": 3, "imgsz": shape[2:], |
| "names": [line.strip() for line in labels.read_text().splitlines() if line.strip()], |
| "target_platform": chip, |
| "coordinates": "pixels", |
| } |
| if family in ("yolov5", "yolov5_seg", "yolov7"): |
| values = [float(value) for value in next(assets.glob("anchors_*.txt")).read_text().split()] |
| info.update(head="anchored", anchors=[values[i:i + 6] for i in range(0, len(values), 6)]) |
| if family == "yolox": |
| info["head"] = "yolox" |
| if family in ("yolov10", "yolo26"): |
| info["end2end"] = True |
| if task == "pose": |
| info["kpt_shape"] = [17, 3] |
| (destination / "metadata.yaml").write_text(yaml.safe_dump(info, sort_keys=False)) |
|
|
|
|
| def convert(root, zoo, asset, chip, precision, refresh_metadata=False): |
| family, name, model, source = asset |
| destination = root / "rknn" / chip / family / name / precision |
| destination.mkdir(parents=True, exist_ok=True) |
| output = destination / "model.rknn" |
| receipt = destination / "conversion.json" |
| if not refresh_metadata and output.is_file() and receipt.is_file() and ( |
| family == "yolo_world" or (destination / "metadata.yaml").is_file()): |
| return f"cached {family}/{name} {chip} {precision}" |
| work_root = root / "scratch" |
| work_root.mkdir(exist_ok=True) |
| with tempfile.TemporaryDirectory(dir=work_root, prefix=f"{name}-{chip}-{precision}-") as scratch: |
| |
| tree = Path(scratch) |
| (tree / "datasets").symlink_to(zoo / "datasets") |
| example = tree / "examples" / family |
| example.mkdir(parents=True) |
| (example / "model").symlink_to(zoo / "examples" / family / "model") |
| working = example / "python" |
| shutil.copytree(zoo / "examples" / family / "python", working) |
| staged = destination / "model.partial.rknn" |
| commands = [] |
| if family == "yolo_world": |
| subdir = "clip_text" if name == "clip_text" else "yolo_world" |
| script = f"{subdir}/convert.py" |
| commands.append([script, str(model), chip, precision, str(staged)]) |
| elif family == "yolov8_pose": |
| if precision == "i8": |
| commands.append(["convert.py", str(model), chip, "i8", "1"]) |
| commands.append(["convert.py", str(model), chip, "i8", "2", str(staged)]) |
| else: |
| commands.append(["convert.py", str(model), chip, "fp", "1", str(staged)]) |
| else: |
| commands.append(["convert.py", str(model), chip, precision, str(staged)]) |
| if not output.is_file(): |
| with (destination / "convert.log").open("w") as log: |
| for command in commands: |
| print(subprocess.list2cmdline(command), file=log, flush=True) |
| subprocess.run([sys.executable, *command], cwd=working, |
| stdout=log, stderr=subprocess.STDOUT, check=True) |
| if not staged.is_file() or not staged.stat().st_size: |
| raise RuntimeError("Converter did not produce an RKNN file") |
| staged.replace(output) |
| write_metadata(zoo, family, model, destination, chip) |
| receipt.write_text(json.dumps({ |
| "source": source, "onnx_sha256": digest(model), "rknn_sha256": digest(output), |
| "commands": commands, "toolkit_python": sys.executable, |
| "toolkit_version": version("rknn-toolkit2"), |
| "zoo_commit": subprocess.check_output( |
| ["git", "-C", str(zoo), "rev-parse", "HEAD"], text=True).strip(), |
| }, indent=2) + "\n") |
| return f"ok {family}/{name} {chip} {precision}" |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| root = Path(__file__).resolve().parent |
| parser.add_argument("--zoo", type=Path, default=root.parent / "rknn_model_zoo", |
| help="Model Zoo checkout (default: ../rknn_model_zoo)") |
| parser.add_argument("--jobs", type=int, default=4) |
| parser.add_argument("--model", action="append", help="Limit the run to named variants") |
| parser.add_argument("--chip", choices=("rk3576", "rk3588"), action="append") |
| parser.add_argument("--precision", choices=("i8", "fp"), action="append") |
| parser.add_argument("--refresh-metadata", action="store_true") |
| args = parser.parse_args() |
| zoo = args.zoo.resolve() |
| inventory_rows = [item for item in inventory(zoo) if not args.model or item[1] in args.model] |
| failures = 0 |
| with concurrent.futures.ThreadPoolExecutor(max_workers=6) as downloads, \ |
| concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as workers: |
| pending = {downloads.submit(download, root, zoo, *item): item for item in inventory_rows} |
| builds = {} |
| for future in concurrent.futures.as_completed(pending): |
| family, name, source = pending[future] |
| try: |
| model = future.result() |
| asset = family, name, model, source or SUPPLIED_ONNX.get(name) or str(model) |
| for chip in args.chip or ("rk3576", "rk3588"): |
| for precision in args.precision or ("i8", "fp"): |
| if name == "clip_text" and precision == "i8": |
| continue |
| job = workers.submit(convert, root, zoo, asset, chip, precision, |
| args.refresh_metadata) |
| builds[job] = family, name, chip, precision |
| except Exception as error: |
| print(f"asset failed {family}/{name}: {error}", flush=True) |
| failures += 1 |
| for future in concurrent.futures.as_completed(builds): |
| try: |
| print(future.result(), flush=True) |
| except Exception as error: |
| print(f"build failed {' '.join(builds[future])}: {error}", flush=True) |
| failures += 1 |
| return bool(failures) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
