import json

import gradio as gr
import yolov5

model = yolov5.load('model/Cinderella.pt')

#てすと

def inference(gr_input):
   with open("json/idol_classes.json", "r")as f:
      dic = json.load(f)
   predictions = model(gr_input).pred[0]
   name = [x for x in dic.values()]
   idols = []
   for i in predictions[:, 5]:  # 検出
      idols.append(name[int(i.item())])
   if len(idols) == 0:
      return "該当するクラスが存在しません"
   else:
      text = "\n".join(idols)
      return text

with gr.Blocks() as app:
    gr.Markdown('<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=alrab222.Cinderella" />')
    gr.Markdown(
                "# <center> デレマス画像分類器\n"
                "## <center> シンデレラガールズのキャラの画像を貼ると、機械学習で判別できるモデルです\n"
    )
    inputs = gr.inputs.Image()
    output = gr.Textbox()
    btn = gr.Button("judge")
    btn.click(inference, inputs, output)
app.launch()