Raptor/validation/tools/yolo_local_image_validation.py

#!/usr/bin/env python3

import argparse
import json
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
from PIL import Image

SCRIPT_DIR = Path(__file__).resolve().parent
VALIDATION_DIR = SCRIPT_DIR.parent
if str(VALIDATION_DIR) not in sys.path:
    sys.path.insert(0, str(VALIDATION_DIR))

if sys.version_info < (3, 10):
    raise SystemExit(
        "yolo_local_image_validation.py requires Python 3.10+ because validation modules use modern type syntax. "
        "Run it with a newer interpreter, for example your project venv Python."
    )

from onnx_utils import _ONNX_TO_NP, onnx_io, write_inputs_to_memory_bin
from validate_one import MODE_COMPILE_ONLY, build_dump_ranges, run_pim_simulator, sanitize_output_name, validate_network
from yolo_real_image_validation import (
    IMAGE_CASES,
    decode_yolo_output,
    download_image,
    draw_detections,
    letterbox_rgb,
    save_tensor_csv,
    top_unique_labels,
)


def resolve_default_paths():
    validation_dir = Path(__file__).resolve().parent.parent
    repo_root = validation_dir.parent
    return {
        "validation_dir": validation_dir,
        "repo_root": repo_root,
        "network_dir": validation_dir / "networks" / "yolo11n" / "depth_51",
        "raptor_path": repo_root / "build_release" / "Release" / "bin" / "onnx-mlir",
        "onnx_include_dir": repo_root / "onnx-mlir" / "include",
        "simulator_dir": repo_root / "backend-simulators" / "pim" / "pim-simulator",
    }


def find_network_onnx(network_dir: Path) -> Path:
    onnx_files = sorted(network_dir.glob("*.onnx"))
    if not onnx_files:
        raise FileNotFoundError(f"No .onnx file found in {network_dir}")
    if len(onnx_files) > 1:
        names = ", ".join(path.name for path in onnx_files)
        raise RuntimeError(f"Expected exactly one .onnx file in {network_dir}, found: {names}")
    return onnx_files[0]


def local_case_paths(network_dir: Path, case_name: str):
    return {
        "root": network_dir,
        "runner": network_dir / "runner" / "build" / "runner",
        "runner_build": network_dir / "runner" / "build",
        "raptor_pim": network_dir / "raptor" / "pim",
        "real_root": network_dir / "real_image_validation",
        "input_csv": network_dir / "real_image_validation" / "inputs" / f"{case_name}.csv",
        "ref_dir": network_dir / "real_image_validation" / "reference" / case_name,
        "sim_dir": network_dir / "real_image_validation" / "simulation" / case_name,
        "sim_bin": network_dir / "real_image_validation" / "simulation" / case_name / "out.bin",
    }


def ensure_local_artifacts(args, network_onnx_path: Path):
    validate_network(
        network_onnx_path=network_onnx_path,
        raptor_path=args.raptor_path,
        onnx_include_dir=args.onnx_include_dir,
        simulator_dir=args.simulator_dir,
        crossbar_size=args.crossbar_size,
        crossbar_count=args.crossbar_count,
        core_count=args.core_count,
        command_timeout_seconds=args.command_timeout_seconds,
        mode=MODE_COMPILE_ONLY,
        verbose=args.verbose,
    )


def ensure_existing_artifacts(network_dir: Path):
    required_paths = [
        network_dir / "runner" / "build" / "runner",
        network_dir / "raptor" / "pim" / "config.json",
        network_dir / "raptor" / "pim" / "memory.bin",
    ]
    missing = [str(path) for path in required_paths if not path.exists()]
    if missing:
        raise FileNotFoundError(
            "Missing compiled local artifacts. Re-run without --skip-compile or restore these paths:\n  "
            + "\n  ".join(missing)
        )


def run_local_reference_and_simulator(args, network_dir: Path, network_onnx_path: Path, case_name: str):
    paths = local_case_paths(network_dir, case_name)
    paths["ref_dir"].mkdir(parents=True, exist_ok=True)
    paths["sim_dir"].mkdir(parents=True, exist_ok=True)

    output_descriptors = onnx_io(network_onnx_path)[1]
    if len(output_descriptors) != 1:
        raise RuntimeError(f"Expected one YOLO output tensor, found {len(output_descriptors)}")

    runner_cmd = [
        str(paths["runner"]),
        "--in0-csv-file",
        str(paths["input_csv"]),
        "--in0-shape",
        "1x3x640x640",
        "--save-csv-dir",
        str(paths["ref_dir"]),
    ]
    subprocess.run(runner_cmd, cwd=paths["runner_build"], check=True)

    tensor = np.loadtxt(paths["input_csv"], delimiter=",", dtype=np.float32).reshape(1, 3, 640, 640)
    write_inputs_to_memory_bin(paths["raptor_pim"] / "memory.bin", paths["raptor_pim"] / "config.json", [tensor])

    dump_ranges = build_dump_ranges(paths["raptor_pim"] / "config.json", output_descriptors)
    run_pim_simulator(
        args.simulator_dir,
        paths["raptor_pim"],
        paths["sim_bin"],
        dump_ranges,
        timeout_sec=args.command_timeout_seconds,
    )
    return paths, output_descriptors[0]


