File size: 2,478 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import argparse
import os
import numpy as np
import axengine as axe
from PIL import Image
def preprocess(image_path, height=None, width=None):
image = Image.open(image_path).convert("RGB")
if height is not None and width is not None:
image = image.resize((width, height), Image.BICUBIC)
# 不做均值归一化，仅转 NCHW/RGB，保存为 uint8
image_np = np.asarray(image).astype(np.uint8)
image_np = np.transpose(image_np, (2, 0, 1))[np.newaxis, ...]
return image_np, image
def postprocess(output):
output = np.squeeze(output, axis=0)
output = np.transpose(output, (1, 2, 0))
output = np.clip(output, 0.0, 1.0)
return Image.fromarray((output * 255.0).round().astype(np.uint8))
def infer_image(args):
session = axe.InferenceSession(args.axmodel, providers=["AxEngineExecutionProvider"])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
input_tensor, hazy_image = preprocess(args.image, args.height, args.width)
output = session.run([output_name], {input_name: input_tensor})[0]
dehaze_image = postprocess(output)
output_dir = os.path.dirname(args.output)
if output_dir:
os.makedirs(output_dir, exist_ok=True)
compare_image = Image.new("RGB", (hazy_image.width + dehaze_image.width, hazy_image.height))
compare_image.paste(hazy_image, (0, 0))
compare_image.paste(dehaze_image, (hazy_image.width, 0))
compare_image.save(args.output)
print("Comparison image saved to:", args.output)
def parse_args():
parser = argparse.ArgumentParser(description="Run Light-DehazeNet axmodel inference on one image.")
parser.add_argument("-m", "--axmodel", default="./LightDehazeNet.axmodel", help="path to axmodel model")
parser.add_argument("-i", "--image", default='query_hazy_images/outdoor_natural/nh(5).png', help="path to input hazy image")
parser.add_argument("-o", "--output", default="axmodel_dehaze.jpg", help="path to save side-by-side comparison image")
parser.add_argument("--height", type=int, default=640, help="resize input to this height before inference")
parser.add_argument("--width", type=int, default=480, help="resize input to this width before inference")
args = parser.parse_args()
if (args.height is None) != (args.width is None):
parser.error("--height and --width must be specified together")
return args
if __name__ == "__main__":
infer_image(parse_args())
|