import sys
from pathlib import Path


sys.path.insert(0, Path(__file__).parent.as_posix())


import cv2
from onnx_runner import OnnxRunner


def main(args):
    onnx_path = Path(args.onnx)
    input_path = Path(args.input)
    out_dir = Path(args.out_dir)

    assert onnx_path.suffix == ".onnx"

    if input_path.is_file():
        input_images_path = [input_path]
    else:
        input_images_path = sorted(
            [
                p
                for p in input_path.rglob("*")
                if p.suffix.lower() in (".png", ".jpg", ".jpeg")
            ]
        )

    out_dir.mkdir(exist_ok=True, parents=True)
    onnx_runner = OnnxRunner(onnx_path, args.latent)
    for input_img_path in input_images_path:
        input_img_path: Path

        input_bgr = cv2.imread(input_img_path.as_posix(), cv2.IMREAD_COLOR)
        assert input_bgr is not None
        out_bgr = onnx_runner.run(input_bgr)

        out_path = out_dir / f"{input_img_path.stem}.png"
        cv2.imwrite(out_path.as_posix(), out_bgr)
        print(f"saved {out_path}")


if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument("--onnx", type=str, required=True)
    parser.add_argument("--input", type=str, required=True)
    parser.add_argument("--out-dir", type=str, required=True)
    parser.add_argument("--latent", type=str, default=None)

    main(parser.parse_args())