Spaces:
Build error
Build error
File size: 1,517 Bytes
1951449 c4d8d8b 1951449 |
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 |
import torch
import onnxruntime
import onnx
import cv2
import argparse
import warnings
import numpy as np
import matplotlib.pyplot as plt
import os
parser = argparse.ArgumentParser()
parser.add_argument('--test_path', type=str, default='/home/arye-stark/zwb/Illumination-Adaptive-Transformer/IAT_enhance/demo_imgs/low_demo.jpg')
parser.add_argument('--pk_path', type=str, default='model_zoo/Low.onnx')
parser.add_argument('--save_path', type=str, default='Results/')
config = parser.parse_args()
if not os.path.isdir(config.save_path):
os.mkdir(config.save_path)
img = plt.imread(config.test_path)
input_image = np.asarray(img) / 255.0
input_image = torch.from_numpy(input_image).float()
input_image = input_image.permute(2, 0, 1).unsqueeze(0)
input_image = input_image.numpy()
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
model_name = 'IAT'
print('-' * 50)
try:
onnx_session = onnxruntime.InferenceSession(config.pk_path, providers=providers)
onnx_input = {'input': input_image}
#onnx_output0, onnx_output1, onnx_output2 = onnx_session.run(['output0', 'output1', 'output2'], onnx_input)
onnx_output = onnx_session.run(['output'], onnx_input)
torch_output = np.squeeze(onnx_output[0], 0)
torch_output = np.transpose(torch_output * 255, [1, 2, 0]).astype(np.uint8)
plt.imsave(config.save_path+'output.png', torch_output)
except Exception as e:
print(f'Input on model:{model_name} failed')
print(e)
else:
print(f'Input on model:{model_name} succeed')
|