Update Lowlight.py

Lowlight.py

@@ -1,46 +1,47 @@
+import os
 import torch
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+import time
 import torchvision
-import onnxruntime
-import onnx
 import cv2
+import torchvision.utils as tvu
+import torch.functional as F
 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()
+from model.IAT_main import IAT
 
-if not os.path.isdir(config.save_path):
-    os.mkdir(config.save_path)
+def inference_img(img_path,Net):
+    
+    low_image = Image.open(img_path).convert('RGB')
+    enhance_transforms = transforms.Compose([
+    transforms.Resize((256,256)),
+    transforms.ToTensor()
+    ])
 
-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()
+    with torch.no_grad():
+        low_image = enhance_transforms(low_image)
+        low_image = low_image.unsqueeze(0)
+        start = time.time()
+        restored2 = Net(low_image)
+        end = time.time()
 
-providers = ['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)
-    torchvision.utils.save_image(torch.from_numpy(onnx_output[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)
-else:
-    print(f'Input on model:{model_name} succeed')
+    return restored2,end-start
 
+if __name__ == '__main__': 
+    parser=argparse.ArgumentParser()    
+    parser.add_argument('--test_path',type=str,required=True,help='Path to test')
+    parser.add_argument('--save_path',type=str,required=True,help='Path to save')
+    parser.add_argument('--pk_path',type=str,default='model_zoo/underwater.pth',help='Path of the checkpoint')
+    opt = parser.parse_args()
+    if not os.path.isdir(opt.save_path):
+       os.mkdir(opt.save_path)
+    Net = IAT()
+    Net.load_state_dict(torch.load(opt.pk_path, map_location=torch.device('cpu')))
+    Net = Net.eval()
+    image = opt.test_path
+    print(image)
+    restored2,time_num = inference_img(image,Net)
+    torchvision.utils.save_image(restored2,opt.save_path+'output.png')