File size: 4,877 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import argparse
import os
import numpy as np
import axengine as axe
from PIL import Image, ImageDraw, ImageFont
HEIGHT = 512
WIDTH = 512
IMG_EXTENSIONS = ('.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP')
def parse_args():
parser = argparse.ArgumentParser(description='Run GCANet axmodel inference with fixed 512x512 input.')
parser.add_argument('--task', default='dehaze', choices=['dehaze', 'derain'])
parser.add_argument('--axmodel', default='GCANet_u16.axmodel', help='Path to axmodel model.')
parser.add_argument('--indir', default='examples')
parser.add_argument('--outdir', default='axmodel_output')
parser.add_argument(
'--input-mode',
default='raw_u8',
choices=['compat_centered_u8', 'raw_u8'],
help='compat_centered_u8 keeps the historical centered input encoding for the current axmodel; raw_u8 follows the declared U8 input processor literally.',
)
return parser.parse_args()
def make_dataset(image_dir):
images = []
assert os.path.isdir(image_dir), '%s is not a valid directory' % image_dir
for root, _, fnames in sorted(os.walk(image_dir)):
for fname in fnames:
if fname.endswith(IMG_EXTENSIONS):
images.append(os.path.join(root, fname))
return images
def edge_compute_np(img_chw):
x_diffx = np.abs(img_chw[:, :, 1:] - img_chw[:, :, :-1])
x_diffy = np.abs(img_chw[:, 1:, :] - img_chw[:, :-1, :])
edge = np.zeros_like(img_chw, dtype=np.float32)
edge[:, :, 1:] += x_diffx
edge[:, :, :-1] += x_diffx
edge[:, 1:, :] += x_diffy
edge[:, :-1, :] += x_diffy
edge = np.sum(edge, axis=0, keepdims=True) / 3.0
edge = edge / 4.0
return edge.astype(np.float32)
def preprocess(img_path, input_mode):
img = Image.open(img_path).convert('RGB')
img = img.resize((WIDTH, HEIGHT), Image.BICUBIC)
img_rgb = np.array(img).astype(np.float32)
img_rgb_chw = np.transpose(img_rgb, (2, 0, 1))
edge = edge_compute_np(img_rgb_chw)
# This 4-channel model is calibrated from RGB+edge NCHW tensors.
# The build config marks src_format=BGR/tensor_format=RGB, but with
# csc_mode=NoCSC there is no runtime color conversion, so feeding BGR here
# would silently swap channels and introduce artifacts.
model_input = np.concatenate((img_rgb_chw, edge), axis=0)[None, :, :, :]
if input_mode == 'compat_centered_u8':
# The original float model consumes (rgb+edge-128). The current axmodel
# was compiled as U8 input without explicit dequant params, so we keep
# the historical wraparound encoding here for compatibility.
model_input = np.mod(np.round(model_input - 128.0), 256.0).astype(np.uint8)
else:
model_input = np.clip(np.round(model_input), 0, 255).astype(np.uint8)
return img_rgb_chw, model_input
def postprocess(pred, img_chw, only_residual):
out = pred[0].astype(np.float32)
if only_residual:
out = out + img_chw
out = np.round(out).clip(0, 255).astype(np.uint8)
out = np.transpose(out, (1, 2, 0))
return out
def make_comparison_image(input_chw, output_hwc, task):
input_hwc = np.transpose(input_chw, (1, 2, 0)).clip(0, 255).astype(np.uint8)
input_img = Image.fromarray(input_hwc, mode='RGB')
output_img = Image.fromarray(output_hwc, mode='RGB')
title_height = 32
comparison = Image.new('RGB', (WIDTH * 2, HEIGHT + title_height), color='white')
comparison.paste(input_img, (0, title_height))
comparison.paste(output_img, (WIDTH, title_height))
draw = ImageDraw.Draw(comparison)
font = ImageFont.load_default()
right_label = 'dehaze' if task == 'dehaze' else task
draw.text((8, 8), 'hazy', fill='black', font=font)
draw.text((WIDTH + 8, 8), right_label, fill='black', font=font)
draw.line((WIDTH, 0, WIDTH, HEIGHT + title_height), fill='black', width=1)
return comparison
def main():
args = parse_args()
axmodel_path = args.axmodel or os.path.join('onnx', 'gcanet_%s_512x512_sim.axmodel' % args.task)
only_residual = args.task == 'dehaze'
os.makedirs(args.outdir, exist_ok=True)
session = axe.InferenceSession(axmodel_path, providers=['AxEngineExecutionProvider'])
input_name = session.get_inputs()[0].name
for img_path in make_dataset(args.indir):
img_chw, model_input = preprocess(img_path, args.input_mode)
pred = session.run(None, {input_name: model_input})[0]
out_img = postprocess(pred, img_chw, only_residual)
comparison = make_comparison_image(img_chw, out_img, args.task)
save_name = os.path.splitext(os.path.basename(img_path))[0] + '_%s_compare.png' % args.task
comparison.save(os.path.join(args.outdir, save_name))
print('Saved:', os.path.join(args.outdir, save_name))
if __name__ == '__main__':
main()
|