Update app.py

app.py

@@ -1,38 +1,31 @@
 import gradio as gr
-
 import kornia as K
 from kornia.core import Tensor
 from kornia import morphology as morph
-
 import torch
 
-def morphological_operators(filepath, operator, kernel, kernel_size): 
-
-   img: Tensor = K.io.load_image(filepath, K.io.ImageLoadType.RGB32)    
-   img = img[None]
+def morphological_operators(img, operator, kernel, kernel_size): 
+    img: Tensor = K.io.image_to_tensor(img)
+    img = img.unsqueeze(0)  # Add batch dimension
+    device = 'cpu'  # 'cuda:0' for GPU
+    kernels = {
+        "Ones": torch.ones(kernel_size, kernel_size).to(device),
+        "Eye": torch.eye(kernel_size).to(device),
+        "Cross": torch.tensor([[0, 1, 0],[1, 1, 1],[0, 1, 0]]).to(device),
+    }
+    operations = {
+        'Dilation': morph.dilation(img, kernels[kernel]),
+        'Erosion': morph.erosion(img, kernels[kernel]),
+        'Open': morph.opening(img, kernels[kernel]),
+        'Close': morph.closing(img, kernels[kernel]),
+        'Gradient': 1. - morph.gradient(img, kernels[kernel]),
+        'Bottom Hat': 1. - morph.bottom_hat(img, kernels[kernel]),
+        'Top Hat': 1. - morph.top_hat(img, kernels[kernel])
+    }
+        
+    output = K.tensor_to_image(operations[operator].squeeze(0))
+    return output
 
-   device = 'cpu' # 'cuda:0' for GPU
-   kernels = {
-      "Ones": torch.ones(kernel_size,kernel_size).to(device),
-      "Eye": torch.eye(kernel_size).to(device),
-      "Cross": torch.tensor([[0, 1, 0],[1, 1, 1],[0, 1, 0]]).to(device),
-      }
-
-   operations = {
-      'Dilation': morph.dilation(img, kernels[kernel]),
-      'Erosion':morph.erosion(img, kernels[kernel]),
-      'Open': morph.opening(img, kernels[kernel]),
-      'Close': morph.closing(img, kernels[kernel]),
-      'Gradient': 1. - morph.gradient(img, kernels[kernel]),
-      'Bottom Hat': 1. - morph.bottom_hat(img, kernels[kernel]),
-      'Top Hat': 1. - morph.top_hat(img, kernels[kernel])
-   }
-
-      
-   output = K.tensor_to_image(operations[operator].squeeze(0))
-   return output
-    
-    
 examples = [
     ["examples/cat.png", "Dilation", "Ones", 3],
     ["examples/huggingface.jpg", "Close", "Eye", 5]
@@ -42,26 +35,47 @@ title = "Kornia Morphological Operators"
 description = "<p style='text-align: center'>This is a Gradio demo for Kornia's Morphological Operators.</p><p style='text-align: center'>To use it, simply upload your image, or click one of the examples to load them, and select any morphological operator to run it! Read more at the links at the bottom.</p>"
 article = "<p style='text-align: center'><a href='https://kornia.readthedocs.io/en/latest/' target='_blank'>Kornia Docs</a> | <a href='https://github.com/kornia/kornia' target='_blank'>Kornia Github Repo</a> | <a href='https://kornia-tutorials.readthedocs.io/en/latest/morphology_101.html' target='_blank'>Kornia Morphological Operators Tutorial</a></p>"
 
-iface = gr.Interface(morphological_operators, 
-            [
-            	gr.Image(type="filepath"),
-	            gr.Dropdown(choices=["Dilation", "Erosion", "Open", "Close", "Gradient", "Bottom Hat", "Top Hat"]),
-               gr.Radio(choices=["Ones", "Eye", "Cross"]),
-               gr.Slider(
-                minimum=1,
-                maximum=7,
-                step=2,
-                value=3,
-                label="Kernel size"
-            )
-            ],
-            "image",
-            examples,
-            title=title,
-            description=description,
-            article=article,
-            live=True
-            
-)
+with gr.Blocks(title=title) as demo:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+    
+    with gr.Row():
+        input_image = gr.Image(label="Input Image")
+        output_image = gr.Image(label="Output Image")
+    
+    with gr.Row():
+        operator = gr.Dropdown(
+            choices=["Dilation", "Erosion", "Open", "Close", "Gradient", "Bottom Hat", "Top Hat"],
+            label="Operator"
+        )
+        kernel = gr.Radio(
+            choices=["Ones", "Eye", "Cross"],
+            label="Kernel Type"
+        )
+        kernel_size = gr.Slider(
+            minimum=1,
+            maximum=7,
+            step=2,
+            value=3,
+            label="Kernel size"
+        )
+    
+    apply_button = gr.Button("Apply Operator")
+    apply_button.click(
+        fn=morphological_operators,
+        inputs=[input_image, operator, kernel, kernel_size],
+        outputs=output_image
+    )
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[input_image, operator, kernel, kernel_size],
+        outputs=output_image,
+        fn=morphological_operators,
+        cache_examples=True
+    )
+    
+    gr.Markdown(article)
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()