Add source

source/convert-official.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""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:
+        # Keep the relative dataset and model paths expected by the official scripts.
+        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())

source/onnx/yolo26m.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:890c8b5241232f21a683e86708d98badb4d579e1998a4f05483819d0a945120c
+size 81963156

source/onnx/yolo26n.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c6e79ecfaa699806e71acdafd06455086280428ad8424a6ec3eeada5689020a
+size 9941956

source/onnx/yolo26s.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88dbcad6f8e1e70e3c32dd3f86d37ab41ec52e5a65bc638407d44c6bee73bb6c
+size 38291128