import gradio as gr

from PIL import Image

from predictor import Predictor

MODEL_SALT_BESTIMAGE = "saltacc/beit-bestimage-salt"
MODEL_CAFE_AESTHETIC = "cafeai/cafe_aesthetic"

MODEL_NAMES = [
    MODEL_SALT_BESTIMAGE,
    MODEL_CAFE_AESTHETIC,
]

models = {
    MODEL_SALT_BESTIMAGE: Predictor(MODEL_SALT_BESTIMAGE),
    MODEL_CAFE_AESTHETIC: Predictor(MODEL_CAFE_AESTHETIC),
}


def predict(image: Image.Image) -> list[dict[str, float]]:
    results = []

    for model_name in MODEL_NAMES:
        results.append(models[model_name].predict([image])[0])

    return results


def construct_ui():
    with gr.Blocks() as ui:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="Input Image",
                    type="pil",
                    source="upload",
                    interactive=True,
                )

                submit_btn = gr.Button(
                    value="Submit",
                    variant="primary",
                )

            with gr.Column():
                result_salt = gr.Label(
                    label=MODEL_SALT_BESTIMAGE,
                )

                result_cafe = gr.Label(
                    label=MODEL_CAFE_AESTHETIC,
                )

        submit_btn.click(
            predict, inputs=[input_image], outputs=[result_salt, result_cafe]
        )

    return ui


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