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


inputs = gr.inputs.Image()

interface = gr.Interface(fn=inference, inputs=inputs, outputs="text")
interface.launch()