Added clamping to delete noise in output

models/RCAN/rcan.py

@@ -106,15 +106,17 @@ class RCAN(nn.Module):
         upscaled_image = self.upscaling_module(deep_feature) # Upscaling module
         reconstructed_image = self.reconstruction_conv(upscaled_image) # Reconstruction
 
-        return reconstructed_image
+        return reconstructed_image.clamp(0, 1)
     
     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'))
+        self.eval()
+        with torch.no_grad():
+            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__':

models/SRGAN/srgan.py

@@ -51,15 +51,17 @@ class GeneratorResnet(nn.Module):
         out = torch.add(out1, out2)
         out = self.upsampling(out)
         out = self.conv3(out)
-        return out
+        return out.clamp(0, 1)
     
     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'))
+        self.eval()
+        with torch.no_grad():
+            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__':