Upload 4 files

SuperResolution.py

@@ -8,7 +8,7 @@ import torchvision
 import argparse
 from models.SCET import SCET
 
-def inference_img(img_path,Net,device):
+def inference_img(img_path,Net):
     
     low_image = Image.open(img_path).convert('RGB')
     enhance_transforms = transforms.Compose([
@@ -19,7 +19,7 @@ def inference_img(img_path,Net,device):
         low_image = enhance_transforms(low_image)
         low_image = low_image.unsqueeze(0)
         start = time.time()
-        restored2 = Net(low_image.to(device))
+        restored2 = Net(low_image)
         end = time.time()
 
 
@@ -31,17 +31,16 @@ if __name__ == '__main__':
     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()
+        Net = SCET(63, 128, opt.scale)
     else:
-        Net = SCET(64, 128, opt.scale).eval()
-    Net.load_state_dict(torch.load(opt.pk_path))
-    Net=Net.to(device)
+        Net = SCET(64, 128, opt.scale)
+    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,device)
+    restored2,time_num=inference_img(image,Net)
     torchvision.utils.save_image(restored2,opt.save_path+'output.png')

Underwater.py

@@ -11,7 +11,7 @@ import torch.functional as F
 import argparse
 from net.Ushape_Trans import *
 
-def inference_img(img_path,Net,device):
+def inference_img(img_path,Net):
     
     low_image = Image.open(img_path).convert('RGB')
     enhance_transforms = transforms.Compose([
@@ -23,7 +23,7 @@ def inference_img(img_path,Net,device):
         low_image = enhance_transforms(low_image)
         low_image = low_image.unsqueeze(0)
         start = time.time()
-        restored2 = Net(low_image.to(device))
+        restored2 = Net(low_image)
         end = time.time()
 
 
@@ -37,11 +37,10 @@ if __name__ == '__main__':
     opt = parser.parse_args()
     if not os.path.isdir(opt.save_path):
        os.mkdir(opt.save_path)
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    Net = Generator().eval()
-    Net.load_state_dict(torch.load(opt.pk_path))
-    Net = Net.to(device)
+    Net = Generator()
+    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,device)
+    restored2,time_num = inference_img(image,Net)
     torchvision.utils.save_image(restored2,opt.save_path+'output.png')