Cinderella / main.py
alrab222's picture
Create main.py
000a6ec
raw
history blame
No virus
610 Bytes
import json
import gradio as gr
import yolov5
model = yolov5.load('Cinderella.pt')
def inference(gr_input):
with open("json/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()