|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """
|
| Re-shapes a published face graph so Imagent's YunetFaceModel can run it through Windows ML. The
|
| weights are not touched: only the graph's declared shapes, and for YuNet the constants of its
|
| final Reshapes. Tooling only: Imagent itself has no Python in it. Run it through
|
| Imagent.ModelTools (`export yunet-onnx`, `export sface-onnx`, `export auraface-onnx`), or:
|
|
|
| uv run export_faces.py --model yunet --repo opencv/face_detection_yunet --revision <commit> \
|
| --file face_detection_yunet_2023mar.onnx --out <folder>
|
|
|
| Why each needs it, found by running them under ONNX Runtime (OpenCV, which they were published
|
| for, reshapes a network at load time; ONNX Runtime runs the shapes the file declares):
|
|
|
| yunet declared [1, 3, 640, 640]. Its convolutions run at any size once the dims are
|
| symbolic, but its last Reshapes are [1, -1, k], which folds a batch into the
|
| anchors: they become [0, -1, k], so each image keeps its own row.
|
| sface declared batch 1, and MXNet's exporter left every intermediate shape in value_info,
|
| so ONNX Runtime refused a batch of 4 ("Shape mismatch attempting to re-use buffer").
|
| auraface batches already; only its output is declared [1, 512].
|
|
|
| The contract, which YunetFaceModel's summary also states:
|
|
|
| yunet.onnx input input float [batch, 3, height, width] BGR 0-255, sides multiples of 32
|
| output cls_S, obj_S [batch, A_S, 1]; bbox_S [batch, A_S, 4]; kps_S [batch, A_S, 10]
|
| for strides S = 8, 16, 32, A_S = (height/S) * (width/S)
|
| sface.onnx input data float [batch, 3, 112, 112] aligned RGB 0-255
|
| output embedding float [batch, 128] not normalised
|
| auraface.onnx input data float [batch, 3, 112, 112] aligned RGB, (px - 127.5) / 127.5
|
| output embedding float [batch, 512] not normalised
|
|
|
| The batch dimension is named "batch" - the name the app's session pool overrides. Every graph is
|
| checked against the file it came from - same input, same output - and at a second batch and
|
| size, image by image, before it is written.
|
| """
|
|
|
| import argparse
|
| import sys
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import onnx
|
| import onnxruntime as ort
|
| from onnx import helper, numpy_helper
|
|
|
|
|
| THREADS = 4
|
|
|
|
|
| def session(model):
|
| options = ort.SessionOptions()
|
| options.intra_op_num_threads = THREADS
|
| options.inter_op_num_threads = 1
|
| options.log_severity_level = 3
|
| return ort.InferenceSession(model.SerializeToString(), options, providers=["CPUExecutionProvider"])
|
|
|
|
|
| def name_dims(value, names):
|
| """Sets the dims of an input or output: a str is symbolic, an int fixed."""
|
| dims = value.type.tensor_type.shape.dim
|
| for i, n in enumerate(names):
|
| if isinstance(n, str):
|
| dims[i].dim_param = n
|
| else:
|
| dims[i].dim_value = n
|
|
|
|
|
| def strip_shapes(model):
|
| """Drops recorded intermediate shapes, and weights listed as inputs, so ONNX Runtime infers its own."""
|
| del model.graph.value_info[:]
|
| weights = {i.name for i in model.graph.initializer}
|
| keep = [i for i in model.graph.input if i.name not in weights]
|
| del model.graph.input[:]
|
| model.graph.input.extend(keep)
|
|
|
|
|
| def rename_output(model, old, new):
|
| """Publishes an output under a stable name through an Identity, leaving the graph's own names alone."""
|
| output = next(o for o in model.graph.output if o.name == old)
|
| model.graph.node.append(helper.make_node("Identity", [old], [new], name=f"{new}_identity"))
|
| output.name = new
|
|
|
|
|
| def reshape_yunet(model):
|
| constants = {i.name: i for i in model.graph.initializer}
|
| for node in model.graph.node:
|
| if node.op_type == "Constant":
|
| for a in node.attribute:
|
| if a.t.dims:
|
| constants[node.output[0]] = a.t
|
| patched = 0
|
| for node in model.graph.node:
|
| if node.op_type != "Reshape" or node.input[1] not in constants:
|
| continue
|
| shape = numpy_helper.to_array(constants[node.input[1]])
|
| if shape.tolist()[:2] != [1, -1]:
|
| continue
|
|
|
| fixed = numpy_helper.from_array(np.array([0, -1, shape[2]], dtype=np.int64), f"{node.name}_shape_batched")
|
| model.graph.initializer.append(fixed)
|
| node.input[1] = fixed.name
|
| patched += 1
|
| if patched != 12:
|
| sys.exit(f"Expected YuNet's 12 output Reshapes, found {patched}: the graph is not the one this script knows.")
|
|
|
| name_dims(model.graph.input[0], ["batch", 3, "height", "width"])
|
| for output in model.graph.output:
|
| stride = output.name.split("_")[1]
|
| name_dims(output, ["batch", f"anchors{stride}", output.type.tensor_type.shape.dim[2].dim_value])
|
| strip_shapes(model)
|
|
|
|
|
| def reshape_recognizer(model, width):
|
| strip_shapes(model)
|
| name_dims(model.graph.input[0], ["batch", 3, 112, 112])
|
| old = model.graph.output[0].name
|
| name_dims(model.graph.output[0], ["batch", width])
|
| rename_output(model, old, "embedding")
|
|
|
|
|
| def check(original, exported, feeds_one, feeds_many, split):
|
| """The export matches the original at its own shape, and a larger run matches image by image."""
|
| ref = session(original)
|
| new = session(exported)
|
| names = [o.name for o in ref.get_outputs()]
|
|
|
| a = ref.run(None, feeds_one)
|
| b = new.run(None, feeds_one)
|
| for name, x, y in zip(names, a, b):
|
| diff = float(np.max(np.abs(x - y)))
|
| print(f" {name}: max |diff| against the original {diff:.2e}")
|
| if diff > 1e-4:
|
| sys.exit(f"{name} differs from the original graph by {diff}.")
|
|
|
| many = new.run(None, feeds_many)
|
| for i, single in enumerate(split(feeds_many)):
|
| one = new.run(None, single)
|
| for name, x, y in zip(names, one, many):
|
| diff = float(np.max(np.abs(x[0] - y[i])))
|
| if diff > 1e-3:
|
| sys.exit(f"{name}: image {i} of the batch differs from running it alone by {diff}.")
|
| print(f" batch of {len(split(feeds_many))}: every image matches its own run")
|
|
|
|
|
| def main():
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument("--model", choices=["yunet", "sface", "auraface"], required=True)
|
| parser.add_argument("--repo", required=True)
|
| parser.add_argument("--revision", required=True, help="a commit, never a branch")
|
| parser.add_argument("--file", required=True, help="the graph's name in the source repo")
|
| parser.add_argument("--out", type=Path, required=True)
|
| args = parser.parse_args()
|
| args.out.mkdir(parents=True, exist_ok=True)
|
|
|
| from huggingface_hub import hf_hub_download
|
| source = hf_hub_download(args.repo, args.file, revision=args.revision)
|
| original = onnx.load(source)
|
| exported = onnx.load(source)
|
| rng = np.random.default_rng(7)
|
| print(f"{args.repo}@{args.revision}/{args.file} -> {args.model}.onnx")
|
|
|
| if args.model == "yunet":
|
| reshape_yunet(exported)
|
| one = {"input": (rng.random((1, 3, 640, 640)) * 255).astype(np.float32)}
|
|
|
| many = {"input": (rng.random((3, 3, 1280, 1280)) * 255).astype(np.float32)}
|
| split = lambda f: [{"input": f["input"][i:i + 1]} for i in range(len(f["input"]))]
|
| else:
|
| first = [n.op_type for n in exported.graph.node[:3]]
|
| print(f" first operators: {first}")
|
| width = 128 if args.model == "sface" else 512
|
| reshape_recognizer(exported, width)
|
| scale = 255.0 if args.model == "sface" else 1.0
|
| one = {"data": (rng.random((1, 3, 112, 112)) * scale).astype(np.float32)}
|
| many = {"data": (rng.random((4, 3, 112, 112)) * scale).astype(np.float32)}
|
| split = lambda f: [{"data": f["data"][i:i + 1]} for i in range(len(f["data"]))]
|
|
|
| rename_output(original, original.graph.output[0].name, "embedding")
|
|
|
| onnx.checker.check_model(exported)
|
| check(original, exported, one, many, split)
|
| onnx.save(exported, args.out / f"{args.model}.onnx")
|
| print(f"wrote {args.out / f'{args.model}.onnx'}")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
