52Hz commited on
Commit
3dd1292
1 Parent(s): 99085b2

Create main_test_CMFNet.py

Browse files
Files changed (1) hide show
  1. main_test_CMFNet.py +88 -0
main_test_CMFNet.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import cv2
3
+ import glob
4
+ import numpy as np
5
+ from collections import OrderedDict
6
+ from skimage import img_as_ubyte
7
+ import os
8
+ import torch
9
+ import requests
10
+ from PIL import Image
11
+ import torchvision.transforms.functional as TF
12
+ import torch.nn.functional as F
13
+ from natsort import natsorted
14
+ from model.CMFNet import CMFNet
15
+
16
+ def main():
17
+ parser = argparse.ArgumentParser(description='Demo Image Deblur')
18
+ parser.add_argument('--input_dir', default='test/', type=str, help='Input images')
19
+ parser.add_argument('--result_dir', default='results/', type=str, help='Directory for results')
20
+ parser.add_argument('--weights',
21
+ default='experiments/pretrained_models/deblur_GoPro_CMFNet.pth', type=str,
22
+ help='Path to weights')
23
+
24
+ args = parser.parse_args()
25
+
26
+ inp_dir = args.input_dir
27
+ out_dir = args.result_dir
28
+
29
+ os.makedirs(out_dir, exist_ok=True)
30
+
31
+ files = natsorted(glob.glob(os.path.join(inp_dir, '*')))
32
+
33
+ if len(files) == 0:
34
+ raise Exception(f"No files found at {inp_dir}")
35
+
36
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
37
+
38
+ # Load corresponding models architecture and weights
39
+ model = CMFNet()
40
+ model = model.to(device)
41
+ model.eval()
42
+ load_checkpoint(model, args.weights)
43
+
44
+
45
+ mul = 8
46
+ for file_ in files:
47
+ img = Image.open(file_).convert('RGB')
48
+ input_ = TF.to_tensor(img).unsqueeze(0).to(device)
49
+
50
+ # Pad the input if not_multiple_of 8
51
+ h, w = input_.shape[2], input_.shape[3]
52
+ H, W = ((h + mul) // mul) * mul, ((w + mul) // mul) * mul
53
+ padh = H - h if h % mul != 0 else 0
54
+ padw = W - w if w % mul != 0 else 0
55
+ input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
56
+
57
+ with torch.no_grad():
58
+ restored = model(input_)
59
+
60
+ restored = torch.clamp(restored, 0, 1)
61
+ restored = restored[:, :, :h, :w]
62
+ restored = restored.permute(0, 2, 3, 1).cpu().detach().numpy()
63
+ restored = img_as_ubyte(restored[0])
64
+
65
+ f = os.path.splitext(os.path.split(file_)[-1])[0]
66
+ save_img((os.path.join(out_dir, f + '.png')), restored)
67
+
68
+
69
+
70
+ def save_img(filepath, img):
71
+ cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
72
+
73
+
74
+ def load_checkpoint(model, weights):
75
+ checkpoint = torch.load(weights, map_location=torch.device('cpu'))
76
+ try:
77
+ model.load_state_dict(checkpoint["state_dict"])
78
+ except:
79
+ state_dict = checkpoint["state_dict"]
80
+ new_state_dict = OrderedDict()
81
+ for k, v in state_dict.items():
82
+ name = k[7:] # remove `module.`
83
+ new_state_dict[name] = v
84
+ model.load_state_dict(new_state_dict)
85
+
86
+
87
+ if __name__ == '__main__':
88
+ main()