File size: 2,733 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 63 64 65 66 67 68 | import argparse
import numpy as np
import axengine as axe
from PIL import Image, ImageDraw, ImageFont
def draw_label(img, text):
draw = ImageDraw.Draw(img)
try:
font = ImageFont.truetype("DejaVuSans-Bold.ttf", max(18, img.height // 36))
except Exception:
font = ImageFont.load_default()
padding = 6
bbox = draw.textbbox((0, 0), text, font=font)
box_w = bbox[2] - bbox[0] + padding * 2
box_h = bbox[3] - bbox[1] + padding * 2
draw.rectangle([0, 0, box_w, box_h], fill=(0, 0, 0))
draw.text((padding, padding), text, fill=(255, 255, 255), font=font)
def main():
parser = argparse.ArgumentParser(description="axmodel single-image dehazing inference.")
parser.add_argument("--axmodel", default="aodnet_1x3x480x640_sim.axmodel", help="axmodel model path")
parser.add_argument("--input_image", default='./pic/canyon2.jpg', help="hazy image path")
parser.add_argument("--output", default="axmodel_compare.png", help="output comparison path (hazy | dehazed)")
parser.add_argument("--height", type=int, default=480, help="resize height (match axmodel input)")
parser.add_argument("--width", type=int, default=640, help="resize width (match axmodel input)")
parser.add_argument("--normalize", action="store_true", help="use (x-0.5)/0.5 normalization")
parser.add_argument("--no_label", action="store_true", help="do not draw hazy/dehazed labels")
args = parser.parse_args()
img = Image.open(args.input_image).convert("RGB")
orig_size = img.size
img_resized = img.resize((args.width, args.height), Image.BILINEAR)
arr = np.asarray(img_resized).astype(np.float32)
arr = arr.transpose(2, 0, 1)[None, :, :, :].astype(np.uint8)
sess = axe.InferenceSession(args.axmodel, providers=["AxEngineExecutionProvider"])
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))
# 还原到原图大小
if out_img.size != orig_size:
out_img = out_img.resize(orig_size, Image.BICUBIC)
hazy_labeled = img.copy()
dehazed_labeled = out_img.copy()
if not args.no_label:
draw_label(hazy_labeled, "hazy")
draw_label(dehazed_labeled, "dehazed")
compare = Image.new("RGB", (img.width + out_img.width, max(img.height, out_img.height)))
compare.paste(hazy_labeled, (0, 0))
compare.paste(dehazed_labeled, (img.width, 0))
compare.save(args.output)
print("Saved:", args.output)
print("Input: {:.0f}x{:.0f} -> axmodel: {}x{}".format(*orig_size, args.height, args.width))
if __name__ == "__main__":
main()
|