Updated loading from RCAN

models/RCAN/rcan.py

@@ -1,4 +1,8 @@
+import os
+import torch
 from torch import nn
+from PIL import Image
+from torchvision.transforms import ToTensor
 
 NUM_RESIDUAL_GROUPS = 8
 NUM_RESIDUAL_BLOCKS = 16
@@ -103,3 +107,16 @@ class RCAN(nn.Module):
         reconstructed_image = self.reconstruction_conv(upscaled_image) # Reconstruction
 
         return reconstructed_image
+
+if __name__ == '__main__':
+    current_dir = os.path.dirname(os.path.realpath(__file__))
+
+    model = RCAN()
+    model.load_state_dict(torch.load(current_dir + '/rcan_checkpoint.pth', map_location=torch.device('cpu')))
+    model.eval()
+    with torch.no_grad():
+        input_image = Image.open('images/demo.png')
+        input_tensor = ToTensor()(input_image).unsqueeze(0)
+        output_tensor = model(input_tensor)
+        print(input_tensor.shape)
+        print(output_tensor.shape)