ai-hub (#3)

Reviewed-on: #3

examples/ai-hub/prepare_inputs.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Prepare Qualcomm AI Hub calibration and validation inputs for the training example."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--dataset-dir",
+        type=Path,
+        default=Path("examples/training/data/flower_photos_sagemaker"),
+        help="ImageFolder-style dataset used for training.",
+    )
+    parser.add_argument(
+        "--calibration-dir",
+        type=Path,
+        default=Path("examples/training/data/aihub_calibration"),
+        help="Directory where .npy calibration samples will be written.",
+    )
+    parser.add_argument(
+        "--input-file",
+        type=Path,
+        default=Path("examples/training/data/inputs.npz"),
+        help="Validation .npz input file for qc-cli ai-hub validate.",
+    )
+    parser.add_argument("--input-name", default="input", help="ONNX input name.")
+    parser.add_argument("--image-size", type=int, default=160, help="Square image size used by training.")
+    parser.add_argument("--samples", type=int, default=16, help="Number of calibration samples to write.")
+    return parser.parse_args()
+
+
+def preprocess_image(path: Path, image_size: int) -> np.ndarray:
+    image = Image.open(path).convert("RGB").resize((image_size, image_size), Image.Resampling.BILINEAR)
+    array = np.asarray(image, dtype=np.float32) / 255.0
+    array = np.transpose(array, (2, 0, 1))
+    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
+    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
+    return ((array - mean) / std)[None, ...].astype("float32")
+
+
+def main() -> None:
+    args = parse_args()
+    images = sorted(p for p in args.dataset_dir.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS)
+    if not images:
+        raise SystemExit(f"No images found under {args.dataset_dir}")
+    if args.samples < 1:
+        raise SystemExit("--samples must be at least 1")
+
+    args.calibration_dir.mkdir(parents=True, exist_ok=True)
+    args.input_file.parent.mkdir(parents=True, exist_ok=True)
+
+    sample_count = min(args.samples, len(images))
+    prepared = []
+    for index, image_path in enumerate(images[:sample_count]):
+        sample = preprocess_image(image_path, args.image_size)
+        np.save(args.calibration_dir / f"sample_{index:03d}.npy", sample)
+        prepared.append(sample)
+
+    np.savez(args.input_file, **{args.input_name: prepared[0]})
+    print(f"Wrote {sample_count} calibration samples to {args.calibration_dir}")
+    print(f"Wrote validation input to {args.input_file}")
+
+
+if __name__ == "__main__":
+    main()