import os

import gradio as gr
from huggingface_hub import HfFileSystem
from imgutils.generic import classify_predict_score
from natsort import natsorted

hf_fs = HfFileSystem()

_REPOSITORY = 'deepghs/anime_aesthetic'
_DEFAULT_MODEL = 'caformer_s36_v0_ls0.2'
_MODELS = natsorted([
    os.path.dirname(os.path.relpath(file, _REPOSITORY))
    for file in hf_fs.glob(f'{_REPOSITORY}/*/model.onnx')
])

LABELS = ["worst", "low", "normal", "good", "great", "best", "masterpiece"]


def _fn_predict(image, model):
    scores = classify_predict_score(
        image=image,
        repo_id=_REPOSITORY,
        model_name=model,
    )
    final_score = sum(i * scores[label] for i, label in enumerate(LABELS))
    return final_score, scores


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(type='pil', label='Original Image')
                gr_model = gr.Dropdown(_MODELS, value=_DEFAULT_MODEL, label='Model')
                gr_submit = gr.Button(value='Submit', variant='primary')

            with gr.Column():
                gr_score = gr.Text(label='Aesthetic Score (0~6)', value='')
                gr_output = gr.Label(label='Aesthetic Classes')

            gr_submit.click(
                _fn_predict,
                inputs=[gr_input_image, gr_model],
                outputs=[gr_score, gr_output],
            )

    demo.queue(os.cpu_count()).launch()