Update Lowlight.py

Lowlight.py

@@ -1,4 +1,5 @@
 import torch
+import torchvision
 import onnxruntime
 import onnx
 import cv2
@@ -32,9 +33,10 @@ try:
     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)
+    torchvision.utils.save_image(torch.from_numpy(onnx_ouput[0]), config.save_path+'output.png')
+    #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)