Update app.py

app.py

@@ -103,7 +103,7 @@ class ImageTransform:
                 transforms.Resize((img_size, img_size)),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=[0.5], std=[0.5])
-           })}
+           ])}
 ​
    def __call__(self, img, phase='train'):
        img = self.transform[phase](img)
@@ -114,6 +114,44 @@ class ImageTransform:
 path = hf_hub_download('huggan/NeonGAN', 'model.bin')
 model_gen_n = torch.load(path, map_location=torch.device('cpu'))
 
+transform = ImageTransform(img_size=256)
+
+inputs = [
+    gr.inputs.Image(type="pil", label="Original Image")
+]
+
+outputs = [
+    gr.outputs.Image(type="pil", label="Neon Image")
+]
+
+def get_output_image(img):
+
+    img = transform(img, phase='test')
+    gen_img = model_gen_n(img.unsqueeze(0))[0]
+
+    # Reverse Normalization
+    gen_img = gen_img * 0.5 + 0.5
+    gen_img = gen_img * 255
+    gen_img = gen_img.detach().cpu().numpy().astype(np.uint8)
+
+    gen_img = np.transpose(gen_img, [1,2,0])
+
+    gen_img = Image.fromarray(gen_img)
+    
+    return gen_img
+    
+gr.Interface(
+    get_output_image,
+    inputs,
+    outputs,
+    examples = examples,
+    title=title,
+    description=description,
+    theme="huggingface",
+).launch(enable_queue=True)
+    
+    
+