import gradio as gr
import os
from model.model import TextureSynthesisCNN
from model.utils import convert_tensor_to_PIL_image

def synth_image(image, epochs=10):
    synthesizer = TextureSynthesisCNN(tex_exemplar_image=image)
    output_tensor = synthesizer.synthesize_texture(num_epochs=int(epochs))
    output_image = convert_tensor_to_PIL_image(output_tensor)
    return output_image

demo = gr.Interface(
    fn=synth_image,
    inputs=[gr.Image(type="numpy"), 
            gr.Slider(label="Num epochs to optimize for", value=1, minimum=1, maximum=400, step=1, info="Optimizing for higher epochs gives better results. However, note that since this is a CNN based texture synthesiszer, it can take upto 5-10mins per synthesis.")],
    outputs=[gr.Image(type="pil")],
    flagging_options=["blurry", "incorrect"],
    examples=[
        [os.path.join(os.path.dirname(__file__), "images/blotchy_0025.png"), 10],
        [os.path.join(os.path.dirname(__file__), "images/blotchy_0027.png"), 10],
        [os.path.join(os.path.dirname(__file__), "images/cracked_0080.png"), 10],
        [os.path.join(os.path.dirname(__file__), "images/scenery.png"), 10],
    ],
)

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