Spaces:
Sleeping
Sleeping
Update main_test_CMFNet.py
Browse files- main_test_CMFNet.py +7 -6
main_test_CMFNet.py
CHANGED
@@ -55,14 +55,15 @@ def main():
|
|
55 |
input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
|
56 |
with torch.no_grad():
|
57 |
restored = model(input_)
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
|
|
63 |
|
64 |
f = os.path.splitext(os.path.split(file_)[-1])[0]
|
65 |
-
save_img((os.path.join(out_dir, f + '.png')),
|
66 |
|
67 |
|
68 |
def save_img(filepath, img):
|
|
|
55 |
input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
|
56 |
with torch.no_grad():
|
57 |
restored = model(input_)
|
58 |
+
|
59 |
+
restored_ = restored[0]
|
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 |
def save_img(filepath, img):
|