# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = [
#   "onnx==1.17.0",
#   "onnxruntime==1.20.1",
#   "numpy<2.2",
#   "huggingface_hub",
# ]
# ///
"""
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

# Small enough to leave a machine usable while it runs; these graphs are tiny.
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
    # A fresh constant per node, in case the exporter shared one between Reshapes.
    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)}
    # The second shape is the app's: a 1280 letterbox. Batched, not just resized.
    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"]))]
    # The original's single output keeps its old name; compare it with ours.
    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()