Full catalogue sync: example.py

example.py

@@ -27,14 +27,14 @@ from pathlib import Path
 
 import numpy as np
 from PIL import Image
+from torchvision.models import GoogLeNet_Weights
 
 
 SCRIPT_DIR = Path(__file__).resolve().parent
 REPO_ROOT = SCRIPT_DIR.parents[2]
 
-MODEL_PATH = SCRIPT_DIR / "deit-tiny-int8-ethos-u-executorch.pte"
+MODEL_PATH = SCRIPT_DIR / "deit-tiny_ethos_ethosu_optimized.pte"
 IMAGE_PATH = SCRIPT_DIR / "sample_input.jpg"
-LABELS_PATH = SCRIPT_DIR / "imagenet_classes.json"
 DEFAULT_WORKDIR = SCRIPT_DIR / "fvp_workdir"
 DEFAULT_RUNNER_ELF = (
     REPO_ROOT / "third_party/executorch-arm-runner/cmake-out/arm_executor_runner"
@@ -49,6 +49,10 @@ TOP_K = 5
 NUM_MACS = 256
 SEMIHOSTING_CMD_LINE_MAX = 256
 
+# ImageNet class labels (1000 classes)
+WEIGHTS = GoogLeNet_Weights.IMAGENET1K_V1
+IMAGENET_CLASSES = WEIGHTS.meta["categories"]
+
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
@@ -56,7 +60,6 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--model", type=Path, default=MODEL_PATH)
     parser.add_argument("--image", type=Path, default=IMAGE_PATH)
-    parser.add_argument("--labels", type=Path, default=LABELS_PATH)
     parser.add_argument("--workdir", type=Path, default=DEFAULT_WORKDIR)
     parser.add_argument(
         "--fvp-bin",
@@ -94,16 +97,6 @@ def preprocess(image_path: Path) -> np.ndarray:
     return np.expand_dims(array.astype(np.float32), axis=0)
 
 
-def load_labels(labels_path: Path) -> list[str] | None:
-    if not labels_path.exists():
-        return None
-    with labels_path.open() as f:
-        labels = json.load(f)
-    if isinstance(labels, dict):
-        return [labels[str(i)] for i in range(len(labels))]
-    return labels
-
-
 def build_fvp_cmd(
     *,
     fvp_bin: str,
@@ -230,7 +223,6 @@ def main() -> None:
     args = parse_args()
     model_path = args.model.resolve()
     image_path = args.image.resolve()
-    labels_path = args.labels.resolve()
     runner_elf = args.runner_elf.resolve()
 
     if not model_path.exists():
@@ -243,7 +235,6 @@ def main() -> None:
             "Build it with make install-corstone, or pass --runner-elf."
         )
 
-    labels = load_labels(labels_path)
     input_tensor = preprocess(image_path)
     raw_output = run_fvp(
         model_path=model_path,
@@ -253,7 +244,7 @@ def main() -> None:
         runner_elf=runner_elf,
         timelimit=args.timelimit,
     )
-    results = postprocess(raw_output, labels)
+    results = postprocess(raw_output, IMAGENET_CLASSES)
 
     print("Top-5 predictions:")
     for rank, item in enumerate(results, start=1):