example.py

"""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()