mecha10-models 0.9.0

Mecha10 model services — ML model management and Minio download support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

#!/usr/bin/env python3
"""
Embed ONNX "class_names" metadata

This script is embedded in the mecha10 CLI binary and automatically called
during model download to embed a model's already-resolved class names into
the model's own ONNX `metadata_props`, under a "class_names" key holding a
JSON array of strings - the exact shape `object-detector`'s
`ObjectDetectorNode::parse_class_names_metadata` reads back
(see packages/nodes/object-detector/src/lib.rs).

Usage:
    python embed_class_names.py <input_model.onnx> <class_names.json> <output_model.onnx>

Requirements:
    pip install onnx
"""

import json
import sys
from pathlib import Path

try:
    import onnx
except ImportError:
    print("Error: onnx not found", file=sys.stderr)
    print("Install with: pip install onnx", file=sys.stderr)
    sys.exit(1)


def embed_class_names(input_path: str, class_names_path: str, output_path: str):
    input_path = Path(input_path)
    class_names_path = Path(class_names_path)
    output_path = Path(output_path)

    if not input_path.exists():
        print(f"Error: Input model not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    if not class_names_path.exists():
        print(f"Error: Class names file not found: {class_names_path}", file=sys.stderr)
        sys.exit(1)

    class_names = json.loads(class_names_path.read_text())

    model = onnx.load(str(input_path))

    for prop in list(model.metadata_props):
        if prop.key == "class_names":
            model.metadata_props.remove(prop)

    entry = model.metadata_props.add()
    entry.key = "class_names"
    entry.value = json.dumps(class_names)

    onnx.save(model, str(output_path))
    print(f"✅ Embedded {len(class_names)} class names into {output_path.name}")


if __name__ == "__main__":
    if len(sys.argv) != 4:
        print(
            "Usage: embed_class_names.py <input_model.onnx> <class_names.json> <output_model.onnx>",
            file=sys.stderr,
        )
        sys.exit(1)

    embed_class_names(sys.argv[1], sys.argv[2], sys.argv[3])