import gradio as gr
from PIL import Image
import tempfile


def predict_depth(model, image):
    depth = model.infer_pil(image)
    return depth


def create_demo():
    with gr.Row():
        with gr.Column(scale=1):
            model_type = gr.Dropdown(["Prismer-Base", "Prismer-Large"], label="Model Size", value="Prismer-Base")
            rgb = gr.Image(label="Input Image", type='pil')
            submit = gr.Button("Submit")
        with gr.Column(scale=2):
            pred = gr.Textbox(label="Model Prediction")
            with gr.Row():
                depth = gr.Image(label="Depth")
                edge = gr.Image(label="Edge")
                normals = gr.Image(label="Normals")
            with gr.Row():
                seg = gr.Image(label="Segmentation")
                obj_det = gr.Image(label="Object Detection")
                ocr_det = gr.Image(label="OCR Detection")

    def on_submit(im, model_type):
        return pred, depth, edge, normals, seg, obj_det, ocr_det

    submit.click(on_submit, inputs=[rgb, model_type], outputs=[pred, depth, edge, normals, seg, obj_det, ocr_det])
    examples = gr.Examples(examples=["examples/1.png"], inputs=[rgb])