File size: 1,770 Bytes
4dca198 | 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 | import argparse
import numpy as np
import onnxruntime as ort
from PIL import Image
def main():
parser = argparse.ArgumentParser(description="ONNX single-image dehazing inference.")
parser.add_argument("--onnx", default="aodnet_1x3x480x640_sim.onnx", help="ONNX model path")
parser.add_argument("--input_image", required=True, help="hazy image path")
parser.add_argument("--output", default="onnx_result.png", help="output path")
parser.add_argument("--height", type=int, default=480, help="resize height (match ONNX input)")
parser.add_argument("--width", type=int, default=640, help="resize width (match ONNX input)")
parser.add_argument("--normalize", action="store_true", help="use (x-0.5)/0.5 normalization")
args = parser.parse_args()
img = Image.open(args.input_image).convert("RGB")
orig_size = img.size
img = img.resize((args.width, args.height), Image.BILINEAR)
arr = np.asarray(img).astype(np.float32) / 255.0
if args.normalize:
arr = (arr - 0.5) / 0.5
arr = arr.transpose(2, 0, 1)[None, :, :, :].astype(np.float32)
sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
out = sess.run(None, {input_name: arr})[0]
out = out[0].transpose(1, 2, 0)
out = np.clip(out, 0.0, 1.0)
out_img = Image.fromarray((out * 255.0).round().astype(np.uint8))
compare = Image.new("RGB", (img.width + out_img.width, max(img.height, out_img.height)))
compare.paste(img, (0, 0))
compare.paste(out_img, (img.width, 0))
compare.save(args.output)
print("Saved:", args.output)
print("Input: {:.0f}x{:.0f} -> ONNX: {}x{}".format(*orig_size, args.height, args.width))
if __name__ == "__main__":
main()
|