File size: 4,864 Bytes
525e655 | 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 125 126 127 128 129 130 131 132 133 134 | import argparse
import os
import time
import axengine as axe
import numpy as np
from PIL import Image, ImageDraw
def get_tensor_dtype(tensor_info):
dtype = getattr(tensor_info, 'dtype', None)
if dtype is None:
dtype = getattr(tensor_info, 'type', None)
return str(dtype).lower() if dtype is not None else ''
def format_tensor_stats(name, tensor):
tensor_np = np.asarray(tensor)
return '{}: dtype={}, shape={}, min={:.6f}, max={:.6f}'.format(
name,
tensor_np.dtype,
tensor_np.shape,
float(tensor_np.min()),
float(tensor_np.max()),
)
def preprocess(image_path, height, width, resize, input_dtype):
original_image = Image.open(image_path).convert('RGB')
original_size = original_image.size
if resize:
model_image = original_image.resize((width, height), Image.BILINEAR)
else:
if original_image.size[0] < width or original_image.size[1] < height:
raise ValueError('Image is smaller than axmodel input size: {}'.format(image_path))
model_image = original_image.crop((0, 0, width, height))
image_np = np.asarray(model_image)
if 'float' in input_dtype:
image_np = image_np.astype(np.float32) / 255.0
else:
image_np = image_np.astype(np.uint8)
image_np = image_np.transpose(2, 0, 1)[None, :, :, :]
return original_image, original_size, image_np
def postprocess(enhanced_np, original_size):
enhanced_np = enhanced_np[0].transpose(1, 2, 0)
max_value = float(enhanced_np.max())
min_value = float(enhanced_np.min())
if max_value <= 1.5 and min_value >= -0.5:
enhanced_np = np.clip(enhanced_np, 0.0, 1.0) * 255.0
else:
enhanced_np = np.clip(enhanced_np, 0.0, 255.0)
enhanced_np = enhanced_np.round().astype(np.uint8)
enhanced_image = Image.fromarray(enhanced_np)
return enhanced_image.resize(original_size, Image.BILINEAR)
def add_label(image, text):
label_height = 32
canvas = Image.new('RGB', (image.width, image.height + label_height), color=(0, 0, 0))
canvas.paste(image, (0, label_height))
draw = ImageDraw.Draw(canvas)
draw.text((10, 8), text, fill=(255, 255, 255))
return canvas
def save_compare(original_image, enhanced_image, result_path):
original_labeled = add_label(original_image, 'Original')
enhanced_labeled = add_label(enhanced_image, 'Result')
compare_image = Image.new('RGB', (original_labeled.width + enhanced_labeled.width, original_labeled.height))
compare_image.paste(original_labeled, (0, 0))
compare_image.paste(enhanced_labeled, (original_labeled.width, 0))
result_dir = os.path.dirname(result_path)
if result_dir and not os.path.exists(result_dir):
os.makedirs(result_dir)
compare_image.save(result_path)
def build_axmodel_session(axmodel_path):
return axe.InferenceSession(axmodel_path, providers=['AxEngineExecutionProvider'])
def infer(config):
if not os.path.isfile(config.input):
raise ValueError('Input image does not exist: {}'.format(config.input))
session = build_axmodel_session(config.axmodel)
input_info = session.get_inputs()[0]
output_info = session.get_outputs()[0]
input_name = input_info.name
input_shape = input_info.shape
input_dtype = get_tensor_dtype(input_info)
output_dtype = get_tensor_dtype(output_info)
height = int(input_shape[2]) if config.height <= 0 else config.height
width = int(input_shape[3]) if config.width <= 0 else config.width
original_image, original_size, input_np = preprocess(
config.input,
height,
width,
bool(config.resize),
input_dtype,
)
start = time.time()
axmodel_outputs = session.run(None, {input_name: input_np})
axmodel_enhanced = axmodel_outputs[0]
elapsed = time.time() - start
enhanced_image = postprocess(axmodel_enhanced, original_size)
save_compare(original_image, enhanced_image, config.output)
print('Input image:', config.input)
print('Output image:', config.output)
print('Model input dtype:', input_dtype or 'unknown')
print('Model output dtype:', output_dtype or 'unknown')
print(format_tensor_stats('Prepared input', input_np))
print(format_tensor_stats('Model output', axmodel_enhanced))
print('axmodel time:', elapsed)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--axmodel', type=str, default='zerodcepp_512_sf8.axmodel')
parser.add_argument('--input', type=str, default='101_3_.png')
parser.add_argument('--output', type=str, default='axmodel_res.png')
parser.add_argument('--height', type=int, default=512)
parser.add_argument('--width', type=int, default=512)
parser.add_argument('--resize', type=int, default=1)
config = parser.parse_args()
infer(config)
|