Upload 3 files

Dehazing.py

@@ -0,0 +1,45 @@
+import os
+import torch
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+import time
+import torchvision
+import cv2
+import torchvision.utils as tvu
+import torch.functional as F
+import argparse
+
+def inference_img(haze_path,Net):
+    
+    haze_image = Image.open(haze_path).convert('RGB')
+    enhance_transforms = transforms.Compose([
+    transforms.Resize((400,400)),
+    transforms.ToTensor()
+    ])
+
+    print(haze_image.size)
+    with torch.no_grad():
+        haze_image = enhance_transforms(haze_image)
+        #print(haze_image)
+        haze_image = haze_image.unsqueeze(0)
+        start = time.time()
+        restored2 = Net(haze_image)
+        end = time.time()
+
+
+    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/Haze4k.tjm',help='Path of the checkpoint')
+    opt = parser.parse_args()
+    if not os.path.isdir(opt.save_path):
+        os.mkdir(opt.save_path)
+    Net=torch.jit.load(opt.pk_path,map_location=torch.device('cpu')).eval()
+    image = opt.test_path
+    print(image)
+    restored2,time_num = inference_img(image,Net)
+    torchvision.utils.save_image(restored2,opt.save_path+os.path.split(image)[-1])

Lowlight.py

@@ -0,0 +1,44 @@
+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+os.path.split(config.test_path)[-1], 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')
+
+

SuperResolution.py

@@ -0,0 +1,47 @@
+import os
+import torch
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+import time
+import torchvision
+import argparse
+from models.SCET import SCET
+
+def inference_img(img_path,Net,device):
+    
+    low_image = Image.open(img_path).convert('RGB')
+    enhance_transforms = transforms.Compose([
+    transforms.ToTensor()
+    ])
+
+    with torch.no_grad():
+        low_image = enhance_transforms(low_image)
+        low_image = low_image.unsqueeze(0)
+        start = time.time()
+        restored2 = Net(low_image.to(device))
+        end = time.time()
+
+
+    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/SRx4.pth',help='Path of the checkpoint')
+    parser.add_argument('--scale',type=int,default=4,help='scale factor')
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    opt = parser.parse_args()
+    if not os.path.isdir(opt.save_path):
+        os.mkdir(opt.save_path)
+    if opt.scale == 3:
+        Net = SCET(63, 128, opt.scale).eval()
+    else:
+        Net = SCET(64, 128, opt.scale).eval()
+    Net.load_state_dict(torch.load(opt.pk_path))
+    Net=Net.to(device)
+    image=opt.test_path
+    print(image)
+    restored2,time_num=inference_img(image,Net,device)
+    torchvision.utils.save_image(restored2,opt.save_path+os.path.split(image)[-1])