mirror of
https://github.com/openmv/openmv.git
synced 2025-09-26 23:09:13 +08:00
148 lines
5.9 KiB
Python
Executable File
148 lines
5.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# This file is part of the OpenMV project.
|
|
#
|
|
# Copyright (C) 2025 OpenMV, LLC.
|
|
#
|
|
# This work is licensed under the MIT license, see the file LICENSE for details.
|
|
#
|
|
# AI models converter.
|
|
|
|
import sys
|
|
import os
|
|
import csv
|
|
import glob
|
|
import argparse
|
|
import binascii
|
|
import subprocess
|
|
import re
|
|
|
|
C_GREEN = '\033[92m'
|
|
C_RED = '\033[91m'
|
|
C_BLUE = '\033[94m'
|
|
C_RESET = '\033[0m'
|
|
|
|
|
|
def vela_compile(model_path, build_dir, vela_args):
|
|
vela_ini = os.path.dirname(os.path.abspath(__file__))
|
|
model = os.path.basename(os.path.splitext(model_path)[0])
|
|
|
|
# Construct the command
|
|
command = [
|
|
'vela',
|
|
*vela_args,
|
|
'--output-dir', build_dir,
|
|
'--config', f'{vela_ini}/vela.ini',
|
|
model_path
|
|
]
|
|
|
|
# Call the command and capture the output
|
|
try:
|
|
result = subprocess.run(command, check=True, text=True, capture_output=True)
|
|
except subprocess.CalledProcessError as e:
|
|
print(e.stderr, file=sys.stderr)
|
|
print(vela_args, file=sys.stderr)
|
|
|
|
csv_file_path = glob.glob(f"{build_dir}/{model}_summary_*.csv")[0]
|
|
with open(csv_file_path, mode='r') as file:
|
|
row = next(csv.DictReader(file))
|
|
stoi = lambda x, d=1: str(int(float(x) / d))
|
|
color = lambda c,x: c + x + C_RESET
|
|
|
|
summary = {
|
|
C_BLUE + "Network:": C_BLUE + row["network"],
|
|
C_BLUE + "Accelerator Configuration:": C_GREEN + row["accelerator_configuration"],
|
|
C_BLUE + "System Configuration:": C_BLUE + row["system_config"],
|
|
C_BLUE + "Memory Mode:": C_BLUE + row["memory_mode"],
|
|
C_BLUE + "Compiler Mode: ": C_RED + vela_args[-1],
|
|
C_BLUE + "Accelerator Clock:": C_BLUE + stoi(row["core_clock"], 10**6) + " MHz",
|
|
C_BLUE + "SRAM Usage:": C_RED + stoi(row["sram_memory_used"]) + " KiB",
|
|
C_BLUE + "Flash Usage:": C_RED + stoi(row["off_chip_flash_memory_used"]) + " KiB",
|
|
C_BLUE + "Inference Time:": C_GREEN + "%.2f ms, %.2f inferences/s"%
|
|
(float(row["inference_time"]) * 1000, float(row["inferences_per_second"])),
|
|
}
|
|
for key, value in summary.items():
|
|
print(f"{key:<{35}} {value:<{50}}", file=sys.stderr)
|
|
print(C_RESET, file=sys.stderr)
|
|
os.rename(f"{build_dir}/{model}_vela.tflite", f"{build_dir}/{model}.tflite")
|
|
|
|
def stedge_compile(model_path, build_dir, profile, stedge_args=None):
|
|
core_dir = os.path.realpath("tools/st/stedgeai/2.1")
|
|
config = os.path.realpath("tools/st/scripts/neuralart.json")
|
|
model_name = os.path.basename(os.path.splitext(model_path)[0])
|
|
model_ext = os.path.splitext(model_path)[1]
|
|
output_dir = os.path.join(build_dir, model_name)
|
|
|
|
# Remove any Make-related variables that could leak
|
|
env = os.environ.copy()
|
|
env["STEDGEAI_CORE_DIR"] = core_dir
|
|
for var in ["RM", "CFLAGS", "CPPFLAGS", "CXXFLAGS", "LDFLAGS", 'MAKEFLAGS']:
|
|
env.pop(var, None)
|
|
|
|
print(f"{C_GREEN}Creating relocatable binary model {model_name}{C_RESET}")
|
|
|
|
# Step 1: stedgeai generate
|
|
generate_command = [
|
|
os.path.join(core_dir, "Utilities/linux/stedgeai"),
|
|
"generate",
|
|
*stedge_args,
|
|
"--model", model_path,
|
|
"--relocatable",
|
|
"--st-neural-art", f"{profile}@{config}",
|
|
"--no-workspace",
|
|
"--output", os.path.join(output_dir, "gen"),
|
|
"--verbosity", "1",
|
|
]
|
|
|
|
try:
|
|
result = subprocess.run(generate_command, check=True, text=True, env=env,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"stedgeai command failed with exit code {e.returncode}", file=sys.stderr)
|
|
print(" ".join(generate_command), file=sys.stderr)
|
|
raise(e)
|
|
|
|
# Step 2: Python relocation script
|
|
reloc_command = [
|
|
sys.executable, # Uses current Python interpreter
|
|
os.path.join(core_dir, "scripts/N6_reloc/npu_driver.py"),
|
|
"--input", os.path.join(output_dir, "gen", "network.c"),
|
|
"--output", output_dir,
|
|
"--verbosity", "1",
|
|
]
|
|
|
|
try:
|
|
result = subprocess.run(reloc_command, check=True, text=True, env=env,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Relocation script failed with exit code {e.returncode}", file=sys.stderr)
|
|
print(" ".join(reloc_command), file=sys.stderr)
|
|
raise(e)
|
|
|
|
match = re.search(
|
|
r"([ \t]+XIP size.*?Table: mempool.*?\n)", result.stdout, re.DOTALL | re.MULTILINE
|
|
)
|
|
if match:
|
|
print(C_BLUE + match.group(1).rstrip() + C_RESET + "\n")
|
|
|
|
os.rename(f"{output_dir}/network_rel.bin", f"{build_dir}/{model_name}{model_ext}")
|
|
|
|
if __name__ == '__main__':
|
|
# python tools/tflite2c.py --input lib/models/fomo_face_detection.tflite --build-dir /tmp/build_st --stedge-args "--target stm32n6"
|
|
# python tools/tflite2c.py --input lib/models/fomo_face_detection.tflite --build-dir /tmp/build --vela-args "--system-config RTSS_HP_DTCM_MRAM --accelerator-config ethos-u55-256 --memory-mode Shared_Sram"
|
|
parser = argparse.ArgumentParser(description='AI models converter.')
|
|
parser.add_argument('--input', action = 'store', help = 'Input model.', required=True)
|
|
parser.add_argument('--build-dir', action = 'store', help='Build directory.', default='build')
|
|
parser.add_argument('--vela-args', action = 'store', help='Vela compiler args.')
|
|
parser.add_argument('--stedge-args', action = 'store', help='STEdge AI tools args.')
|
|
parser.add_argument('--stedge-profile', action = 'store', help='STEdge AI tools profile.', default="default")
|
|
args = parser.parse_args()
|
|
|
|
if args.vela_args:
|
|
# Compile the model using Vela.
|
|
vela_compile(args.input, args.build_dir, args.vela_args.split())
|
|
elif args.stedge_args:
|
|
# Compile the model using STEdge AI tools.
|
|
stedge_compile(args.input, args.build_dir, args.stedge_profile, args.stedge_args.split())
|
|
else:
|
|
parser.print_help(sys.stderr)
|