Update predict.py

predict.py

@@ -8,7 +8,7 @@ import glob
 import torch
 from skimage import img_as_ubyte
 from PIL import Image
-from model.SRMNet import SRMNet
+from model.CMFNet import CMFNet
 from main_test_SRMNet import save_img, setup
 import torchvision.transforms.functional as TF
 import torch.nn.functional as F
@@ -16,7 +16,7 @@ import torch.nn.functional as F
 
 class Predictor(cog.Predictor):
     def setup(self):
-        model_dir = 'experiments/pretrained_models/AWGN_denoising_SRMNet.pth'
+        model_dir = 'experiments/pretrained_models/deraindrop_model.pth'
 
         parser = argparse.ArgumentParser(description='Demo Image Denoising')
         parser.add_argument('--input_dir', default='./test/', type=str, help='Input images')
@@ -38,7 +38,7 @@ class Predictor(cog.Predictor):
         shutil.copy(str(image), input_path)
 
         # Load corresponding models architecture and weights
-        model = SRMNet()
+        model = CMFNet()
         model.eval()
         model = model.to(self.device)
 
@@ -46,7 +46,7 @@ class Predictor(cog.Predictor):
         os.makedirs(save_dir, exist_ok=True)
 
         out_path = Path(tempfile.mkdtemp()) / "out.png"
-        mul = 16
+        mul = 8
         for file_ in sorted(glob.glob(os.path.join(folder, '*.PNG'))):
             img = Image.open(file_).convert('RGB')
             input_ = TF.to_tensor(img).unsqueeze(0).cuda()
@@ -60,7 +60,7 @@ class Predictor(cog.Predictor):
             with torch.no_grad():
                 restored = model(input_)
 
-            restored = torch.clamp(restored, 0, 1)
+            restored = torch.clamp(restored[0], 0, 1)
             restored = restored[:, :, :h, :w]
             restored = restored.permute(0, 2, 3, 1).cpu().detach().numpy()
             restored = img_as_ubyte(restored[0])