93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare Qualcomm AI Hub calibration and validation inputs for the meter detector."""
|
|
|
|
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/meter-detection/data/electric-meter-detection"),
|
|
help="Root of the extracted Roboflow dataset.",
|
|
)
|
|
parser.add_argument(
|
|
"--calibration-dir",
|
|
type=Path,
|
|
default=Path("examples/meter-detection/data/aihub_calibration"),
|
|
help="Directory where .npy calibration samples will be written.",
|
|
)
|
|
parser.add_argument(
|
|
"--input-file",
|
|
type=Path,
|
|
default=Path("examples/meter-detection/data/inputs.npz"),
|
|
help="Validation .npz input file for qc-cli ai-hub validate.",
|
|
)
|
|
parser.add_argument("--input-name", default="images", help="ONNX input name.")
|
|
parser.add_argument("--image-size", type=int, default=640, help="Square image size used for ONNX export.")
|
|
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:
|
|
"""Apply Ultralytics-style letterboxing and produce an NCHW float32 tensor."""
|
|
with Image.open(path) as source:
|
|
image = source.convert("RGB")
|
|
|
|
scale = min(image_size / image.width, image_size / image.height)
|
|
resized_width = round(image.width * scale)
|
|
resized_height = round(image.height * scale)
|
|
image = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
|
|
|
|
canvas = Image.new("RGB", (image_size, image_size), (114, 114, 114))
|
|
left = round((image_size - resized_width) / 2 - 0.1)
|
|
top = round((image_size - resized_height) / 2 - 0.1)
|
|
canvas.paste(image, (left, top))
|
|
|
|
array = np.asarray(canvas, dtype=np.float32) / 255.0
|
|
return np.transpose(array, (2, 0, 1))[None, ...].astype(np.float32)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.image_size < 1:
|
|
raise SystemExit("--image-size must be at least 1")
|
|
if args.samples < 1:
|
|
raise SystemExit("--samples must be at least 1")
|
|
|
|
images = sorted(
|
|
path
|
|
for path in args.dataset_dir.rglob("*")
|
|
if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS and path.parent.name == "images"
|
|
)
|
|
if not images:
|
|
raise SystemExit(f"No images found under {args.dataset_dir}")
|
|
|
|
args.calibration_dir.mkdir(parents=True, exist_ok=True)
|
|
args.input_file.parent.mkdir(parents=True, exist_ok=True)
|
|
for stale_sample in args.calibration_dir.glob("sample_*.npy"):
|
|
stale_sample.unlink()
|
|
|
|
prepared: list[np.ndarray] = []
|
|
for index, image_path in enumerate(images[: args.samples]):
|
|
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]}) # pyright: ignore[reportArgumentType]
|
|
print(f"Wrote {len(prepared)} calibration samples to {args.calibration_dir}")
|
|
print(f"Wrote validation input to {args.input_file}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|