File size: 1,431 Bytes
d1faacc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | 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())
|