def analyze_case(args, network_dir: Path, network_onnx_path: Path, case, work_dir: Path):
    image_path = work_dir / f"{case.name}{Path(case.url).suffix or '.img'}"
    csv_path = work_dir / f"{case.name}.csv"
    annotated_dir = args.annotated_dir
    annotated_dir.mkdir(parents=True, exist_ok=True)
    download_image(case.url, image_path)
    tensor = letterbox_rgb(Image.open(image_path))
    save_tensor_csv(tensor, csv_path)

    paths = local_case_paths(network_dir, case.name)
    paths["input_csv"].parent.mkdir(parents=True, exist_ok=True)
    paths["input_csv"].write_bytes(csv_path.read_bytes())
    paths, output_descriptor = run_local_reference_and_simulator(args, network_dir, network_onnx_path, case.name)

    output_index, output_name, output_dtype_code, output_shape = output_descriptor
    output_dtype = np.dtype(_ONNX_TO_NP[output_dtype_code])
    ref_csv_path = paths["ref_dir"] / f"output{output_index}_{sanitize_output_name(output_name)}.csv"
    ref = np.loadtxt(ref_csv_path, delimiter=",", dtype=output_dtype).reshape(output_shape)
    sim = np.frombuffer(
        paths["sim_bin"].read_bytes(),
        dtype=output_dtype,
        count=int(np.prod(output_shape)),
    ).reshape(output_shape)

    abs_diff = np.abs(sim.astype(np.float64) - ref.astype(np.float64))
    rel_diff = abs_diff / np.maximum(np.abs(ref.astype(np.float64)), 1e-12)

    ref_detections = decode_yolo_output(ref)
    sim_detections = decode_yolo_output(sim)
    ref_labels = top_unique_labels(ref_detections)
    sim_labels = top_unique_labels(sim_detections)
    ref_image_path = annotated_dir / f"{case.name}_reference.png"
    sim_image_path = annotated_dir / f"{case.name}_simulator.png"
    draw_detections(image_path, ref_detections, ref_image_path)
    draw_detections(image_path, sim_detections, sim_image_path)

    return {
        "case": case.name,
        "expected_label": case.expected_label,
        "ref_top_labels": ref_labels,
        "sim_top_labels": sim_labels,
        "top1_match": bool(ref_labels and sim_labels and ref_labels[0] == sim_labels[0]),
        "expected_in_ref": case.expected_label in ref_labels,
        "expected_in_sim": case.expected_label in sim_labels,
        "max_abs_diff": float(abs_diff.max()),
        "mean_abs_diff": float(abs_diff.mean()),
        "max_rel_diff": float(rel_diff.max()),
        "mean_rel_diff": float(rel_diff.mean()),
        "reference_annotated_image": str(ref_image_path),
        "simulator_annotated_image": str(sim_image_path),
        "ref_top_detections": ref_detections[:5],
        "sim_top_detections": sim_detections[:5],
    }


def main():
    defaults = resolve_default_paths()

    parser = argparse.ArgumentParser(description="Validate YOLO detections on real images using local compilation and simulator execution.")
    parser.add_argument("--network-dir", type=Path, default=defaults["network_dir"])
    parser.add_argument("--network-onnx", type=Path, default=None)
    parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"])
    parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"])
    parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"])
    parser.add_argument("--crossbar-size", type=int, default=2048)
    parser.add_argument("--crossbar-count", type=int, default=256)
    parser.add_argument("--core-count", type=int, default=1000)
    parser.add_argument("--command-timeout-seconds", type=float, default=7200.0)
    parser.add_argument("--skip-compile", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument(
        "--annotated-dir",
        type=Path,
        default=defaults["network_dir"] / "real_image_validation" / "annotated",
    )
    args = parser.parse_args()

    args.network_dir = args.network_dir.resolve()
    args.network_onnx = args.network_onnx.resolve() if args.network_onnx else find_network_onnx(args.network_dir)
    args.raptor_path = args.raptor_path.resolve()
    args.onnx_include_dir = args.onnx_include_dir.resolve()
    args.simulator_dir = args.simulator_dir.resolve()
    args.annotated_dir = args.annotated_dir.resolve()

    if not args.skip_compile:
        ensure_local_artifacts(args, args.network_onnx)
    else:
        ensure_existing_artifacts(args.network_dir)

    reports = []
    with tempfile.TemporaryDirectory(prefix="yolo_local_images_") as tmp_dir:
        work_dir = Path(tmp_dir)
        for case in IMAGE_CASES:
            reports.append(analyze_case(args, args.network_dir, args.network_onnx, case, work_dir))

    print(json.dumps({"network_dir": str(args.network_dir), "network_onnx": str(args.network_onnx), "cases": reports}, indent=2))


if __name__ == "__main__":
    main()