import gradio as gr
import cv2
import numpy as np


class ImageToSketchEffect:
    """Transform an image to mimic the effect of a sketch drawn with a pencil."""

    def __init__(
            self,
            blur_strength: int = 5,
            blur_kernel_size: tuple = (0, 0),
            edge_enhance_level: int = None,
            custom_sharpen_kernel: np.ndarray = None,
    ):
        self.blur_strength = blur_strength
        self.blur_kernel_size = blur_kernel_size
        self.edge_enhance_level = edge_enhance_level
        self.sharpen_kernel = np.array([[0, -1, 0], [-1, edge_enhance_level, -1],
                                        [0, -1, 0]]) if custom_sharpen_kernel is None else custom_sharpen_kernel

    def blend_mode_dodge(self, top_layer: np.ndarray, bottom_layer: np.ndarray) -> np.ndarray:
        output = bottom_layer * 255.0 / (255.0 - top_layer)
        output[output > 255] = 255
        output[bottom_layer == 255] = 255
        return output.astype('uint8')

    def enhance_edges(self, img: np.ndarray) -> np.ndarray:
        if self.edge_enhance_level is not None and isinstance(self.edge_enhance_level, int):
            inverse_image = 255 - img
            return 255 - cv2.filter2D(src=inverse_image, ddepth=-1, kernel=self.sharpen_kernel)
        return img

    def apply_sketch_effect(self, image: np.ndarray) -> (np.ndarray, np.ndarray):
        grayscale_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        inverted_image = 255 - grayscale_image
        blurred_image = cv2.GaussianBlur(inverted_image, ksize=self.blur_kernel_size, sigmaX=self.blur_strength)
        sketch_image = self.blend_mode_dodge(blurred_image, grayscale_image)
        sketch_image = self.enhance_edges(sketch_image)
        blurred_for_alpha = cv2.GaussianBlur(sketch_image, (5, 5), 0)
        if len(blurred_for_alpha.shape) == 3 and blurred_for_alpha.shape[2] == 3:
            blurred_for_alpha = cv2.cvtColor(blurred_for_alpha, cv2.COLOR_BGR2GRAY)
        _, alpha_mask = cv2.threshold(blurred_for_alpha, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

        final_sketch_transparent = cv2.merge((blurred_for_alpha, blurred_for_alpha, blurred_for_alpha, alpha_mask))

        return sketch_image, final_sketch_transparent


def process_image(image, blur_strength, blur_kernel_size, edge_enhance_level):
    sketch_effect = ImageToSketchEffect(blur_strength, (blur_kernel_size, blur_kernel_size), edge_enhance_level)
    return sketch_effect.apply_sketch_effect(image)


iface = gr.Interface(fn=process_image,
                     inputs=[gr.Image(), gr.Slider(1, 10, value=7), gr.Slider(3, 25, step=2, value=15),
                             gr.Slider(1, 10, value=7)],
                     outputs=[gr.Image(type="numpy", label="Pencil Sketch"),gr.Image(type="numpy", label="Transparent Pencil Sketch")],
                     title="Image to Sketch Converter",
                     description="Upload an image to convert it into a pencil sketch.")

if __name__ == "__main__":
    iface.launch()