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]

   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]
]

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
            
)

iface.launch()