import gradio as gr
from setup import setup
import cv2
from PIL import Image
from manga_line_extraction.model import MangaLineExtractor
from anime2sketch.model import Anime2Sketch

setup()

print("Setup finished")

extractor = MangaLineExtractor("./models/erika.pth", "cpu")
to_sketch = Anime2Sketch("./models/netG.pth", "cpu")

print("Model loaded")


def extract(image):
    return extractor.predict(image)


def convert_to_sketch(image):
    return to_sketch.predict(image)


def start(image):
    return [extract(image), convert_to_sketch(Image.fromarray(image).convert("RGB"))]


def ui():
    with gr.Blocks() as blocks:
        gr.Markdown(
            """
        # Anime to Sketch 
        Unofficial demo for converting illustrations into sketches. 
        Original repos:
        - [MangaLineExtraction_PyTorch](https://github.com/ljsabc/MangaLineExtraction_PyTorch)
        - [Anime2Sketch](https://github.com/Mukosame/Anime2Sketch)
        """
        )

        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input", interactive=True)

                extract_btn = gr.Button("Extract", variant="primary")

            with gr.Column():
                # with gr.Row():
                extract_output_img = gr.Image(
                    label="MangaLineExtraction", interactive=False
                )
                to_sketch_output_img = gr.Image(label="Anime2Sketch", interactive=False)

        gr.Examples(
            fn=start,
            examples=[
                ["./examples/1.jpg"],
                ["./examples/2.jpg"],
                ["./examples/3.jpg"],
                ["./examples/4.jpg"],
                ["./examples/5.jpg"],
                ["./examples/6.jpg"],
                ["./examples/7.jpg"],
                ["./examples/8.jpg"],
            ],
            inputs=[input_img],
            outputs=[extract_output_img, to_sketch_output_img],
            label="Examples",
            cache_examples=True,
        )

        gr.Markdown("Images are from nijijourney.")

        extract_btn.click(
            fn=start,
            inputs=[input_img],
            outputs=[extract_output_img, to_sketch_output_img],
        )

    return blocks


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