import gradio as gr

import cv2
import numpy as np

import torch
import kornia as K
from kornia.core import Tensor

def edge_detection(file, detector): 

    img_bgr: np.ndarray = cv2.imread(file, cv2.IMREAD_COLOR)

    x_bgr: torch.Tensor = K.utils.image_to_tensor(img_bgr)  # CxHxWx
    x_bgr = x_bgr[None,...].float() / 255.

    x_rgb: torch.Tensor = K.color.bgr_to_rgb(x_bgr)
    x_gray = K.color.rgb_to_grayscale(x_rgb)
    

    if detector == '1st order derivates in x':
        grads: torch.Tensor = K.filters.spatial_gradient(x_gray, order=1)  # BxCx2xHxW
        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: torch.Tensor = K.filters.spatial_gradient(x_gray, order=1)  # BxCx2xHxW
        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: torch.Tensor = K.filters.spatial_gradient(x_gray, order=2)  # BxCx2xHxW
        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: torch.Tensor = K.filters.spatial_gradient(x_gray, order=2)  # BxCx2xHxW
        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: torch.Tensor = K.filters.sobel(x_gray)
        output = K.utils.tensor_to_image(1. - x_sobel)

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

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

    return output
    
    
examples = [
    ["examples/huggingface.jpg"],
    ["examples/doraemon.jpg"]
]

title = "Kornia Edge Detection"
description = "<p style='text-align: center'>This is a Gradio demo for Kornia's Edge Detection.</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 edge detector 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/filtering_edges.html' target='_blank'>Kornia Edge Detection Tutorial</a></p>"

iface = gr.Interface(fn=edge_detection, 
             inputs=[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"],
	      	 	  value="1st order derivates in x")
                    ],
            outputs="image",
            examples=examples
            
)

iface.launch()