import gradio as gr
from PIL import Image
import os
import numpy as np
# from outpaint import outpainting
# from model import colorazation, UNETmodel, utils1
# from model import inference, model
# from model import colorazation, deeplabmodel, utils 
from model import MainModel
import inference as inf


# pretrained model
def colorize_image(image):
     # Load the model
    # file_path = 'ImageColorizationModel10.pth'
    file_path = './model/model_final.pth'
    model_2 = inf.load_model(model_class=MainModel, file_path=file_path)
    output_img = inf.predict_color(model_2, image=image)
    return output_img



# pretrained model
colorization_interface = gr.Interface(
    colorize_image,
    gr.Image(type="pil", label="Input Image"),
    [gr.Image(type="pil", label="Output Image")],
    title="Image Colorization",
    description="Upload an image to perform colorization.",

)

# deeplab model
# depinterface = gr.Interface(
#     depColorize_image,
#     gr.Image(type="pil", label="Input Image"),
#     [gr.Image(type="pil", label="Output Image")],
#     title="Image Colorization",
#     description="Upload an image to perform colorization.",

# )

# scratch mod

# Launch the interface
# interface.launch(share=True)
with gr.TabbedInterface([ colorization_interface ], ["Colorization_pretrain_unet"]) as tabs:
    tabs.launch(share=True)