59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Create Horizon hb_mapper calibration feature maps from the YOLO dataset."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
|
|
|
|
|
|
def resize_rgb(image: Image.Image, size: int) -> np.ndarray:
|
|
image = image.convert("RGB")
|
|
resized = image.resize((size, size), Image.Resampling.BILINEAR)
|
|
array = np.asarray(resized, dtype=np.float32)
|
|
return np.transpose(array, (2, 0, 1))[None, ...]
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--samples", type=int, default=128)
|
|
parser.add_argument("--size", type=int, default=640)
|
|
args = parser.parse_args()
|
|
|
|
roots = [args.dataset / "images" / "train", args.dataset / "images" / "vel"]
|
|
images = sorted(
|
|
p for root in roots if root.is_dir() for p in root.iterdir() if p.suffix.lower() in IMAGE_EXTS
|
|
)
|
|
if not images:
|
|
raise SystemExit(f"no images found below {args.dataset}")
|
|
count = min(args.samples, len(images))
|
|
# Evenly spread samples over the sorted set so calibration is not dominated by one split.
|
|
indices = np.linspace(0, len(images) - 1, count, dtype=np.int64)
|
|
selected = [images[int(i)] for i in indices]
|
|
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
for old in args.output.iterdir():
|
|
if old.is_file() and old.suffix in {".bin", ".rgbchw", ".txt"}:
|
|
old.unlink()
|
|
manifest = []
|
|
for index, path in enumerate(selected):
|
|
data = resize_rgb(Image.open(path), args.size)
|
|
data.tofile(args.output / f"{index:05d}.rgbchw")
|
|
manifest.append(str(path.relative_to(args.dataset)))
|
|
(args.output.parent / "calibration_manifest.txt").write_text(
|
|
"\n".join(manifest) + "\n", encoding="utf-8"
|
|
)
|
|
print(f"wrote {count} samples ({args.size}x{args.size}) to {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|