Added more kernels and examples

app.py

@@ -6,35 +6,36 @@ from kornia import morphology as morph
 
 import torch
 
-def morphological_operators(filepath, operator): 
+def morphological_operators(filepath, operator, kernel, kernel_size): 
 
-    img: Tensor = K.io.load_image(filepath, K.io.ImageLoadType.RGB32)    
-    img = img[None]
-    
-    device = 'cpu' # 'cuda:0' for GPU
-    kernel = torch.tensor([[0, 1, 0],[1, 1, 1],[0, 1, 0]]).to(device)
-
-    if operator == 'Dilation':
-       opt = morph.dilation(img, kernel)
-    elif operator == 'Erosion':
-       opt = morph.erosion(img, kernel)
-    elif operator == 'Open':
-       opt = morph.opening(img, kernel)
-    elif operator == 'Close':
-       opt = morph.closing(img, kernel)
-    elif operator == 'Gradient':
-       opt = 1. - morph.gradient(img, kernel)
-    elif operator == 'Bottom Hat':
-       opt = 1. - morph.bottom_hat(img, kernel)
-    else:
-       opt = 1. - morph.top_hat(img, kernel)
-       
-    output = K.tensor_to_image(opt.squeeze(0))
-    return output
+   img: Tensor = K.io.load_image(filepath, K.io.ImageLoadType.RGB32)    
+   img = img[None]
+
+   device = 'cpu' # 'cuda:0' for GPU
+   kernels = {
+      "Ones": torch.ones(kernel_size,kernel_size).to(device),
+      "Eye": torch.eye(kernel_size).to(device)
+      }
+   #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"]
+    ["examples/cat.png", "Dilation", "Ones"],
+    ["examples/huggingface.jpg", "Close", "Eye"]
 ]
 
 title = "Kornia Morphological Operators"
@@ -44,10 +45,22 @@ article = "<p style='text-align: center'><a href='https://kornia.readthedocs.io/
 iface = gr.Interface(morphological_operators, 
             [
             	gr.Image(type="filepath"),
-	            gr.Dropdown(choices=["Dilation", "Erosion", "Open", "Close", "Gradient", "Bottom Hat", "Top Hat"])
+	            gr.Dropdown(choices=["Dilation", "Erosion", "Open", "Close", "Gradient", "Bottom Hat", "Top Hat"]),
+               gr.Radio(choices=["Ones", "Eye"]),
+               gr.Slider(
+                minimum=1,
+                maximum=7,
+                step=2,
+                value=3,
+                label="Kernel size"
+            )
             ],
             "image",
-            examples
+            examples,
+            title=title,
+            description=description,
+            article=article,
+            live=True
             
 )