File size: 2,881 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 | import argparse
import os
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw
LABEL_HEIGHT = 32
def parse_args():
parser = argparse.ArgumentParser(description='Run Zero-DCE ONNX inference.')
parser.add_argument('--image', type=str, default='data/test_data/DICM/25.jpg')
parser.add_argument('--onnx', type=str, default='zero_dce_1x3x256x256_sim.onnx')
parser.add_argument('--output', type=str, default='onnx_result.png')
parser.add_argument('--height', type=int, default=256)
parser.add_argument('--width', type=int, default=256)
return parser.parse_args()
def load_image(image_path, height, width):
org_image = Image.open(image_path).convert('RGB')
input_image = org_image.resize((width, height), Image.BILINEAR)
image_array = np.asarray(input_image).astype(np.float32) / 255.0
input_array = np.transpose(image_array, (2, 0, 1))
input_array = np.expand_dims(input_array, axis=0)
return org_image, input_array
def run_onnx(input_array, onnx_path):
session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
output = session.run([output_name], {input_name: input_array.astype(np.float32)})[0]
return output
def output_to_pil(output_array, org_size):
output = np.squeeze(output_array, axis=0)
output = np.transpose(output, (1, 2, 0))
output = np.clip(output, 0.0, 1.0)
output = (output * 255.0).astype(np.uint8)
res_image = Image.fromarray(output)
return res_image.resize(org_size, Image.BILINEAR)
def add_label(image, text):
canvas = Image.new('RGB', (image.width, image.height + LABEL_HEIGHT), color=(255, 255, 255))
canvas.paste(image, (0, LABEL_HEIGHT))
draw = ImageDraw.Draw(canvas)
draw.text((10, 8), text, fill=(255, 0, 0))
return canvas
def save_compare_image(org_image, res_image, output_path):
org_labeled = add_label(org_image, 'org')
res_labeled = add_label(res_image, 'res')
compare = Image.new('RGB', (org_labeled.width + res_labeled.width, org_labeled.height), color=(255, 255, 255))
compare.paste(org_labeled, (0, 0))
compare.paste(res_labeled, (org_labeled.width, 0))
output_dir = os.path.dirname(output_path)
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
compare.save(output_path)
def main():
args = parse_args()
org_image, input_array = load_image(args.image, args.height, args.width)
onnx_output = run_onnx(input_array, args.onnx)
res_image = output_to_pil(onnx_output, org_image.size)
print('image:', args.image)
print('onnx:', args.onnx)
print('input shape:', input_array.shape)
print('onnx output shape:', onnx_output.shape)
print('org size:', org_image.size)
save_compare_image(org_image, res_image, args.output)
print('Saved result image:', args.output)
print('Comparison layout: org | res')
if __name__ == '__main__':
main()
|