Added inference method for RCAN

models/RCAN/rcan.py

@@ -107,6 +107,15 @@ class RCAN(nn.Module):
         reconstructed_image = self.reconstruction_conv(upscaled_image) # Reconstruction
 
         return reconstructed_image
+    
+    def inference(self, x):
+        """
+        x is a PIL image
+        """
+        x = ToTensor()(x).unsqueeze(0)
+        x = self.forward(x)
+        x = Image.fromarray((x.squeeze(0).permute(1, 2, 0).detach().numpy() * 255).astype('uint8'))
+        return x
 
 if __name__ == '__main__':
     current_dir = os.path.dirname(os.path.realpath(__file__))
@@ -116,7 +125,4 @@ if __name__ == '__main__':
     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)        
+        output_image = model.inference(input_image)