Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -7,8 +7,6 @@ from skimage import img_as_ubyte
 from tqdm import tqdm
 from natsort import natsorted
 from glob import glob
-from utils.image_utils import save_img
-from utils.model_utils import load_checkpoint
 import argparse
 from model_arch.SRMNet_SWFF import SRMNet_SWFF
 from model_arch.SRMNet import SRMNet
@@ -89,6 +87,22 @@ def define_model(args):
 
     return model
 
+def save_img(filepath, img):
+    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+
+
+def load_checkpoint(model, weights):
+    checkpoint = torch.load(weights, map_location=torch.device('cpu'))
+    try:
+        model.load_state_dict(checkpoint["state_dict"])
+    except:
+        state_dict = checkpoint["state_dict"]
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:]  # remove `module.`
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict)
+
 
 if __name__ == '__main__':
     main()