Spaces:
Running
Running
Update main_test_CMFNet.py
Browse files- main_test_CMFNet.py +5 -6
main_test_CMFNet.py
CHANGED
@@ -55,12 +55,11 @@ 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 |
-
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, 'out.png')), restored_)
|
|
|
55 |
input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
|
56 |
with torch.no_grad():
|
57 |
restored = model(input_)
|
58 |
+
restored_ = restored[0]
|
59 |
+
restored_= torch.clamp(restored_, 0, 1)
|
60 |
+
restored_= restored_[:, :, :h, :w]
|
61 |
+
restored_= restored_.permute(0, 2, 3, 1).cpu().detach().numpy()
|
62 |
+
restored_= img_as_ubyte(restored_[0])
|
|
|
63 |
|
64 |
f = os.path.splitext(os.path.split(file_)[-1])[0]
|
65 |
save_img((os.path.join(out_dir, 'out.png')), restored_)
|