| 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) |
| |
| 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()) |
|
|
