| |
| """Export task heads that have no Model Zoo example and compile them with RKNN.""" |
|
|
| import argparse |
| import hashlib |
| import os |
| from pathlib import Path |
| import sys |
| import tempfile |
|
|
|
|
| TASKS = { |
| "yolov5n-cls": "classify", |
| "yolov8n-cls": "classify", |
| "yolo11n-cls": "classify", |
| } |
| YOLO26_TASKS = {"cls": "classify", "seg": "segment", "pose": "pose", "obb": "obb", |
| "sem": "semantic", "depth": "depth"} |
| TASKS.update({f"{family}{size}-cls": "classify" |
| for family in ("yolov5", "yolov8", "yolo11") for size in "sm"}) |
| TASKS.update({f"yolo26{size}-{suffix}": task |
| for size in "nsm" for suffix, task in YOLO26_TASKS.items()}) |
| SIZES = {"classify": 224, "semantic": 512, "depth": 384} |
| CHIPS = ("rk3576", "rk3588") |
| IMAGENET_MEAN = (0.485, 0.456, 0.406) |
| IMAGENET_STD = (0.229, 0.224, 0.225) |
|
|
|
|
| def load_yolov5_classifier(checkpoint, source): |
| """Bake YOLOv5's ImageNet normalization into the input graph.""" |
| import torch |
|
|
| sys.path.insert(0, str(source)) |
| try: |
| from models.experimental import attempt_load |
|
|
| classifier = attempt_load(str(checkpoint), device=torch.device("cpu"), fuse=True) |
| finally: |
| sys.path.pop(0) |
|
|
| class NormalizeClassifier(torch.nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.model = model |
| self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).reshape(1, 3, 1, 1)) |
| self.register_buffer("std", torch.tensor(IMAGENET_STD).reshape(1, 3, 1, 1)) |
|
|
| def forward(self, images): |
| return self.model((images - self.mean) / self.std).softmax(dim=1) |
|
|
| return NormalizeClassifier(classifier).eval(), classifier.names |
|
|
|
|
| def split_yolo26(head, task, size): |
| """Expose one-to-one heads before TopK, following the Zoo YOLO26 detection layout.""" |
| import torch |
| from ultralytics.utils.tal import make_anchors |
|
|
| def forward(features): |
| predictions = head.forward_head(features, **head.one2one) |
| anchors, strides = make_anchors(features, head.stride, 0.5) |
| head.anchors = anchors.transpose(0, 1) |
| head.strides = strides.transpose(0, 1) |
| raw_boxes = predictions["boxes"] |
| left_top, right_bottom = raw_boxes.split(2, dim=1) |
| anchor = head.anchors.unsqueeze(0) |
| if task == "obb": |
| angle = predictions["angle"] |
| offset_x, offset_y = ((right_bottom - left_top) / 2).split(1, dim=1) |
| cosine, sine = angle.cos(), angle.sin() |
| rotated = torch.cat((offset_x * cosine - offset_y * sine, |
| offset_x * sine + offset_y * cosine), dim=1) |
| boxes = torch.cat((anchor + rotated, left_top + right_bottom), dim=1) |
| else: |
| boxes = torch.cat((anchor - left_top, anchor + right_bottom), dim=1) |
|
|
| outputs = (predictions["scores"].sigmoid(), boxes * (head.strides / size)) |
| if task == "segment": |
| return (*outputs, predictions["mask_coefficient"], head.proto(features)) |
| if task == "obb": |
| return (*outputs, angle) |
|
|
| points = head.kpts_decode(predictions["kpts"]).reshape(1, *head.kpt_shape, -1) |
| |
| points = torch.cat((points[:, :, :2] / size, points[:, :, 2:]), dim=2) |
| return (*outputs, points.reshape(1, head.nk, -1)) |
|
|
| return forward |
|
|
|
|
| def export_checkpoint(checkpoint, destination, size, yolov5_source): |
| os.environ.setdefault("YOLO_AUTOINSTALL", "false") |
| import onnx |
| import torch |
| import ultralytics |
| import yaml |
| from ultralytics import YOLO |
|
|
| task = TASKS[checkpoint.stem] |
| if checkpoint.stem.startswith("yolov5"): |
| model, names = load_yolov5_classifier(checkpoint, yolov5_source) |
| head = None |
| else: |
| model = YOLO(str(checkpoint), task=task).model.cpu().float().eval() |
| head = model.model[-1] |
| names = getattr(model, "names", None) or {0: "depth"} |
| split = task in {"segment", "pose", "obb"} |
| if split: |
| if head.reg_max != 1 or getattr(head, "one2one_cv2", None) is None: |
| raise ValueError("Expected a DFL-free, one-to-one YOLO26 head") |
| |
| head.end2end = True |
| if head is not None: |
| model.fuse(verbose=False) |
| with checkpoint.open("rb") as file: |
| digest = hashlib.file_digest(file, "sha256").hexdigest() |
| metadata = { |
| "task": task, |
| "batch": 1, |
| "channels": 3, |
| "imgsz": [size, size], |
| "names": names, |
| "preprocess": "center_crop" if task == "classify" else "letterbox", |
| "coordinates": "pixels", |
| "end2end": False, |
| "export_source": "model-zoo-inspired-pytorch", |
| "checkpoint": checkpoint.name, |
| "checkpoint_sha256": digest, |
| "ultralytics": ultralytics.__version__, |
| "pytorch": str(torch.__version__), |
| } |
| if split: |
| head.export = True |
| head.forward = split_yolo26(head, task, size) |
| metadata.update(coordinates="normalized", end2end=True) |
| output_names = ["Sigmoid_scores", "boxes"] |
| if task == "segment": |
| output_names.extend(("mask_coefficients", "prototypes")) |
| elif task == "pose": |
| output_names.append("keypoints") |
| metadata["kpt_shape"] = list(head.kpt_shape) |
| else: |
| output_names.append("angle") |
| elif task == "classify": |
| if head is not None: |
| head.export = True |
| else: |
| metadata["checkpoint_source"] = ( |
| f"https://github.com/ultralytics/yolov5/releases/download/v7.0/{checkpoint.name}" |
| ) |
| metadata["normalization"] = "ImageNet mean/std baked into ONNX" |
| output_names = ["probabilities"] |
| else: |
| |
| head.export = False |
| metadata["dense_resize"] = "half_pixel" if task == "semantic" else "align_corners" |
| output_names = ["logits" if task == "semantic" else "depth"] |
|
|
| image = torch.zeros(1, 3, size, size) |
| with torch.inference_mode(): |
| expected = model(image) |
| tensors = expected if isinstance(expected, tuple) else (expected,) |
| if len(tensors) != len(output_names) or not all(torch.is_tensor(x) for x in tensors): |
| raise ValueError("Unexpected output contract from checkpoint head") |
| for name, tensor in zip(output_names, tensors): |
| print(f"{name}: {tuple(tensor.shape)}") |
| torch.onnx.export(model, image, str(destination / "model.onnx"), opset_version=13, |
| input_names=["images"], output_names=output_names, |
| dynamo=False, do_constant_folding=True) |
|
|
| graph = onnx.load(destination / "model.onnx") |
| onnx.checker.check_model(graph) |
| (destination / "metadata.yaml").write_text(yaml.safe_dump(metadata, sort_keys=False)) |
|
|
|
|
| def compile_graph(source, destination, chip, precision, dataset): |
| import yaml |
| from rknn.api import RKNN |
|
|
| metadata = yaml.safe_load((source / "metadata.yaml").read_text()) |
| metadata.update(target_platform=chip, precision="int8" if precision == "i8" else "fp16") |
|
|
| def check(code, step): |
| if code != 0: |
| raise RuntimeError(f"RKNN {step} failed ({code})") |
|
|
| toolkit = RKNN(verbose=False) |
| try: |
| check(toolkit.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], |
| target_platform=chip), "config") |
| check(toolkit.load_onnx(model=str(source / "model.onnx")), "load ONNX") |
| with tempfile.TemporaryDirectory(prefix="nvr-calibration-") as directory: |
| if precision == "i8" and metadata["checkpoint"].startswith("yolov5"): |
| dataset = center_crop_calibration(dataset, Path(directory), metadata["imgsz"][0]) |
| check(toolkit.build(do_quantization=precision == "i8", |
| dataset=str(dataset) if dataset else None), "build") |
| check(toolkit.export_rknn(str(destination / "model.rknn")), "export") |
| finally: |
| toolkit.release() |
| (destination / "metadata.yaml").write_text(yaml.safe_dump(metadata, sort_keys=False)) |
|
|
|
|
| def center_crop_calibration(dataset, directory, size): |
| """Use the same square crop and linear resize as YOLOv5 classify/predict.py.""" |
| import cv2 |
|
|
| paths = [] |
| for index, line in enumerate(dataset.read_text().splitlines()): |
| if not line.strip(): |
| continue |
| source = dataset.parent / line.strip() |
| image = cv2.imread(str(source), cv2.IMREAD_COLOR) |
| if image is None: |
| raise ValueError(f"Cannot read calibration image: {source}") |
| height, width = image.shape[:2] |
| side = min(height, width) |
| top, left = (height - side) // 2, (width - side) // 2 |
| image = cv2.resize(image[top:top + side, left:left + side], (size, size)) |
| output = directory / f"{index:04}.png" |
| if not cv2.imwrite(str(output), image): |
| raise RuntimeError(f"Cannot write calibration image: {output}") |
| paths.append(str(output)) |
| cropped = directory / "images.txt" |
| cropped.write_text("\n".join(paths) + "\n") |
| return cropped |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| commands = parser.add_subparsers(dest="command", required=True) |
| export = commands.add_parser("export", help="Export one original checkpoint to ONNX") |
| export.add_argument("checkpoint", type=Path) |
| export.add_argument("output", type=Path, help="New directory for the ONNX and metadata") |
| export.add_argument("--imgsz", type=int, help="Override the square input size") |
| export.add_argument("--yolov5-source", type=Path, help="YOLOv5 v7.0 source checkout") |
| compile_parser = commands.add_parser("compile", help="Build one RKNN from an exported directory") |
| compile_parser.add_argument("onnx_directory", type=Path) |
| compile_parser.add_argument("output", type=Path, help="New directory for the RKNN and metadata") |
| compile_parser.add_argument("--chip", required=True, choices=CHIPS) |
| compile_parser.add_argument("--precision", required=True, choices=("fp", "i8")) |
| compile_parser.add_argument("--dataset", type=Path, help="RKNN calibration image list") |
| args = parser.parse_args() |
|
|
| destination = args.output.resolve() |
| if destination.exists(): |
| parser.error(f"Output already exists: {destination}") |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(prefix=".convert-", dir=destination.parent) as directory: |
| work = Path(directory) / "output" |
| work.mkdir() |
| if args.command == "export": |
| checkpoint = args.checkpoint.resolve() |
| if not checkpoint.is_file() or checkpoint.stem not in TASKS: |
| parser.error("Provide one of these local checkpoints: " + ", ".join(TASKS)) |
| size = args.imgsz or SIZES.get(TASKS[checkpoint.stem], 640) |
| if size < 32 or size > 4096 or size % 32: |
| parser.error("--imgsz must be a multiple of 32 from 32 to 4096") |
| yolov5_source = args.yolov5_source.resolve() if args.yolov5_source else None |
| if checkpoint.stem.startswith("yolov5") and ( |
| yolov5_source is None or not (yolov5_source / "models/experimental.py").is_file() |
| ): |
| parser.error("YOLOv5 classification requires a --yolov5-source v7.0 checkout") |
| export_checkpoint(checkpoint, work, size, yolov5_source) |
| else: |
| dataset = args.dataset.resolve() if args.dataset else None |
| if args.precision == "i8" and (dataset is None or not dataset.is_file()): |
| parser.error("INT8 requires an existing --dataset image list") |
| compile_graph(args.onnx_directory.resolve(), work, args.chip, args.precision, dataset) |
| work.rename(destination) |
| print(destination) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
