import gradio as gr

import kornia as K
from kornia.core import Tensor

def edge_detection(filepath, detector): 

    img: Tensor = K.io.load_image(filepath, K.io.ImageLoadType.RGB32)    
    img = img[None]

    x_gray = K.color.rgb_to_grayscale(img)
    

    if detector == '1st order derivates in x':
        grads: Tensor = K.filters.spatial_gradient(x_gray, order=1) 
        grads_x = grads[:, :, 0]
        grads_y = grads[:, :, 1]
        
        output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
    
    elif detector == '1st order derivates in y':
        grads: Tensor = K.filters.spatial_gradient(x_gray, order=1)
        grads_x = grads[:, :, 0]
        grads_y = grads[:, :, 1]
        
        output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
        
    elif detector == '2nd order derivatives in x':
        grads: Tensor = K.filters.spatial_gradient(x_gray, order=2) 
        grads_x = grads[:, :, 0]
        grads_y = grads[:, :, 1]
        
        output = K.utils.tensor_to_image(1. - grads_x.clamp(0., 1.))
        
    elif detector == '2nd order derivatives in y':
        grads: Tensor = K.filters.spatial_gradient(x_gray, order=2)
        grads_x = grads[:, :, 0]
        grads_y = grads[:, :, 1]
        
        output = K.utils.tensor_to_image(1. - grads_y.clamp(0., 1.))
        
    elif detector == 'Sobel':
        x_sobel: Tensor = K.filters.sobel(x_gray)
        output = K.utils.tensor_to_image(1. - x_sobel)

    elif detector == 'Laplacian':
        x_laplacian: Tensor = K.filters.laplacian(x_gray, kernel_size=5)
        output = K.utils.tensor_to_image(1. - x_laplacian.clamp(0., 1.))

    else:
        x_canny: Tensor = K.filters.canny(x_gray)[0]
        output = K.utils.tensor_to_image(1. - x_canny.clamp(0., 1.0))

    return output
    
title = "Tactile Graphic Generator"
description = "<p style='text-align: center'>To use it, simply upload your image. Press Enter, them click on the tactile graphic to download.</p>"

iface = gr.Interface(edge_detection, 
            [
            	gr.Image(type="filepath"),
	            gr.Dropdown(choices=["1st order derivates in x", "1st order derivates in y", "2nd order derivatives in x", "2nd order derivatives in y", "Sobel", "Laplacian", "Canny"])
            ],
            "image",
            title=title,
            description=description,
            
)

iface.launch()