File size: 1,935 Bytes
3135a01 d3d4eba 3135a01 d3d4eba 3135a01 |
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 |
import sys
import pathlib
CURRENT_DIR = pathlib.Path(__file__).parent
sys.path.append(str(CURRENT_DIR))
import onnxruntime
import cv2
import numpy as np
from data.data_tiling import tiling_inference
import argparse
def main(args):
onnx_file_name = args.onnx_path
image_path = args.image_path
output_path = args.output_path
if args.ipu:
providers = ["VitisAIExecutionProvider"]
provider_options = [{"config_file": args.provider_config}]
else:
providers = ['CPUExecutionProvider']
provider_options = None
ort_session = onnxruntime.InferenceSession(onnx_file_name, providers=providers, provider_options=provider_options)
lr = cv2.imread(image_path)[np.newaxis,:,:,:].astype(np.float32)
# Tiled inference
sr = tiling_inference(ort_session, lr, 8, (56, 56))
sr = np.clip(sr, 0, 255)
sr = sr.squeeze().astype(np.uint8)
cv2.imwrite(output_path, sr)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='PAN')
parser.add_argument('--onnx_path',
type=str,
default='PAN_int8.onnx',
help='Path to onnx model')
parser.add_argument('--image_path',
type=str,
default='test_data/test.png',
help='Path to your low resolution input image.')
parser.add_argument('--output_path',
type=str,
default='test_data/sr.png',
help='Path to your upscaled output image.')
parser.add_argument('--provider_config',
type=str,
default="vaip_config.json",
help="Path of the config file for seting provider_options.")
parser.add_argument('--ipu', action='store_true', help='Use Ipu for interence.')
args = parser.parse_args()
main(args)
|