import gradio as gr
import cv2

import torch
import kornia as K

def load_torch_image(fname):
    img = K.image_to_tensor(fname, False).float() / 255.
    img = K.color.bgr_to_rgb(img)
    return img


def enhance(file, brightness, contrast, saturation, gamma, hue):
    fname = file.name
    im = cv2.imread(fname)                 
    img = load_torch_image(im)

    x_out: torch.Tensor = K.enhance.adjust_brightness(img, float(brightness))
    x_out = K.enhance.adjust_contrast(x_out, float(contrast))
    x_out = K.enhance.adjust_saturation(x_out, float(saturation))
    x_out = K.enhance.adjust_gamma(x_out, float(gamma))
    x_out = K.enhance.adjust_hue(x_out, float(hue))

    return K.utils.tensor_to_image(x_out)


examples = [
    ["examples/ninja_turtles.jpg", 0, 1, 1, 1, 0],
    ["examples/kitty.jpg", 0, 1, 1, 1, 0],
]

title = "Kornia Image Enhancements"
description = "Gradio demo for Kornia's Image Enhancements. To use it, simply upload your image, or click one of the examples to load them, and use the sliders to enhance! Read more at the links below."
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/image_enhancement.html' target='_blank'>Kornia Enhancements Tutorial</a></p>"

iface = gr.Interface(
    enhance,
    [
        gr.inputs.Image(type="file"),
        gr.inputs.Slider(minimum=0, maximum=1, step=0.1, default=0, label="Brightness"),
        gr.inputs.Slider(minimum=0, maximum=4, step=0.1, default=1, label="Contrast"),
        gr.inputs.Slider(minimum=0, maximum=4, step=0.1, default=1, label="Saturation"),
        gr.inputs.Slider(minimum=0, maximum=1, step=0.1, default=1, label="Gamma"),
        gr.inputs.Slider(minimum=0, maximum=4, step=0.1, default=0, label="Hue"),
    ],
    "image",
    examples=examples,
    title=title,
    description=description,
    article=article,
    live=True
)

iface.launch()