Zero-DCE/python/onnx_infer.py

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