"""Run an exported image-classification Ethos-U85 PTE on Corstone-320.""" from __future__ import annotations import argparse import json import os import shutil import subprocess import uuid from pathlib import Path import numpy as np from PIL import Image BUNDLE_DIR = Path(__file__).resolve().parent SHARED_RUNTIME_DIR = Path("/opt/ethos-u85-graviton") LOCAL_RUNTIME_DIR = BUNDLE_DIR / "runtime" def _default_runtime_dir() -> Path: """Mirror install_ethos_u85_graviton.sh's own fallback: it prefers the shared /opt install, but silently falls back to a local, per-bundle one when sudo isn't available. Detect whichever one actually got built.""" override = os.environ.get("ETHOS_RUNTIME_DIR") if override: return Path(override) if (SHARED_RUNTIME_DIR / "bin" / "arm_executor_runner").exists(): return SHARED_RUNTIME_DIR return LOCAL_RUNTIME_DIR RUNTIME_DIR = _default_runtime_dir() DEFAULT_MODEL = BUNDLE_DIR / "deit-tiny_ethos_ethosu_optimized.pte" DEFAULT_IMAGE = BUNDLE_DIR / "sample_input.jpg" DEFAULT_FVP = RUNTIME_DIR / "bin" / "FVP_Corstone_SSE-320" DEFAULT_RUNNER = RUNTIME_DIR / "bin" / "arm_executor_runner" DEFAULT_WORKDIR = RUNTIME_DIR / "output" / "fvp" MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) IMAGENET_CLASSES = json.loads((BUNDLE_DIR / "imagenet_classes.json").read_text(encoding="utf-8")) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Execute an Ethos-U85 image-classification PTE on Corstone-320." ) parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) parser.add_argument( "--image", type=Path, default=DEFAULT_IMAGE, help="RGB image to classify (default: sample_input.jpg next to this script).", ) parser.add_argument("--fvp-bin", type=Path, default=DEFAULT_FVP) parser.add_argument("--runner-elf", type=Path, default=DEFAULT_RUNNER) parser.add_argument("--workdir", type=Path, default=DEFAULT_WORKDIR) parser.add_argument( "--output", type=Path, help="Optional JSON destination; omit to only print predictions.", ) parser.add_argument("--timelimit", type=int, default=1800) return parser.parse_args() def load_image(path: Path) -> Image.Image: print(f"Using image: {path}") return Image.open(path).convert("RGB") def preprocess(path: Path) -> np.ndarray: image = load_image(path) width, height = image.size if width < height: new_width, new_height = 256, round(height * 256 / width) else: new_height, new_width = 256, round(width * 256 / height) image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) left = (new_width - 224) // 2 top = (new_height - 224) // 2 image = image.crop((left, top, left + 224, top + 224)) array = np.asarray(image, dtype=np.float32) / 255.0 array = array.transpose(2, 0, 1) array = (array - MEAN[:, None, None]) / STD[:, None, None] return np.expand_dims(array.astype(np.float32), axis=0) def fvp_environment(fvp: Path) -> dict[str, str]: env = os.environ.copy() resolved = fvp.resolve() for ancestor in resolved.parents: candidate = ancestor / "python" / "lib" if candidate.is_dir() and any(candidate.glob("libpython*.so*")): current = env.get("LD_LIBRARY_PATH") env["LD_LIBRARY_PATH"] = f"{candidate}:{current}" if current else str(candidate) break return env def main() -> None: args = parse_args() model = args.model.resolve() fvp = args.fvp_bin.resolve() runner = args.runner_elf.resolve() for kind, path in (("model", model), ("FVP", fvp), ("runner", runner)): if not path.exists(): raise FileNotFoundError(f"Missing {kind}: {path}") if not args.image.is_file(): raise FileNotFoundError(f"Missing image: {args.image}") workdir = args.workdir.resolve() workdir.mkdir(parents=True, exist_ok=True) run_id = uuid.uuid4().hex[:8] # workdir defaults under the SHARED ETHOS_RUNTIME_DIR, so a fixed # "model.pte" name would collide with a concurrent run of another model # bundle. Keep it unique per run_id like the input/output files. staged_model = workdir / f"model_{run_id}.pte" input_path = workdir / f"input_{run_id}.bin" output_base = f"out_{run_id}" output_path = workdir / f"{output_base}-0.bin" shutil.copyfile(model, staged_model) input_path.write_bytes(preprocess(args.image).tobytes()) command_line = ( f"arm_executor_runner -m {staged_model.name} -i {input_path.name} -o {output_base}" ) command = [ str(fvp), "-C", "mps4_board.subsystem.ethosu.num_macs=256", "-C", "mps4_board.visualisation.disable-visualisation=1", "-C", "vis_hdlcd.disable_visualisation=1", "-C", "mps4_board.telnetterminal0.start_telnet=0", "-C", "mps4_board.uart0.out_file=-", "-C", "mps4_board.uart0.shutdown_on_eot=1", "-C", "mps4_board.subsystem.cpu0.semihosting-enable=1", "-C", "mps4_board.subsystem.ethosu.extra_args='--fast'", "-C", "mps4_board.subsystem.cpu0.semihosting-stack_base=0", "-C", "mps4_board.subsystem.cpu0.semihosting-heap_limit=0", "-C", f"mps4_board.subsystem.cpu0.semihosting-cwd={workdir}", "-C", f"mps4_board.subsystem.cpu0.semihosting-cmd_line='{command_line}'", "-a", str(runner), "--timelimit", str(args.timelimit), ] result = subprocess.run( command, capture_output=True, text=True, timeout=args.timelimit + 30, check=False, env=fvp_environment(fvp), ) print(result.stdout, end="") if result.returncode != 0: raise RuntimeError( f"FVP exited with {result.returncode}\n{result.stderr[-2048:]}" ) if not output_path.is_file(): raise RuntimeError(f"FVP did not produce {output_path}") logits = np.fromfile(output_path, dtype=np.float32) if logits.size != 1000: raise ValueError(f"Expected 1000 float32 logits, got {logits.size}") probabilities = np.exp(logits.astype(np.float64) - logits.max()) probabilities /= probabilities.sum() top5 = np.argsort(probabilities)[::-1][:5] predictions = [ { "rank": rank, "class_index": int(index), "class_name": IMAGENET_CLASSES[int(index)], "probability": float(probabilities[index]), } for rank, index in enumerate(top5, 1) ] print("Top-5 ImageNet predictions:") for prediction in predictions: print( f" {prediction['rank']}. index={prediction['class_index']}, " f"class={prediction['class_name']}, " f"probability={prediction['probability']:.6f}" ) print(f"Raw output: {output_path}") if args.output is not None: predictions_path = args.output.resolve() predictions_path.parent.mkdir(parents=True, exist_ok=True) predictions_path.write_text( json.dumps(predictions, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) print(f"Predictions JSON: {predictions_path}") if __name__ == "__main__": main()