import numpy as np
import cv2
from PIL import Image
from cv2.ximgproc import guidedFilter
import gradio as gr


def clean_image(input_image: Image) -> Image:
    img = np.array(input_image).astype(np.float32)
    y = img.copy()

    for _ in range(64):
        y = cv2.bilateralFilter(y, 5, 8, 8)

    for _ in range(4):
        y = guidedFilter(img, y, 4, 16)

    output_image = Image.fromarray(y.clip(0, 255).astype(np.uint8))
    return output_image


def example(_image):
    pass


def ui():
    with gr.Blocks() as app:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Input Image")
                start_btn = gr.Button(value="Start", variant="primary")

                gr.Examples(
                    examples=[
                        ["./examples/sample1.jpg"],
                        ["./examples/sample2.jpg"],
                        ["./examples/sample3.jpg"],
                    ],
                    inputs=[input_image],
                    outputs=[],
                    fn=example,
                    cache_examples=True,
                )

            with gr.Column():
                output_image = gr.Image(type="pil", label="Output Image")

        gr.Markdown(
            "The lllyasviel's original repo is [here](https://github.com/lllyasviel/AdverseCleaner/tree/main)."
        )

        start_btn.click(fn=clean_image, inputs=[input_image], outputs=[output_image])

    app.launch()


if __name__ == "__main__":
    ui()