Update app.py

app.py

@@ -16,6 +16,7 @@ import os
 import matplotlib.pyplot as plt
 import gradio as gr
 import io
+from torchvision.transforms.functional import to_pil_image
 
 class mydata(Dataset):
     def __init__(self, LR_path, GT_path, in_memory = True, transform = None):
@@ -556,45 +557,98 @@ config = {
 
 #train(config)
 
-def preprocess_input(image_path):
-    # Load the input image
-    input_image = Image.open(image_path).convert("RGB")
-    input_image = np.array(input_image) / 127.5 - 1.0
+# def preprocess_input(image_path):
+#     # Load the input image
+#     input_image = Image.open(image_path).convert("RGB")
+#     input_image = np.array(input_image) / 127.5 - 1.0
+#     input_image = input_image.transpose(2, 0, 1).astype(np.float32)
+#     return torch.tensor(input_image).unsqueeze(0)
+
+
+def preprocess_input(image):
+    # Convert Gradio image to PIL format
+    image_pil = to_pil_image(image)
+    # Convert image to numpy array and normalize
+    input_image = np.array(image_pil) / 127.5 - 1.0
+    # Transpose dimensions
     input_image = input_image.transpose(2, 0, 1).astype(np.float32)
-    return torch.tensor(input_image).unsqueeze(0)
+    # Convert to PyTorch tensor and add batch dimension
+    input_tensor = torch.tensor(input_image).unsqueeze(0)
+    return input_tensor
+
 
-def test_single_image(generator_path, input_image_path):
+
+    
+# def test_single_image(generator_path, input_image_path):
+#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#     num_gpus = torch.cuda.device_count()
+#     # Load the generator model
+#     generator = Generator()
+    
+#     if num_gpus == 2:
+#         state_dict = torch.load(generator_path)
+#         if 'module.' in list(state_dict.keys())[0]:
+#             state_dict = {k[7:]: v for k, v in state_dict.items()}
+#         # Load the state dictionary to the model
+#         generator.load_state_dict(state_dict)
+#         generator = DataParallel(generator, device_ids= [0,1]).to(device)
+        
+#     else:
+#         generator = generator.to(device)
+#         generator.load_state_dict(torch.load(generator_path))
+#     generator.eval()
+
+#     # Preprocess the input image
+#     input_image = preprocess_input(input_image_path).to(device)
+
+#     # Perform inference
+#     with torch.no_grad():
+#         output = generator(input_image)
+#         output = output[0].cpu().numpy()
+#         output = (output + 1.0) / 2.0
+#         output = output.transpose(1, 2, 0)
+#         plt.imshow(output)
+#         plt.show()
+
+def test_single_image(input_image):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     num_gpus = torch.cuda.device_count()
+    generator_path = 'pre_trained_model_064.pt'
+
     # Load the generator model
     generator = Generator()
-    
+
     if num_gpus == 2:
         state_dict = torch.load(generator_path)
         if 'module.' in list(state_dict.keys())[0]:
             state_dict = {k[7:]: v for k, v in state_dict.items()}
         # Load the state dictionary to the model
         generator.load_state_dict(state_dict)
-        generator = DataParallel(generator, device_ids= [0,1]).to(device)
-        
+        generator = torch.nn.DataParallel(generator, device_ids=[0, 1]).to(device)
     else:
         generator = generator.to(device)
         generator.load_state_dict(torch.load(generator_path))
     generator.eval()
 
     # Preprocess the input image
-    input_image = preprocess_input(input_image_path).to(device)
+    input_tensor = preprocess_input(input_image).to(device)
 
     # Perform inference
     with torch.no_grad():
-        output = generator(input_image)
+        output = generator(input_tensor)
         output = output[0].cpu().numpy()
         output = (output + 1.0) / 2.0
         output = output.transpose(1, 2, 0)
-        plt.imshow(output)
-        plt.show()
+        # Convert output to PIL image format
+        output_image = to_pil_image(output)
+        return output_image
+
+
+
+
+
 # Load the generator model path
-generator_path = '/pre_trained_model_064.pt'
+#generator_path = 'pre_trained_model_064.pt'