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( "#
デレマス画像分類器\n" "##
シンデレラガールズのキャラの画像を貼ると、機械学習で判別できるモデルです\n" ) inputs = gr.inputs.Image() output = gr.Textbox() btn = gr.Button("judge") btn.click(inference, inputs, output) app.launch()